'''simple docstring'''

import copy
import os
from typing import Union

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


lowerCAmelCase_ : List[str]						  =      logging.get_logger(__name__)

lowerCAmelCase_ : Optional[int]						  =      {
    "BAAI/AltCLIP": "https://huggingface.co/BAAI/AltCLIP/resolve/main/config.json",
    # See all AltCLIP models at https://huggingface.co/models?filter=altclip
}





class 							lowercase						(						__lowerCamelCase ):
   lowerCamelCase_    					='altclip_text_model'


   def __init__(   self				:	Any   ,		__lowerCAmelCase				:	List[str]=25_0002   ,		__lowerCAmelCase				:	Union[str, Any]=1024   ,		__lowerCAmelCase				:	Dict=24   ,		__lowerCAmelCase				:	int=16   ,		__lowerCAmelCase				:	str=4096   ,		__lowerCAmelCase				:	List[Any]="gelu"   ,		__lowerCAmelCase				:	Union[str, Any]=0.1   ,		__lowerCAmelCase				:	Tuple=0.1   ,		__lowerCAmelCase				:	Tuple=514   ,		__lowerCAmelCase				:	Optional[int]=1   ,		__lowerCAmelCase				:	Optional[Any]=0.02   ,		__lowerCAmelCase				:	Optional[int]=0.02   ,		__lowerCAmelCase				:	Optional[Any]=1e-05   ,		__lowerCAmelCase				:	List[str]=1   ,		__lowerCAmelCase				:	Dict=0   ,		__lowerCAmelCase				:	List[str]=2   ,		__lowerCAmelCase				:	Optional[int]="absolute"   ,		__lowerCAmelCase				:	Optional[int]=True   ,		__lowerCAmelCase				:	Union[str, Any]=768   ,		**__lowerCAmelCase				:	List[str]   ,		)      ->					Dict:
         super().__init__(pad_token_id=__lowerCAmelCase   ,		bos_token_id=__lowerCAmelCase   ,		eos_token_id=__lowerCAmelCase   ,		**__lowerCAmelCase)

         lowercase_     =      vocab_size
         lowercase_     =      hidden_size
         lowercase_     =      num_hidden_layers
         lowercase_     =      num_attention_heads
         lowercase_     =      hidden_act
         lowercase_     =      intermediate_size
         lowercase_     =      hidden_dropout_prob
         lowercase_     =      attention_probs_dropout_prob
         lowercase_     =      max_position_embeddings
         lowercase_     =      type_vocab_size
         lowercase_     =      initializer_range
         lowercase_     =      initializer_factor
         lowercase_     =      layer_norm_eps
         lowercase_     =      position_embedding_type
         lowercase_     =      use_cache
         lowercase_     =      project_dim





class 							lowercase						(						__lowerCamelCase ):
   lowerCamelCase_    					='altclip_vision_model'


   def __init__(   self				:	List[str]   ,		__lowerCAmelCase				:	Optional[Any]=768   ,		__lowerCAmelCase				:	Union[str, Any]=3072   ,		__lowerCAmelCase				:	List[Any]=512   ,		__lowerCAmelCase				:	Optional[Any]=12   ,		__lowerCAmelCase				:	Optional[Any]=12   ,		__lowerCAmelCase				:	Optional[Any]=3   ,		__lowerCAmelCase				:	List[str]=224   ,		__lowerCAmelCase				:	Optional[Any]=32   ,		__lowerCAmelCase				:	str="quick_gelu"   ,		__lowerCAmelCase				:	List[str]=1e-5   ,		__lowerCAmelCase				:	Optional[Any]=0.0   ,		__lowerCAmelCase				:	Tuple=0.02   ,		__lowerCAmelCase				:	str=1.0   ,		**__lowerCAmelCase				:	List[str]   ,		)      ->					Union[str, Any]:
         super().__init__(**__lowerCAmelCase)

         lowercase_     =      hidden_size
         lowercase_     =      intermediate_size
         lowercase_     =      projection_dim
         lowercase_     =      num_hidden_layers
         lowercase_     =      num_attention_heads
         lowercase_     =      num_channels
         lowercase_     =      patch_size
         lowercase_     =      image_size
         lowercase_     =      initializer_range
         lowercase_     =      initializer_factor
         lowercase_     =      attention_dropout
         lowercase_     =      layer_norm_eps
         lowercase_     =      hidden_act


   @classmethod
   def __UpperCAmelCase       (   cls				:	Dict   ,		__lowerCAmelCase				:	Union[str, os.PathLike]   ,		**__lowerCAmelCase				:	Union[str, Any])      ->					"PretrainedConfig":
         cls._set_token_in_kwargs(__lowerCAmelCase)

         lowercase_ ,							lowercase_     =      cls.get_config_dict(__lowerCAmelCase   ,		**__lowerCAmelCase)

         # get the vision config dict if we are loading from AltCLIPConfig
         if config_dict.get("model_type") == "altclip":
               lowercase_     =      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 							lowercase						(						__lowerCamelCase ):
   lowerCamelCase_    					='altclip'
   lowerCamelCase_    					=True


   def __init__(   self				:	Tuple   ,		__lowerCAmelCase				:	List[Any]=None   ,		__lowerCAmelCase				:	List[Any]=None   ,		__lowerCAmelCase				:	Tuple=768   ,		__lowerCAmelCase				:	List[Any]=2.6592   ,		**__lowerCAmelCase				:	Union[str, Any])      ->					List[str]:
         # If `_config_dict` exist, we use them for the backward compatibility.
         # We pop out these 2 attributes before calling `super().__init__` to avoid them being saved (which causes a lot
         # of confusion!).
         lowercase_     =      kwargs.pop("text_config_dict"   ,		__lowerCAmelCase)
         lowercase_     =      kwargs.pop("vision_config_dict"   ,		__lowerCAmelCase)

         super().__init__(**__lowerCAmelCase)

         # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
         # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
         # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
         if text_config_dict is not None:
               if text_config is None:
                     lowercase_     =      {}

               # This is the complete result when using `text_config_dict`.
               lowercase_     =      AltCLIPTextConfig(**__lowerCAmelCase).to_dict()

               # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
               for key, value in _text_config_dict.items():
                     if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
                           # If specified in `text_config_dict`
                           if key in text_config_dict:
                                 lowercase_     =      (
                                     F'`{key}` is found in both `text_config_dict` and `text_config` but with different values. '
                                     F'The value `text_config_dict["{key}"]` will be used instead.'
                                 )
                           # If inferred from default argument values (just to be super careful)
                           else:
                                 lowercase_     =      (
                                     F'`text_config_dict` is provided which will be used to initialize `AltCLIPTextConfig`. The '
                                     F'value `text_config["{key}"]` will be overriden.'
                                 )
                           logger.warning(__lowerCAmelCase)

            # Update all values in `text_config` with the ones in `_text_config_dict`.
               text_config.update(_text_config_dict)

         if vision_config_dict is not None:
               if vision_config is None:
                     lowercase_     =      {}

               # This is the complete result when using `vision_config_dict`.
               lowercase_     =      AltCLIPVisionConfig(**__lowerCAmelCase).to_dict()
               # convert keys to string instead of integer
               if "id2label" in _vision_config_dict:
                     lowercase_     =      {
                         str(__lowerCAmelCase): value for key, value in _vision_config_dict["id2label"].items()
                     }

               # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
               for key, value in _vision_config_dict.items():
                     if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
                           # If specified in `vision_config_dict`
                           if key in vision_config_dict:
                                 lowercase_     =      (
                                     F'`{key}` is found in both `vision_config_dict` and `vision_config` but with different '
                                     F'values. The value `vision_config_dict["{key}"]` will be used instead.'
                                 )
                           # If inferred from default argument values (just to be super careful)
                           else:
                                 lowercase_     =      (
                                     F'`vision_config_dict` is provided which will be used to initialize `AltCLIPVisionConfig`. '
                                     F'The value `vision_config["{key}"]` will be overriden.'
                                 )
                           logger.warning(__lowerCAmelCase)

            # Update all values in `vision_config` with the ones in `_vision_config_dict`.
               vision_config.update(_vision_config_dict)

         if text_config is None:
               lowercase_     =      {}
               logger.info("`text_config` is `None`. Initializing the `AltCLIPTextConfig` with default values.")

         if vision_config is None:
               lowercase_     =      {}
               logger.info("`vision_config` is `None`. initializing the `AltCLIPVisionConfig` with default values.")

         lowercase_     =      AltCLIPTextConfig(**__lowerCAmelCase)
         lowercase_     =      AltCLIPVisionConfig(**__lowerCAmelCase)

         lowercase_     =      projection_dim
         lowercase_     =      logit_scale_init_value
         lowercase_     =      1.0


   @classmethod
   def __UpperCAmelCase       (   cls				:	str   ,		__lowerCAmelCase				:	AltCLIPTextConfig   ,		__lowerCAmelCase				:	AltCLIPVisionConfig   ,		**__lowerCAmelCase				:	str)      ->					str:
         return cls(text_config=text_config.to_dict()   ,		vision_config=vision_config.to_dict()   ,		**__lowerCAmelCase)


   def __UpperCAmelCase       (   self				:	Union[str, Any])      ->					List[str]:
         lowercase_     =      copy.deepcopy(self.__dict__)
         lowercase_     =      self.text_config.to_dict()
         lowercase_     =      self.vision_config.to_dict()
         lowercase_     =      self.__class__.model_type
         return output







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


   def __UpperCAmelCase       (   self				:	int)      ->					List[str]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()


   def __UpperCAmelCase       (   self				:	Dict)      ->					Any:
         lowercase_ ,							lowercase_     =      FlaxControlNetModel.from_pretrained(
             "lllyasviel/sd-controlnet-canny"   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_ ,							lowercase_     =      FlaxStableDiffusionControlNetPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5"   ,		controlnet=__lowerCAmelCase   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_     =      controlnet_params

         lowercase_     =      "bird"
         lowercase_     =      jax.device_count()
         lowercase_     =      pipe.prepare_text_inputs([prompts] * num_samples)

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

         lowercase_     =      jax.random.PRNGKey(0)
         lowercase_     =      jax.random.split(__lowerCAmelCase   ,		jax.device_count())

         lowercase_     =      replicate(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)

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

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

         lowercase_     =      jnp.asarray(jax.device_get(image_slice.flatten()))
         lowercase_     =      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 __UpperCAmelCase       (   self				:	Tuple)      ->					List[str]:
         lowercase_ ,							lowercase_     =      FlaxControlNetModel.from_pretrained(
             "lllyasviel/sd-controlnet-openpose"   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_ ,							lowercase_     =      FlaxStableDiffusionControlNetPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5"   ,		controlnet=__lowerCAmelCase   ,		from_pt=__lowerCAmelCase   ,		dtype=jnp.bfloataa)
         lowercase_     =      controlnet_params

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

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

         lowercase_     =      jax.random.PRNGKey(0)
         lowercase_     =      jax.random.split(__lowerCAmelCase   ,		jax.device_count())

         lowercase_     =      replicate(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)
         lowercase_     =      shard(__lowerCAmelCase)

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

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

         lowercase_     =      jnp.asarray(jax.device_get(image_slice.flatten()))
         lowercase_     =      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






