



'''simple docstring'''




import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
						import torch

						from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin






class 		_UpperCamelCase						:


			'''simple docstring'''

			def __init__(							self  :				List[str]							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				Union[str, Any]=None							,  lowerCAmelCase__  :				int=None							,  lowerCAmelCase__  :				Optional[int]=None							,  lowerCAmelCase__  :				Any="resnet50"							,  lowerCAmelCase__  :				Any=3							,  lowerCAmelCase__  :				Optional[Any]=3_2							,  lowerCAmelCase__  :				int=3							,  lowerCAmelCase__  :				List[str]=True							,  lowerCAmelCase__  :				List[Any]=True							,  ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[str]					 =	parent
					__SCREAMING_SNAKE_CASE :	List[str]					 =	out_indices if out_indices is not None else [4]
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	stage_names
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	out_features
					__SCREAMING_SNAKE_CASE :	Tuple					 =	backbone
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	batch_size
					__SCREAMING_SNAKE_CASE :	int					 =	image_size
					__SCREAMING_SNAKE_CASE :	List[str]					 =	num_channels
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	use_pretrained_backbone
					__SCREAMING_SNAKE_CASE :	Tuple					 =	is_training

			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)
					__SCREAMING_SNAKE_CASE :	int					 =	self.get_config()

					return config, pixel_values

			def UpperCamelCase__      (							self  :				Tuple		):





					"""simple docstring"""






					return TimmBackboneConfig(
					    image_size=self.image_size							,  num_channels=self.num_channels							,  out_features=self.out_features							,  out_indices=self.out_indices							,  stage_names=self.stage_names							,  use_pretrained_backbone=self.use_pretrained_backbone							,  backbone=self.backbone							,  )

			def UpperCamelCase__      (							self  :				Any							,  lowerCAmelCase__  :				List[Any]							,  lowerCAmelCase__  :				Optional[Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	TimmBackbone(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model(lowerCAmelCase__		)
					self.parent.assertEqual(
					    result.feature_map[-1].shape							,  (self.batch_size, model.channels[-1], 1_4, 1_4)							,  )



			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.prepare_config_and_inputs()
					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	config_and_inputs
					__SCREAMING_SNAKE_CASE :	List[str]					 =	{"""pixel_values""": pixel_values}
					return config, inputs_dict






@require_torch
@require_timm
class 		_UpperCamelCase						(							lowerCamelCase__					,      lowerCamelCase__					,      lowerCamelCase__					,      unittest.TestCase					):


			'''simple docstring'''

			_A   :	Union[str, Any]    					=	(TimmBackbone,) if is_torch_available() else ()
			_A   :	Tuple    					=	{'''feature-extraction''': TimmBackbone} if is_torch_available() else {}
			_A   :	Tuple    					=	False
			_A   :	str    					=	False
			_A   :	Optional[int]    					=	False
			_A   :	str    					=	False

			def UpperCamelCase__      (							self  :				Dict		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	TimmBackboneModelTester(self		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	ConfigTester(self							,  config_class=lowerCAmelCase__							,  has_text_modality=lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					self.config_tester.create_and_test_config_to_json_string()
					self.config_tester.create_and_test_config_to_json_file()
					self.config_tester.create_and_test_config_from_and_save_pretrained()
					self.config_tester.create_and_test_config_with_num_labels()
					self.config_tester.check_config_can_be_init_without_params()
					self.config_tester.check_config_arguments_init()

			def UpperCamelCase__      (							self  :				Dict		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	"""resnet18"""
					__SCREAMING_SNAKE_CASE :	List[str]					 =	"""microsoft/resnet-18"""

					__SCREAMING_SNAKE_CASE :	Dict					 =	AutoBackbone.from_pretrained(lowerCAmelCase__							,  use_timm_backbone=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	int					 =	AutoBackbone.from_pretrained(lowerCAmelCase__		)

					self.assertEqual(len(timm_model.out_features		)							,  len(transformers_model.out_features		)		)
					self.assertEqual(len(timm_model.stage_names		)							,  len(transformers_model.stage_names		)		)
					self.assertEqual(timm_model.channels							,  transformers_model.channels		)
					# Out indices are set to the last layer by default. For timm models, we don't know
					# the number of layers in advance, so we set it to (-1,), whereas for transformers
					# models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
					self.assertEqual(timm_model.out_indices							,  (-1,)		)
					self.assertEqual(transformers_model.out_indices							,  [len(timm_model.stage_names		) - 1]		)

					__SCREAMING_SNAKE_CASE :	Any					 =	AutoBackbone.from_pretrained(lowerCAmelCase__							,  use_timm_backbone=lowerCAmelCase__							,  out_indices=[1, 2, 3]		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	AutoBackbone.from_pretrained(lowerCAmelCase__							,  out_indices=[1, 2, 3]		)

					self.assertEqual(timm_model.out_indices							,  transformers_model.out_indices		)
					self.assertEqual(len(timm_model.out_features		)							,  len(transformers_model.out_features		)		)
					self.assertEqual(timm_model.channels							,  transformers_model.channels		)

			@unittest.skip("""TimmBackbone doesn't support feed forward chunking"""		)
			def UpperCamelCase__      (							self  :				Optional[int]		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone doesn't have num_hidden_layers attribute"""		)
			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone initialization is managed on the timm side"""		)
			def UpperCamelCase__      (							self  :				List[Any]		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone models doesn't have inputs_embeds"""		)
			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone models doesn't have inputs_embeds"""		)
			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone model cannot be created without specifying a backbone checkpoint"""		)
			def UpperCamelCase__      (							self  :				Union[str, Any]		):





					"""simple docstring"""






					pass

			@unittest.skip("""Only checkpoints on timm can be loaded into TimmBackbone"""		)
			def UpperCamelCase__      (							self  :				Dict		):





					"""simple docstring"""






					pass

			@unittest.skip("""model weights aren't tied in TimmBackbone."""		)
			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					pass

			@unittest.skip("""model weights aren't tied in TimmBackbone."""		)
			def UpperCamelCase__      (							self  :				List[Any]		):





					"""simple docstring"""






					pass

			@unittest.skip("""Only checkpoints on timm can be loaded into TimmBackbone"""		)
			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					pass

			@unittest.skip("""Only checkpoints on timm can be loaded into TimmBackbone"""		)
			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone doesn't have hidden size info in its configuration."""		)
			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					pass

			@unittest.skip("""TimmBackbone doesn't support output_attentions."""		)
			def UpperCamelCase__      (							self  :				List[Any]		):





					"""simple docstring"""






					pass

			@unittest.skip("""Safetensors is not supported by timm."""		)
			def UpperCamelCase__      (							self  :				str		):





					"""simple docstring"""






					pass

			@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""		)
			def UpperCamelCase__      (							self  :				Tuple		):





					"""simple docstring"""






					pass

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





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Any					 =	self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE :	Tuple					 =	model_class(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	inspect.signature(model.forward		)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							__SCREAMING_SNAKE_CASE :	int					 =	[*signature.parameters.keys()]

							__SCREAMING_SNAKE_CASE :	List[Any]					 =	["""pixel_values"""]
							self.assertListEqual(arg_names[:1]							,  lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				str		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	int					 =	self.model_tester.prepare_config_and_inputs_for_common()
					__SCREAMING_SNAKE_CASE :	Any					 =	True
					__SCREAMING_SNAKE_CASE :	List[str]					 =	self.has_attentions

					# no need to test all models as different heads yield the same functionality
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.all_model_classes[0]
					__SCREAMING_SNAKE_CASE :	Any					 =	model_class(lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	Any					 =	self._prepare_for_class(lowerCAmelCase__							,  lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	List[str]					 =	model(**lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	outputs[0][-1]

					# Encoder-/Decoder-only models
					__SCREAMING_SNAKE_CASE :	Tuple					 =	outputs.hidden_states[0]
					hidden_states.retain_grad()

					if self.has_attentions:
							__SCREAMING_SNAKE_CASE :	str					 =	outputs.attentions[0]
							attentions.retain_grad()

					output.flatten()[0].backward(retain_graph=lowerCAmelCase__		)

					self.assertIsNotNone(hidden_states.grad		)

					if self.has_attentions:
							self.assertIsNotNone(attentions.grad		)



			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model_class(lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__SCREAMING_SNAKE_CASE :	List[str]					 =	model(**lowerCAmelCase__		)

							self.assertEqual(len(result.feature_maps		)							,  len(config.out_indices		)		)
							self.assertEqual(len(model.channels		)							,  len(config.out_indices		)		)

							# Check output of last stage is taken if out_features=None, out_indices=None
							__SCREAMING_SNAKE_CASE :	List[Any]					 =	copy.deepcopy(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	str					 =	None
							__SCREAMING_SNAKE_CASE :	Dict					 =	model_class(lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	model(**lowerCAmelCase__		)

							self.assertEqual(len(result.feature_maps		)							,  1		)
							self.assertEqual(len(model.channels		)							,  1		)

							# Check backbone can be initialized with fresh weights
							__SCREAMING_SNAKE_CASE :	str					 =	copy.deepcopy(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	str					 =	False
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model_class(lowerCAmelCase__		)
							model.to(lowerCAmelCase__		)
							model.eval()
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model(**lowerCAmelCase__		)




'''simple docstring'''




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 		_UpperCamelCase						(							unittest.TestCase					):


			'''simple docstring'''

			def UpperCamelCase__      (							self  :				Optional[Any]		):





					"""simple docstring"""






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

			def UpperCamelCase__      (							self  :				List[Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	int					 =	FlaxControlNetModel.from_pretrained(
					    """lllyasviel/sd-controlnet-canny"""							,  from_pt=lowerCAmelCase__							,  dtype=jnp.bfloataa		)
					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Tuple					 =	FlaxStableDiffusionControlNetPipeline.from_pretrained(
					    """runwayml/stable-diffusion-v1-5"""							,  controlnet=lowerCAmelCase__							,  from_pt=lowerCAmelCase__							,  dtype=jnp.bfloataa		)
					__SCREAMING_SNAKE_CASE :	str					 =	controlnet_params

					__SCREAMING_SNAKE_CASE :	str					 =	"""bird"""
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	jax.device_count()
					__SCREAMING_SNAKE_CASE :	str					 =	pipe.prepare_text_inputs([prompts] * num_samples		)

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

					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	jax.random.PRNGKey(0		)
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	jax.random.split(lowerCAmelCase__							,  jax.device_count()		)

					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	replicate(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	shard(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	shard(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	List[Any]					 =	pipe(
					    prompt_ids=lowerCAmelCase__							,  image=lowerCAmelCase__							,  params=lowerCAmelCase__							,  prng_seed=lowerCAmelCase__							,  num_inference_steps=5_0							,  jit=lowerCAmelCase__							,  ).images
					assert images.shape == (jax.device_count(), 1, 7_6_8, 5_1_2, 3)

					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	jnp.asarray(jax.device_get(image_slice.flatten()		)		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	jnp.array(
					    [0.16_79_69, 0.11_66_99, 0.08_15_43, 0.15_42_97, 0.13_28_12, 0.10_88_87, 0.16_99_22, 0.16_99_22, 0.20_50_78]		)
					print(F"output_slice: {output_slice}"		)
					assert jnp.abs(output_slice - expected_slice		).max() < 1E-2



			def UpperCamelCase__      (							self  :				str		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	FlaxControlNetModel.from_pretrained(
					    """lllyasviel/sd-controlnet-openpose"""							,  from_pt=lowerCAmelCase__							,  dtype=jnp.bfloataa		)
					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Any					 =	FlaxStableDiffusionControlNetPipeline.from_pretrained(
					    """runwayml/stable-diffusion-v1-5"""							,  controlnet=lowerCAmelCase__							,  from_pt=lowerCAmelCase__							,  dtype=jnp.bfloataa		)
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	controlnet_params

					__SCREAMING_SNAKE_CASE :	List[str]					 =	"""Chef in the kitchen"""
					__SCREAMING_SNAKE_CASE :	int					 =	jax.device_count()
					__SCREAMING_SNAKE_CASE :	Tuple					 =	pipe.prepare_text_inputs([prompts] * num_samples		)

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

					__SCREAMING_SNAKE_CASE :	Dict					 =	jax.random.PRNGKey(0		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	jax.random.split(lowerCAmelCase__							,  jax.device_count()		)

					__SCREAMING_SNAKE_CASE :	List[Any]					 =	replicate(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	shard(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	shard(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	Tuple					 =	pipe(
					    prompt_ids=lowerCAmelCase__							,  image=lowerCAmelCase__							,  params=lowerCAmelCase__							,  prng_seed=lowerCAmelCase__							,  num_inference_steps=5_0							,  jit=lowerCAmelCase__							,  ).images
					assert images.shape == (jax.device_count(), 1, 7_6_8, 5_1_2, 3)

					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

					__SCREAMING_SNAKE_CASE :	Tuple					 =	jnp.asarray(jax.device_get(image_slice.flatten()		)		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	jnp.array(
					    [[0.27_14_84, 0.26_17_19, 0.27_53_91, 0.27_73_44, 0.27_92_97, 0.29_10_16, 0.29_49_22, 0.30_27_34, 0.30_27_34]]		)
					print(F"output_slice: {output_slice}"		)
					assert jnp.abs(output_slice - expected_slice		).max() < 1E-2