




"""simple docstring"""

import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
       from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image





class 							lowercase_			:



 '''simple docstring'''

 def __init__(					self			:			Tuple       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			List[Any]=3       ,				_UpperCAmelCase			:			List[Any]=32       ,				_UpperCAmelCase			:			Optional[Any]=3       ,				_UpperCAmelCase			:			int=10       ,				_UpperCAmelCase			:			int=[8, 16, 32, 64]       ,				_UpperCAmelCase			:			Optional[Any]=[1, 1, 2, 1]       ,				_UpperCAmelCase			:			Optional[Any]=True       ,				_UpperCAmelCase			:			Union[str, Any]=True       ,				_UpperCAmelCase			:			Union[str, Any]="relu"       ,				_UpperCAmelCase			:			Optional[int]=3       ,				_UpperCAmelCase			:			Tuple=None       ,				_UpperCAmelCase			:			Dict=["stage2", "stage3", "stage4"]       ,				_UpperCAmelCase			:			Union[str, Any]=[2, 3, 4]       ,				_UpperCAmelCase			:			Optional[int]=1       ,				):
       _A     =  parent
       _A     =  batch_size
       _A     =  image_size
       _A     =  num_channels
       _A     =  embeddings_size
       _A     =  hidden_sizes
       _A     =  depths
       _A     =  is_training
       _A     =  use_labels
       _A     =  hidden_act
       _A     =  num_labels
       _A     =  scope
       _A     =  len(_UpperCAmelCase						)
       _A     =  out_features
       _A     =  out_indices
       _A     =  num_groups

 def      lowerCAmelCase_				(					self			:			str						):
       _A     =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

       _A     =  None
       if self.use_labels:
             _A     =  ids_tensor([self.batch_size]       ,				self.num_labels						)

       _A     =  self.get_config()

       return config, pixel_values, labels

 def      lowerCAmelCase_				(					self			:			int						):
       return BitConfig(
           num_channels=self.num_channels       ,				embeddings_size=self.embeddings_size       ,				hidden_sizes=self.hidden_sizes       ,				depths=self.depths       ,				hidden_act=self.hidden_act       ,				num_labels=self.num_labels       ,				out_features=self.out_features       ,				out_indices=self.out_indices       ,				num_groups=self.num_groups       ,				)

 def      lowerCAmelCase_				(					self			:			List[Any]       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			List[str]						):
       _A     =  BitModel(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(_UpperCAmelCase						)
       self.parent.assertEqual(
           result.last_hidden_state.shape       ,				(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,				)

 def      lowerCAmelCase_				(					self			:			Tuple       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			Tuple						):
       _A     =  self.num_labels
       _A     =  BitForImageClassification(_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(_UpperCAmelCase       ,				labels=_UpperCAmelCase						)
       self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_labels)						)

 def      lowerCAmelCase_				(					self			:			List[str]       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			int						):
       _A     =  BitBackbone(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(_UpperCAmelCase						)

       # verify feature maps
       self.parent.assertEqual(len(result.feature_maps						)       ,				len(config.out_features						)						)
       self.parent.assertListEqual(list(result.feature_maps[0].shape						)       ,				[self.batch_size, self.hidden_sizes[1], 4, 4]						)

       # verify channels
       self.parent.assertEqual(len(model.channels						)       ,				len(config.out_features						)						)
       self.parent.assertListEqual(model.channels       ,				config.hidden_sizes[1:]						)

       # verify backbone works with out_features=None
       _A     =  None
       _A     =  BitBackbone(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(_UpperCAmelCase						)

       # verify feature maps
       self.parent.assertEqual(len(result.feature_maps						)       ,				1						)
       self.parent.assertListEqual(list(result.feature_maps[0].shape						)       ,				[self.batch_size, self.hidden_sizes[-1], 1, 1]						)

       # verify channels
       self.parent.assertEqual(len(model.channels						)       ,				1						)
       self.parent.assertListEqual(model.channels       ,				[config.hidden_sizes[-1]]						)



 def      lowerCAmelCase_				(					self			:			Union[str, Any]						):
       _A     =  self.prepare_config_and_inputs()
       _A  ,       _A  ,       _A     =  config_and_inputs
       _A     =  {'pixel_values': pixel_values}
       return config, inputs_dict





@require_torch
class 							lowercase_			(	__lowerCAmelCase     ,      __lowerCAmelCase     ,      unittest.TestCase				):



 '''simple docstring'''
 UpperCAmelCase       :	str     =      (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
 UpperCAmelCase       :	Tuple     =      (
     {'''feature-extraction''': BitModel, '''image-classification''': BitForImageClassification}
     if is_torch_available()
     else {}
 )

 UpperCAmelCase       :	Optional[int]     =      False
 UpperCAmelCase       :	Any     =      False
 UpperCAmelCase       :	List[str]     =      False
 UpperCAmelCase       :	List[Any]     =      False
 UpperCAmelCase       :	str     =      False

 def      lowerCAmelCase_				(					self			:			int						):
       _A     =  BitModelTester(self						)
       _A     =  ConfigTester(self       ,				config_class=_UpperCAmelCase       ,				has_text_modality=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Optional[Any]						):
       self.create_and_test_config_common_properties()
       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      lowerCAmelCase_				(					self			:			Any						):
       return

 @unittest.skip(reason='Bit does not output attentions'						)
 def      lowerCAmelCase_				(					self			:			List[str]						):
       pass

 @unittest.skip(reason='Bit does not use inputs_embeds'						)
 def      lowerCAmelCase_				(					self			:			List[Any]						):
       pass

 @unittest.skip(reason='Bit does not support input and output embeddings'						)
 def      lowerCAmelCase_				(					self			:			str						):
       pass

 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()

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

             _A     =  ['pixel_values']
             self.assertListEqual(arg_names[:1]       ,				_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			str						):
       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             _A     =  model_class(config=_UpperCAmelCase						)
             for name, module in model.named_modules():
                   if isinstance(_UpperCAmelCase       ,				(nn.BatchNormad, nn.GroupNorm)						):
                         self.assertTrue(
                             torch.all(module.weight == 1						)       ,				msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''       ,				)
                         self.assertTrue(
                             torch.all(module.bias == 0						)       ,				msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''       ,				)

 def      lowerCAmelCase_				(					self			:			List[str]						):
       def check_hidden_states_output(_UpperCAmelCase			:			Any       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			Any						):
             _A     =  model_class(_UpperCAmelCase						)
             model.to(_UpperCAmelCase						)
             model.eval()

             with torch.no_grad():
                   _A     =  model(**self._prepare_for_class(_UpperCAmelCase       ,				_UpperCAmelCase						)						)

             _A     =  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

             _A     =  self.model_tester.num_stages
             self.assertEqual(len(_UpperCAmelCase						)       ,				expected_num_stages + 1						)

             # Bit's feature maps are of shape (batch_size, num_channels, height, width)
             self.assertListEqual(
                 list(hidden_states[0].shape[-2:]						)       ,				[self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,				)

       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()
       _A     =  ['preactivation', 'bottleneck']
       for model_class in self.all_model_classes:
             for layer_type in layers_type:
                   _A     =  layer_type
                   _A     =  True
                   check_hidden_states_output(_UpperCAmelCase       ,				_UpperCAmelCase       ,				_UpperCAmelCase						)

                   # check that output_hidden_states also work using config
                   del inputs_dict["output_hidden_states"]
                   _A     =  True

                   check_hidden_states_output(_UpperCAmelCase       ,				_UpperCAmelCase       ,				_UpperCAmelCase						)

 @unittest.skip(reason='Bit does not use feedforward chunking'						)
 def      lowerCAmelCase_				(					self			:			int						):
       pass

 def      lowerCAmelCase_				(					self			:			List[str]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase						)



 @slow
 def      lowerCAmelCase_				(					self			:			List[str]						):
       for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _A     =  BitModel.from_pretrained(_UpperCAmelCase						)
             self.assertIsNotNone(_UpperCAmelCase						)






def 		_snake_case	(   )       ->   str:






      '''simple docstring'''


      _A     =  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
      return image





@require_torch
@require_vision
class 							lowercase_			(	unittest.TestCase				):



 '''simple docstring'''

 @cached_property
 def      lowerCAmelCase_				(					self			:			List[str]						):
       return (
           BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						) if is_vision_available() else None
       )



 @slow
 def      lowerCAmelCase_				(					self			:			Tuple						):
       _A     =  BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						).to(_UpperCAmelCase						)

       _A     =  self.default_image_processor
       _A     =  prepare_img()
       _A     =  image_processor(images=_UpperCAmelCase       ,				return_tensors='pt'						).to(_UpperCAmelCase						)

       # forward pass
       with torch.no_grad():
             _A     =  model(**_UpperCAmelCase						)

       # verify the logits
       _A     =  torch.Size((1, 1_000)						)
       self.assertEqual(outputs.logits.shape       ,				_UpperCAmelCase						)

       _A     =  torch.tensor([[-0.6526, -0.5263, -1.4398]]						).to(_UpperCAmelCase						)

       self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,				_UpperCAmelCase       ,				atol=1E-4						)						)





@require_torch
class 							lowercase_			(	__lowerCAmelCase     ,      unittest.TestCase				):



 '''simple docstring'''
 UpperCAmelCase       :	Tuple     =      (BitBackbone,) if is_torch_available() else ()
 UpperCAmelCase       :	List[Any]     =      BitConfig

 UpperCAmelCase       :	Optional[int]     =      False

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  BitModelTester(self						)





'''simple docstring'''

import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
 lowerCamelCase				:  int    		=      argparse.ArgumentParser()

 parser.add_argument(
     '--checkpoint_path', default=None, type=str, required=True, help='Path to the checkpoint to convert.'
 )
 # !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
 parser.add_argument(
     '--original_config_file',
     default=None,
     type=str,
     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(
     '--scheduler_type',
     default='pndm',
     type=str,
     help='Type of scheduler to use. Should be one of [\'pndm\', \'lms\', \'ddim\', \'euler\', \'euler-ancestral\', \'dpm\']',
 )
 parser.add_argument(
     '--pipeline_type',
     default=None,
     type=str,
     help=(
         'The pipeline type. One of \'FrozenOpenCLIPEmbedder\', \'FrozenCLIPEmbedder\', \'PaintByExample\''
         '. If `None` pipeline will be automatically inferred.'
     ),
 )
 parser.add_argument(
     '--image_size',
     default=None,
     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(
     '--prediction_type',
     default=None,
     type=str,
     help=(
         'The prediction type that the model was trained on. Use \'epsilon\' for Stable Diffusion v1.X and Stable'
         ' Diffusion v2 Base. Use \'v_prediction\' 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.)')
 parser.add_argument(
     '--stable_unclip',
     type=str,
     default=None,
     required=False,
     help='Set if this is a stable unCLIP model. One of \'txt2img\' or \'img2img\'.',
 )
 parser.add_argument(
     '--stable_unclip_prior',
     type=str,
     default=None,
     required=False,
     help='Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.',
 )
 parser.add_argument(
     '--clip_stats_path',
     type=str,
     help='Path to the clip stats file. Only required if the stable unclip model\'s config specifies `model.params.noise_aug_config.params.clip_stats_path`.',
     required=False,
 )
 parser.add_argument(
     '--controlnet', action='store_true', default=None, help='Set flag if this is a controlnet checkpoint.'
 )
 parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
 parser.add_argument(
     '--vae_path',
     type=str,
     default=None,
     required=False,
     help='Set to a path, hub id to an already converted vae to not convert it again.',
 )
 lowerCamelCase				:  str    		=      parser.parse_args()

 lowerCamelCase				:  Tuple    		=      download_from_original_stable_diffusion_ckpt(
     checkpoint_path=args.checkpoint_path,
     original_config_file=args.original_config_file,
     image_size=args.image_size,
     prediction_type=args.prediction_type,
     model_type=args.pipeline_type,
     extract_ema=args.extract_ema,
     scheduler_type=args.scheduler_type,
     num_in_channels=args.num_in_channels,
     upcast_attention=args.upcast_attention,
     from_safetensors=args.from_safetensors,
     device=args.device,
     stable_unclip=args.stable_unclip,
     stable_unclip_prior=args.stable_unclip_prior,
     clip_stats_path=args.clip_stats_path,
     controlnet=args.controlnet,
     vae_path=args.vae_path,
 )

 if args.half:
  pipe.to(torch_dtype=torch.floataa)

 if args.controlnet:
  # only save the controlnet model
  pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
 else:
  pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)




