
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        __UpperCamelCase					(    lowerCAmelCase__							:    List[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)


import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


lowercase__          =50000
lowercase__          =5000

lowercase__       ,      lowercase__          =os.path.split(__file__)
lowercase__          =os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))




@get_duration
def        __UpperCamelCase					(    lowerCAmelCase__							:    datasets.Dataset  ,       lowerCAmelCase__							:    List[str]  ):
 for i in range(lowerCAmelCase__  ):
  __a			:  str		    =  dataset[i]




@get_duration
def        __UpperCamelCase					(    lowerCAmelCase__							:    datasets.Dataset  ,       lowerCAmelCase__							:    Optional[Any]  ,       lowerCAmelCase__							:    Tuple  ):
 for i in range(0  ,       len(lowerCAmelCase__  )  ,       lowerCAmelCase__  ):
  __a			:  Optional[int]		    =  dataset[i : i + batch_size]




@get_duration
def        __UpperCamelCase					(    lowerCAmelCase__							:    datasets.Dataset  ,       lowerCAmelCase__							:    Union[str, Any]  ,       lowerCAmelCase__							:    str  ):
 with dataset.formatted_as(type=lowerCAmelCase__  ):
  for i in range(lowerCAmelCase__  ):
   __a			:  Dict		    =  dataset[i]




@get_duration
def        __UpperCamelCase					(    lowerCAmelCase__							:    datasets.Dataset  ,       lowerCAmelCase__							:    List[str]  ,       lowerCAmelCase__							:    Optional[Any]  ,       lowerCAmelCase__							:    int  ):
 with dataset.formatted_as(type=lowerCAmelCase__  ):
  for i in range(0  ,       lowerCAmelCase__  ,       lowerCAmelCase__  ):
   __a			:  int		    =  dataset[i : i + batch_size]




def        __UpperCamelCase					(    ):
 __a			:  Any		    =  {'''num examples''': SPEED_TEST_N_EXAMPLES}
 __a			:  List[Any]		    =  [
     (read, {'''length''': SMALL_TEST}),
     (read, {'''length''': SPEED_TEST_N_EXAMPLES}),
     (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1_0}),
     (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1_0_0}),
     (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1_0_0_0}),
     (read_formatted, {'''type''': '''numpy''', '''length''': SMALL_TEST}),
     (read_formatted, {'''type''': '''pandas''', '''length''': SMALL_TEST}),
     (read_formatted, {'''type''': '''torch''', '''length''': SMALL_TEST}),
     (read_formatted, {'''type''': '''tensorflow''', '''length''': SMALL_TEST}),
     (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 1_0}),
     (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 1_0_0_0}),
 ]

 __a			:  Union[str, Any]		    =  [
     (read, {'''length''': SMALL_TEST}),
     (read, {'''length''': SPEED_TEST_N_EXAMPLES}),
     (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1_0}),
     (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1_0_0}),
     (read_batch, {'''length''': SPEED_TEST_N_EXAMPLES, '''batch_size''': 1_0_0_0}),
     (read_formatted, {'''type''': '''numpy''', '''length''': SMALL_TEST}),
     (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 1_0}),
     (read_formatted_batch, {'''type''': '''numpy''', '''length''': SMALL_TEST, '''batch_size''': 1_0_0_0}),
 ]
 with tempfile.TemporaryDirectory() as tmp_dir:
  print('''generating dataset'''  )
  __a			:  Optional[Any]		    =  datasets.Features(
      {'''list''': datasets.Sequence(datasets.Value('''float32'''  )  ), '''numbers''': datasets.Value('''float32'''  )}  )
  __a			:  Optional[int]		    =  generate_example_dataset(
      os.path.join(lowerCAmelCase__  ,       '''dataset.arrow'''  )  ,       lowerCAmelCase__  ,       num_examples=lowerCAmelCase__  ,       seq_shapes={'''list''': (1_0_0,)}  ,       )
  print('''first set of iterations'''  )
  for func, kwargs in functions:
   print(func.__name__  ,       str(lowerCAmelCase__  )  )
   __a			:  str		    =  func(lowerCAmelCase__  ,       **lowerCAmelCase__  )

  print('''shuffling dataset'''  )
  __a			:  int		    =  dataset.shuffle()
  print('''Second set of iterations (after shuffling'''  )
  for func, kwargs in functions_shuffled:
   print('''shuffled '''  ,       func.__name__  ,       str(lowerCAmelCase__  )  )
   __a			:  List[Any]		    =  func(
       lowerCAmelCase__  ,       **lowerCAmelCase__  )

 with open(lowerCAmelCase__  ,       '''wb'''  ) as f:
  f.write(json.dumps(lowerCAmelCase__  ).encode('''utf-8'''  )  )


if __name__ == "__main__":  # useful to run the profiler
      benchmark_iterating()
