





'''simple docstring'''





import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowercase       :Union[str, Any]       =  logging.get_logger(__name__)

__lowercase       :str       =  {
    "Salesforce/blip-vqa-base": "https://huggingface.co/Salesforce/blip-vqa-base/resolve/main/config.json",
    "Salesforce/blip-vqa-capfit-large": (
        "https://huggingface.co/Salesforce/blip-vqa-base-capfit/resolve/main/config.json"
    ),
    "Salesforce/blip-image-captioning-base": (
        "https://huggingface.co/Salesforce/blip-image-captioning-base/resolve/main/config.json"
    ),
    "Salesforce/blip-image-captioning-large": (
        "https://huggingface.co/Salesforce/blip-image-captioning-large/resolve/main/config.json"
    ),
    "Salesforce/blip-itm-base-coco": "https://huggingface.co/Salesforce/blip-itm-base-coco/resolve/main/config.json",
    "Salesforce/blip-itm-large-coco": "https://huggingface.co/Salesforce/blip-itm-large-coco/resolve/main/config.json",
    "Salesforce/blip-itm-base-flikr": "https://huggingface.co/Salesforce/blip-itm-base-flikr/resolve/main/config.json",
    "Salesforce/blip-itm-large-flikr": (
        "https://huggingface.co/Salesforce/blip-itm-large-flikr/resolve/main/config.json"
    ),
}



class    _a     (				__lowerCAmelCase							):






							"""simple docstring"""

							snake_case_							 =						'''blip_text_model'''

							def __init__(      self				:     List[str]							,	a				:     Any=3_05_24							,	a				:     Optional[Any]=7_68							,	a				:     Tuple=7_68							,	a				:     Optional[Any]=30_72							,	a				:     Optional[Any]=7_68							,	a				:     Any=12							,	a				:     Optional[int]=8							,	a				:     List[str]=5_12							,	a				:     Any="gelu"							,	a				:     Union[str, Any]=1E-12							,	a				:     List[str]=0.0							,	a				:     Union[str, Any]=0.0							,	a				:     Union[str, Any]=0.02							,	a				:     Any=3_05_22							,	a				:     Optional[Any]=2							,	a				:     Tuple=0							,	a				:     int=1_02							,	a				:     Any=True							,	a				:     Any=True							,	**a				:     int							,	)	->Union[str, Any]:
													super().__init__(
													    pad_token_id=lowerCAmelCase_							,	bos_token_id=lowerCAmelCase_							,	eos_token_id=lowerCAmelCase_							,	sep_token_id=lowerCAmelCase_							,	**lowerCAmelCase_							,	)

													SCREAMING_SNAKE_CASE__     :				Any   =					vocab_size
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					hidden_size
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					encoder_hidden_size
													SCREAMING_SNAKE_CASE__     :				str   =					intermediate_size
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					projection_dim
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					hidden_dropout_prob
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					num_hidden_layers
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					num_attention_heads
													SCREAMING_SNAKE_CASE__     :				str   =					max_position_embeddings
													SCREAMING_SNAKE_CASE__     :				str   =					layer_norm_eps
													SCREAMING_SNAKE_CASE__     :				int   =					hidden_act
													SCREAMING_SNAKE_CASE__     :				List[str]   =					initializer_range
													SCREAMING_SNAKE_CASE__     :				Any   =					attention_probs_dropout_prob
													SCREAMING_SNAKE_CASE__     :				Any   =					is_decoder
													SCREAMING_SNAKE_CASE__     :				Dict   =					use_cache






							@classmethod
							def      A_  (      cls				:     Union[str, Any]							,	a				:     Optional[Any]							,	**a				:     Tuple    )	->"PretrainedConfig":
													cls._set_token_in_kwargs(lowerCAmelCase_    )

													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					cls.get_config_dict(lowerCAmelCase_							,	**lowerCAmelCase_    )

													# get the text config dict if we are loading from BlipConfig
													if config_dict.get("model_type"    ) == "blip":
																			SCREAMING_SNAKE_CASE__     :				List[str]   =					config_dict["text_config"]

													if "model_type" in config_dict and hasattr(cls							,	"model_type"    ) and config_dict["model_type"] != cls.model_type:
																			logger.warning(
																			    f"""You are using a model of type {config_dict['model_type']} to instantiate a model of type """
																			    f"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""    )

													return cls.from_dict(lowerCAmelCase_							,	**lowerCAmelCase_    )







class    _a     (				__lowerCAmelCase							):






							"""simple docstring"""

							snake_case_							 =						'''blip_vision_model'''

							def __init__(      self				:     Any							,	a				:     Tuple=7_68							,	a				:     Union[str, Any]=30_72							,	a				:     List[str]=5_12							,	a				:     Optional[int]=12							,	a				:     Dict=12							,	a				:     Union[str, Any]=3_84							,	a				:     Union[str, Any]=16							,	a				:     Optional[Any]="gelu"							,	a				:     Union[str, Any]=1E-5							,	a				:     Tuple=0.0							,	a				:     int=1E-10							,	**a				:     int							,	)	->Optional[Any]:
													super().__init__(**lowerCAmelCase_    )

													SCREAMING_SNAKE_CASE__     :				int   =					hidden_size
													SCREAMING_SNAKE_CASE__     :				Dict   =					intermediate_size
													SCREAMING_SNAKE_CASE__     :				Tuple   =					projection_dim
													SCREAMING_SNAKE_CASE__     :				str   =					num_hidden_layers
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					num_attention_heads
													SCREAMING_SNAKE_CASE__     :				Tuple   =					patch_size
													SCREAMING_SNAKE_CASE__     :				List[str]   =					image_size
													SCREAMING_SNAKE_CASE__     :				int   =					initializer_range
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					attention_dropout
													SCREAMING_SNAKE_CASE__     :				List[str]   =					layer_norm_eps
													SCREAMING_SNAKE_CASE__     :				int   =					hidden_act






							@classmethod
							def      A_  (      cls				:     Union[str, Any]							,	a				:     int							,	**a				:     Optional[int]    )	->"PretrainedConfig":
													cls._set_token_in_kwargs(lowerCAmelCase_    )

													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				List[Any]   =					cls.get_config_dict(lowerCAmelCase_							,	**lowerCAmelCase_    )

													# get the vision config dict if we are loading from BlipConfig
													if config_dict.get("model_type"    ) == "blip":
																			SCREAMING_SNAKE_CASE__     :				Any   =					config_dict["vision_config"]

													if "model_type" in config_dict and hasattr(cls							,	"model_type"    ) and config_dict["model_type"] != cls.model_type:
																			logger.warning(
																			    f"""You are using a model of type {config_dict['model_type']} to instantiate a model of type """
																			    f"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""    )

													return cls.from_dict(lowerCAmelCase_							,	**lowerCAmelCase_    )







class    _a     (				__lowerCAmelCase							):






							"""simple docstring"""

							snake_case_							 =						'''blip'''
							snake_case_							 =						True

							def __init__(      self				:     List[Any]							,	a				:     Dict=None							,	a				:     Tuple=None							,	a				:     Any=5_12							,	a				:     List[str]=2.6592							,	a				:     Optional[int]=2_56							,	**a				:     List[str]							,	)	->int:
													super().__init__(**lowerCAmelCase_    )

													if text_config is None:
																			SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					{}
																			logger.info("`text_config` is `None`. Initializing the `BlipTextConfig` with default values."    )

													if vision_config is None:
																			SCREAMING_SNAKE_CASE__     :				Dict   =					{}
																			logger.info("`vision_config` is `None`. Initializing the `BlipVisionConfig` with default values."    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					BlipTextConfig(**lowerCAmelCase_    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					BlipVisionConfig(**lowerCAmelCase_    )

													SCREAMING_SNAKE_CASE__     :				List[str]   =					self.vision_config.hidden_size

													SCREAMING_SNAKE_CASE__     :				List[str]   =					projection_dim
													SCREAMING_SNAKE_CASE__     :				Any   =					logit_scale_init_value
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					1.0
													SCREAMING_SNAKE_CASE__     :				int   =					0.02
													SCREAMING_SNAKE_CASE__     :				int   =					image_text_hidden_size

							@classmethod
							def      A_  (      cls				:     List[str]							,	a				:     Optional[Any]							,	a				:     Optional[int]							,	**a				:     List[str]    )	->Dict:
													return cls(text_config=text_config.to_dict()							,	vision_config=vision_config.to_dict()							,	**lowerCAmelCase_    )






							def      A_  (      self				:     Optional[int]    )	->Dict:
													SCREAMING_SNAKE_CASE__     :				str   =					copy.deepcopy(self.__dict__    )
													SCREAMING_SNAKE_CASE__     :				str   =					self.text_config.to_dict()
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					self.vision_config.to_dict()
													SCREAMING_SNAKE_CASE__     :				int   =					self.__class__.model_type
													return output






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






							"""simple docstring"""

							def      A_  (      self				:     Dict    )	->List[Any]:
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()

							def      A_  (      self				:     Dict    )	->Tuple:
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Dict   =					FlaxControlNetModel.from_pretrained(
													    "lllyasviel/sd-controlnet-canny"							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Dict   =					FlaxStableDiffusionControlNetPipeline.from_pretrained(
													    "runwayml/stable-diffusion-v1-5"							,	controlnet=a							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					controlnet_params

													SCREAMING_SNAKE_CASE__     :				Dict   =					"bird"
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.device_count()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe.prepare_text_inputs([prompts] * num_samples    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					load_image(
													    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					pipe.prepare_image_inputs([canny_image] * num_samples    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.random.PRNGKey(0    )
													SCREAMING_SNAKE_CASE__     :				int   =					jax.random.split(a							,	jax.device_count()    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					replicate(a    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					shard(a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					shard(a    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					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, 7_68, 5_12, 3)

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					images[0, 2_53:2_56, 2_53:2_56, -1]

													SCREAMING_SNAKE_CASE__     :				Tuple   =					jnp.asarray(jax.device_get(image_slice.flatten()    )    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					jnp.array(
													    [0.16_7969, 0.11_6699, 0.08_1543, 0.15_4297, 0.13_2812, 0.10_8887, 0.16_9922, 0.16_9922, 0.20_5078]    )
													print(f"""output_slice: {output_slice}"""    )
													assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2






							def      A_  (      self				:     List[Any]    )	->Optional[Any]:
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				int   =					FlaxControlNetModel.from_pretrained(
													    "lllyasviel/sd-controlnet-openpose"							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					FlaxStableDiffusionControlNetPipeline.from_pretrained(
													    "runwayml/stable-diffusion-v1-5"							,	controlnet=a							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					controlnet_params

													SCREAMING_SNAKE_CASE__     :				Any   =					"Chef in the kitchen"
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jax.device_count()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe.prepare_text_inputs([prompts] * num_samples    )

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

													SCREAMING_SNAKE_CASE__     :				Any   =					jax.random.PRNGKey(0    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					jax.random.split(a							,	jax.device_count()    )

													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					replicate(a    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					shard(a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					shard(a    )

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					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, 7_68, 5_12, 3)

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
													SCREAMING_SNAKE_CASE__     :				str   =					images[0, 2_53:2_56, 2_53:2_56, -1]

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.asarray(jax.device_get(image_slice.flatten()    )    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.array(
													    [[0.27_1484, 0.26_1719, 0.27_5391, 0.27_7344, 0.27_9297, 0.29_1016, 0.29_4922, 0.30_2734, 0.30_2734]]    )
													print(f"""output_slice: {output_slice}"""    )
													assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2