





import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import (
				    LiltForQuestionAnswering,
				    LiltForSequenceClassification,
				    LiltForTokenClassification,
				    LiltModel,
				)
				from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST

class 		SCREAMING_SNAKE_CASE__    :


						def __init__(self	:				List[str]    ,      a__	:				List[Any]    ,      a__	:				Optional[Any]=13    ,      a__	:				List[Any]=7    ,      a__	:				Union[str, Any]=True    ,      a__	:				Optional[int]=True    ,      a__	:				List[Any]=True    ,      a__	:				List[str]=True    ,      a__	:				Optional[int]=99    ,      a__	:				Dict=24    ,      a__	:				Optional[int]=2    ,      a__	:				Optional[int]=6    ,      a__	:				List[str]=37    ,      a__	:				Optional[Any]="gelu"    ,      a__	:				str=0.1    ,      a__	:				Dict=0.1    ,      a__	:				Optional[Any]=512    ,      a__	:				Dict=16    ,      a__	:				Union[str, Any]=2    ,      a__	:				int=0.0_2    ,      a__	:				int=3    ,      a__	:				Optional[int]=None    ,      a__	:				str=1000    ,      ):





									"""simple docstring"""



									__snake_case								=       parent
									__snake_case								=       batch_size
									__snake_case								=       seq_length
									__snake_case								=       is_training
									__snake_case								=       use_input_mask
									__snake_case								=       use_token_type_ids
									__snake_case								=       use_labels
									__snake_case								=       vocab_size
									__snake_case								=       hidden_size
									__snake_case								=       num_hidden_layers
									__snake_case								=       num_attention_heads
									__snake_case								=       intermediate_size
									__snake_case								=       hidden_act
									__snake_case								=       hidden_dropout_prob
									__snake_case								=       attention_probs_dropout_prob
									__snake_case								=       max_position_embeddings
									__snake_case								=       type_vocab_size
									__snake_case								=       type_sequence_label_size
									__snake_case								=       initializer_range
									__snake_case								=       num_labels
									__snake_case								=       scope
									__snake_case								=       range_bbox


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





									"""simple docstring"""



									__snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.vocab_size				)

									__snake_case								=       ids_tensor([self.batch_size, self.seq_length, 4]    ,      self.range_bbox				)
									# Ensure that bbox is legal
									for i in range(bbox.shape[0]				):
												for j in range(bbox.shape[1]				):
															if bbox[i, j, 3] < bbox[i, j, 1]:
																		__snake_case								=       bbox[i, j, 3]
																		__snake_case								=       bbox[i, j, 1]
																		__snake_case								=       t
															if bbox[i, j, 2] < bbox[i, j, 0]:
																		__snake_case								=       bbox[i, j, 2]
																		__snake_case								=       bbox[i, j, 0]
																		__snake_case								=       t

									__snake_case								=       None
									if self.use_input_mask:
												__snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      vocab_size=2				)

									__snake_case								=       None
									if self.use_token_type_ids:
												__snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.type_vocab_size				)

									__snake_case								=       None
									__snake_case								=       None
									if self.use_labels:
												__snake_case								=       ids_tensor([self.batch_size]    ,      self.type_sequence_label_size				)
												__snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.num_labels				)

									__snake_case								=       self.get_config()

									return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels


						def        a			(self	:				Tuple				):





									"""simple docstring"""



									return LiltConfig(
									    vocab_size=self.vocab_size    ,      hidden_size=self.hidden_size    ,      num_hidden_layers=self.num_hidden_layers    ,      num_attention_heads=self.num_attention_heads    ,      intermediate_size=self.intermediate_size    ,      hidden_act=self.hidden_act    ,      hidden_dropout_prob=self.hidden_dropout_prob    ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,      max_position_embeddings=self.max_position_embeddings    ,      type_vocab_size=self.type_vocab_size    ,      initializer_range=self.initializer_range    ,      )


						def        a			(self	:				Optional[Any]    ,      a__	:				Any    ,      a__	:				Optional[Any]    ,      a__	:				int    ,      a__	:				List[str]    ,      a__	:				str    ,      a__	:				Tuple    ,      a__	:				List[Any]    ,      ):





									"""simple docstring"""



									__snake_case								=       LiltModel(config=lowerCamelCase__				)
									model.to(lowerCamelCase__				)
									model.eval()
									__snake_case								=       model(lowerCamelCase__    ,      bbox=lowerCamelCase__    ,      attention_mask=lowerCamelCase__    ,      token_type_ids=lowerCamelCase__				)
									__snake_case								=       model(lowerCamelCase__    ,      bbox=lowerCamelCase__    ,      token_type_ids=lowerCamelCase__				)
									__snake_case								=       model(lowerCamelCase__    ,      bbox=lowerCamelCase__				)

									self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.hidden_size)				)
									self.parent.assertEqual(result.pooler_output.shape    ,      (self.batch_size, self.hidden_size)				)


						def        a			(self	:				Tuple    ,      a__	:				List[Any]    ,      a__	:				Any    ,      a__	:				Dict    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[Any]    ,      a__	:				str    ,      a__	:				Optional[int]    ,      ):





									"""simple docstring"""



									__snake_case								=       self.num_labels
									__snake_case								=       LiltForTokenClassification(config=lowerCamelCase__				)
									model.to(lowerCamelCase__				)
									model.eval()
									__snake_case								=       model(
									    lowerCamelCase__    ,      bbox=lowerCamelCase__    ,      attention_mask=lowerCamelCase__    ,      token_type_ids=lowerCamelCase__    ,      labels=lowerCamelCase__				)
									self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.seq_length, self.num_labels)				)


						def        a			(self	:				Union[str, Any]    ,      a__	:				Optional[Any]    ,      a__	:				Optional[Any]    ,      a__	:				List[Any]    ,      a__	:				Dict    ,      a__	:				str    ,      a__	:				Union[str, Any]    ,      a__	:				Any    ,      ):





									"""simple docstring"""



									__snake_case								=       LiltForQuestionAnswering(config=lowerCamelCase__				)
									model.to(lowerCamelCase__				)
									model.eval()
									__snake_case								=       model(
									    lowerCamelCase__    ,      bbox=lowerCamelCase__    ,      attention_mask=lowerCamelCase__    ,      token_type_ids=lowerCamelCase__    ,      start_positions=lowerCamelCase__    ,      end_positions=lowerCamelCase__    ,      )
									self.parent.assertEqual(result.start_logits.shape    ,      (self.batch_size, self.seq_length)				)
									self.parent.assertEqual(result.end_logits.shape    ,      (self.batch_size, self.seq_length)				)


						def        a			(self	:				Optional[Any]				):





									"""simple docstring"""



									__snake_case								=       self.prepare_config_and_inputs()
									(
									    __snake_case			
									)								=       config_and_inputs
									__snake_case								=       {
									    '''input_ids''': input_ids,
									    '''bbox''': bbox,
									    '''token_type_ids''': token_type_ids,
									    '''attention_mask''': input_mask,
									}
									return config, inputs_dict






@require_torch
class 		SCREAMING_SNAKE_CASE__    (	__lowerCAmelCase						,				__lowerCAmelCase						,				__lowerCAmelCase						,				unittest.TestCase  ):
						A_						:    Tuple							=   (
						    (
						        LiltModel,
						        LiltForSequenceClassification,
						        LiltForTokenClassification,
						        LiltForQuestionAnswering,
						    )
						    if is_torch_available()
						    else ()
						)
						A_						:    Optional[int]							=   (
						    {
						        "feature-extraction": LiltModel,
						        "question-answering": LiltForQuestionAnswering,
						        "text-classification": LiltForSequenceClassification,
						        "token-classification": LiltForTokenClassification,
						        "zero-shot": LiltForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						A_						:    str							=   False
						A_						:    str							=   False


						def        a			(self	:				Dict    ,      a__	:				int    ,      a__	:				int    ,      a__	:				Union[str, Any]    ,      a__	:				int    ,      a__	:				Optional[int]				):





									"""simple docstring"""



									return True


						def        a			(self	:				List[Any]				):





									"""simple docstring"""



									__snake_case								=       LiltModelTester(self				)
									__snake_case								=       ConfigTester(self    ,      config_class=lowerCamelCase__    ,      hidden_size=37				)


						def        a			(self	:				List[str]				):





									"""simple docstring"""



									self.config_tester.run_common_tests()


						def        a			(self	:				List[Any]				):





									"""simple docstring"""



									__snake_case								=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCamelCase__				)


						def        a			(self	:				Optional[int]				):





									"""simple docstring"""



									__snake_case								=       self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
												__snake_case								=       type
												self.model_tester.create_and_check_model(*lowerCamelCase__				)


						def        a			(self	:				Optional[Any]				):





									"""simple docstring"""



									__snake_case								=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__				)


						def        a			(self	:				Optional[int]				):





									"""simple docstring"""



									__snake_case								=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowerCamelCase__				)


						@slow
						def        a			(self	:				Tuple				):





									"""simple docstring"""



									for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												__snake_case								=       LiltModel.from_pretrained(lowerCamelCase__				)
												self.assertIsNotNone(lowerCamelCase__				)






@require_torch
@slow
class 		SCREAMING_SNAKE_CASE__    (	unittest.TestCase  ):


						def        a			(self	:				int				):





									"""simple docstring"""



									__snake_case								=       LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''				).to(lowerCamelCase__				)

									__snake_case								=       torch.tensor([[1, 2]]    ,      device=lowerCamelCase__				)
									__snake_case								=       torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]    ,      device=lowerCamelCase__				)

									# forward pass
									with torch.no_grad():
												__snake_case								=       model(input_ids=lowerCamelCase__    ,      bbox=lowerCamelCase__				)

									__snake_case								=       torch.Size([1, 2, 768]				)
									__snake_case								=       torch.tensor(
									    [[-0.0_6_5_3, 0.0_9_5_0, -0.0_0_6_1], [-0.0_5_4_5, 0.0_9_2_6, -0.0_3_2_4]]    ,      device=lowerCamelCase__    ,      )

									self.assertTrue(outputs.last_hidden_state.shape    ,      lowerCamelCase__				)
									self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :, :3]    ,      lowerCamelCase__    ,      atol=1E-3				)				)








import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
    import jax
    import jax.numpy as jnp
    from flax.jax_utils import replicate
    from flax.training.common_utils import shard

@slow
@require_flax
class 		SCREAMING_SNAKE_CASE__    (	unittest.TestCase  ):


      def        a			(self	:				int				):





         """simple docstring"""



         super().tearDown()
         gc.collect()


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       FlaxControlNetModel.from_pretrained(
             '''lllyasviel/sd-controlnet-canny'''    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case	,     __snake_case								=       FlaxStableDiffusionControlNetPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''    ,      controlnet=a__    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case								=       controlnet_params

         __snake_case								=       '''bird'''
         __snake_case								=       jax.device_count()
         __snake_case								=       pipe.prepare_text_inputs([prompts] * num_samples				)

         __snake_case								=       load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''				)
         __snake_case								=       pipe.prepare_image_inputs([canny_image] * num_samples				)

         __snake_case								=       jax.random.PRNGKey(0				)
         __snake_case								=       jax.random.split(a__    ,      jax.device_count()				)

         __snake_case								=       replicate(a__				)
         __snake_case								=       shard(a__				)
         __snake_case								=       shard(a__				)

         __snake_case								=       pipe(
             prompt_ids=a__    ,      image=a__    ,      params=a__    ,      prng_seed=a__    ,      num_inference_steps=50    ,      jit=a__    ,      ).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

         __snake_case								=       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
         __snake_case								=       images[0, 253:256, 253:256, -1]

         __snake_case								=       jnp.asarray(jax.device_get(image_slice.flatten()				)				)
         __snake_case								=       jnp.array(
             [0.1_6_7_9_6_9, 0.1_1_6_6_9_9, 0.0_8_1_5_4_3, 0.1_5_4_2_9_7, 0.1_3_2_8_1_2, 0.1_0_8_8_8_7, 0.1_6_9_9_2_2, 0.1_6_9_9_2_2, 0.2_0_5_0_7_8]				)
         print(f"""output_slice: {output_slice}"""				)
         assert jnp.abs(output_slice - expected_slice				).max() < 1E-2


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       FlaxControlNetModel.from_pretrained(
             '''lllyasviel/sd-controlnet-openpose'''    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case	,     __snake_case								=       FlaxStableDiffusionControlNetPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''    ,      controlnet=a__    ,      from_pt=a__    ,      dtype=jnp.bfloataa				)
         __snake_case								=       controlnet_params

         __snake_case								=       '''Chef in the kitchen'''
         __snake_case								=       jax.device_count()
         __snake_case								=       pipe.prepare_text_inputs([prompts] * num_samples				)

         __snake_case								=       load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png'''				)
         __snake_case								=       pipe.prepare_image_inputs([pose_image] * num_samples				)

         __snake_case								=       jax.random.PRNGKey(0				)
         __snake_case								=       jax.random.split(a__    ,      jax.device_count()				)

         __snake_case								=       replicate(a__				)
         __snake_case								=       shard(a__				)
         __snake_case								=       shard(a__				)

         __snake_case								=       pipe(
             prompt_ids=a__    ,      image=a__    ,      params=a__    ,      prng_seed=a__    ,      num_inference_steps=50    ,      jit=a__    ,      ).images
         assert images.shape == (jax.device_count(), 1, 768, 512, 3)

         __snake_case								=       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
         __snake_case								=       images[0, 253:256, 253:256, -1]

         __snake_case								=       jnp.asarray(jax.device_get(image_slice.flatten()				)				)
         __snake_case								=       jnp.array(
             [[0.2_7_1_4_8_4, 0.2_6_1_7_1_9, 0.2_7_5_3_9_1, 0.2_7_7_3_4_4, 0.2_7_9_2_9_7, 0.2_9_1_0_1_6, 0.2_9_4_9_2_2, 0.3_0_2_7_3_4, 0.3_0_2_7_3_4]]				)
         print(f"""output_slice: {output_slice}"""				)
         assert jnp.abs(output_slice - expected_slice				).max() < 1E-2

