



from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule



def 		lowerCAmelCase__    (					UpperCamelCase_				:       Tuple ,       UpperCamelCase_				:       str ,       UpperCamelCase_				:       str=1E-12	)->  List[str]:
    A__		 =					jnp.divide(emb_a.T ,       jnp.clip(jnp.linalg.norm(UpperCamelCase_ ,       axis=1	) ,       a_min=UpperCamelCase_	)	).T
    A__		 =					jnp.divide(emb_a.T ,       jnp.clip(jnp.linalg.norm(UpperCamelCase_ ,       axis=1	) ,       a_min=UpperCamelCase_	)	).T
    return jnp.matmul(UpperCamelCase_ ,       norm_emb_a.T	)



class     _UpperCAmelCase			(       nn.Module     ):
     UpperCamelCase__		    =				42
     UpperCamelCase__		    =				jnp.floataa




     def 		snake_case_		(				self):
         A__		 =					FlaxCLIPVisionModule(self.config.vision_config)
         A__		 =					nn.Dense(self.config.projection_dim ,						use_bias=a__ ,						dtype=self.dtype)

         A__		 =					self.param('''concept_embeds''' ,						jax.nn.initializers.ones ,						(1_7, self.config.projection_dim))
         A__		 =					self.param(
             '''special_care_embeds''' ,						jax.nn.initializers.ones ,						(3, self.config.projection_dim))

         A__		 =					self.param('''concept_embeds_weights''' ,						jax.nn.initializers.ones ,						(1_7,))
         A__		 =					self.param('''special_care_embeds_weights''' ,						jax.nn.initializers.ones ,						(3,))




     def __call__(				self ,						a__):
         A__		 =					self.vision_model(a__)[1]
         A__		 =					self.visual_projection(a__)

         A__		 =					jax_cosine_distance(a__ ,						self.special_care_embeds)
         A__		 =					jax_cosine_distance(a__ ,						self.concept_embeds)

         # increase this value to create a stronger `nfsw` filter
         # at the cost of increasing the possibility of filtering benign image inputs
         A__		 =					0.0

         A__		 =					special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
         A__		 =					jnp.round(a__ ,						3)
         A__		 =					jnp.any(special_scores > 0 ,						axis=1 ,						keepdims=a__)
         # Use a lower threshold if an image has any special care concept
         A__		 =					is_special_care * 0.0_1

         A__		 =					cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
         A__		 =					jnp.round(a__ ,						3)
         A__		 =					jnp.any(concept_scores > 0 ,						axis=1)

         return has_nsfw_concepts







class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				CLIPConfig
     UpperCamelCase__		    =				'''clip_input'''
     UpperCamelCase__		    =				FlaxStableDiffusionSafetyCheckerModule




     def __init__(				self ,						a__ ,						a__ = None ,						a__ = 0 ,						a__ = jnp.floataa ,						a__ = True ,						**a__ ,						):
         if input_shape is None:
             A__		 =					(1, 2_2_4, 2_2_4, 3)
         A__		 =					self.module_class(config=a__ ,						dtype=a__ ,						**a__)
         super().__init__(a__ ,						a__ ,						input_shape=a__ ,						seed=a__ ,						dtype=a__ ,						_do_init=_do_init)




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ = None):
         # init input tensor
         A__		 =					jax.random.normal(a__ ,						a__)

         A__ ,	A__		 =					jax.random.split(a__)
         A__		 =					{'''params''': params_rng, '''dropout''': dropout_rng}

         A__		 =					self.module.init(a__ ,						a__)['''params''']

         return random_params




     def __call__(				self ,						a__ ,						a__ = None ,						):
         A__		 =					jnp.transpose(a__ ,						(0, 2, 3, 1))

         return self.module.apply(
             {'''params''': params or self.params} ,						jnp.array(a__ ,						dtype=jnp.floataa) ,						rngs={} ,						)





import argparse

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_controlnet_from_original_ckpt


if __name__ == "__main__":
   _lowercase						  =		argparse.ArgumentParser()

   parser.add_argument(
       "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert."
   )
   parser.add_argument(
       "--original_config_file",
       type=str,
       required=True,
       help="The YAML config file corresponding to the original architecture.",
   )
   parser.add_argument(
       "--num_in_channels",
       default=None,
       type=int,
       help="The number of input channels. If `None` number of input channels will be automatically inferred.",
   )
   parser.add_argument(
       "--image_size",
       default=512,
       type=int,
       help=(
           "The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2"
           " Base. Use 768 for Stable Diffusion v2."
       ),
   )
   parser.add_argument(
       "--extract_ema",
       action="store_true",
       help=(
           "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights"
           " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield"
           " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning."
       ),
   )
   parser.add_argument(
       "--upcast_attention",
       action="store_true",
       help=(
           "Whether the attention computation should always be upcasted. This is necessary when running stable"
           " diffusion 2.1."
       ),
   )
   parser.add_argument(
       "--from_safetensors",
       action="store_true",
       help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.",
   )
   parser.add_argument(
       "--to_safetensors",
       action="store_true",
       help="Whether to store pipeline in safetensors format or not.",
   )
   parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")
   parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)")



   def 		lowerCAmelCase__    (					UpperCamelCase_				:       Any	)->  Union[str, Any]:
       if string == "True":
           return True
       elif string == "False":
           return False
       else:
           raise ValueError(f"could not parse string as bool {string}"	)

   parser.add_argument(
       "--use_linear_projection", help="Override for use linear projection", required=False, type=parse_bool
   )

   parser.add_argument("--cross_attention_dim", help="Override for cross attention_dim", required=False, type=int)

   _lowercase						  =		parser.parse_args()

   _lowercase						  =		download_controlnet_from_original_ckpt(
       checkpoint_path=args.checkpoint_path,
       original_config_file=args.original_config_file,
       image_size=args.image_size,
       extract_ema=args.extract_ema,
       num_in_channels=args.num_in_channels,
       upcast_attention=args.upcast_attention,
       from_safetensors=args.from_safetensors,
       device=args.device,
       use_linear_projection=args.use_linear_projection,
       cross_attention_dim=args.cross_attention_dim,
   )



   controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
