

from io import BytesIO
from typing import List, Union

import requests

from ..utils import add_end_docstrings, is_decord_available, is_torch_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_decord_available():
    import numpy as np
    from decord import VideoReader


if is_torch_available():
    from ..models.auto.modeling_auto import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING

_SCREAMING_SNAKE_CASE					:  int             =      logging.get_logger(__name__)






@add_end_docstrings(A__	)
class 	UpperCAmelCase__		(	A__	):






  """simple docstring"""


  def __init__(							self						:		Tuple     ,       *__lowerCamelCase						:		Dict     ,       **__lowerCamelCase						:		List[str] )		->							List[Any]:
       super().__init__(*__lowerCamelCase     ,       **__lowerCamelCase )
       requires_backends(self     ,       '''decord''' )
       self.check_model_type(__lowerCamelCase )






  def 				lowercase_		(							self						:		Any     ,       __lowerCamelCase						:		List[str]=None     ,       __lowerCamelCase						:		Union[str, Any]=None     ,       __lowerCamelCase						:		int=None )		->							List[Any]:
       SCREAMING_SNAKE_CASE__   =   {}
       if frame_sampling_rate is not None:
            SCREAMING_SNAKE_CASE__   =   frame_sampling_rate
       if num_frames is not None:
            SCREAMING_SNAKE_CASE__   =   num_frames

       SCREAMING_SNAKE_CASE__   =   {}
       if top_k is not None:
            SCREAMING_SNAKE_CASE__   =   top_k
       return preprocess_params, {}, postprocess_params






  def __call__(							self						:		Tuple     ,       __lowerCamelCase						:		Union[str, List[str]]     ,       **__lowerCamelCase						:		int )		->							Tuple:
       return super().__call__(__lowerCamelCase     ,       **__lowerCamelCase )






  def 				lowercase_		(							self						:		Dict     ,       __lowerCamelCase						:		str     ,       __lowerCamelCase						:		Optional[int]=None     ,       __lowerCamelCase						:		Optional[int]=1 )		->							List[str]:
       if num_frames is None:
            SCREAMING_SNAKE_CASE__   =   self.model.config.num_frames

       if video.startswith('''http://''' ) or video.startswith('''https://''' ):
            SCREAMING_SNAKE_CASE__   =   BytesIO(requests.get(__lowerCamelCase ).content )

       SCREAMING_SNAKE_CASE__   =   VideoReader(__lowerCamelCase )
       videoreader.seek(0 )

       SCREAMING_SNAKE_CASE__   =   0
       SCREAMING_SNAKE_CASE__   =   num_frames * frame_sampling_rate - 1
       SCREAMING_SNAKE_CASE__   =   np.linspace(__lowerCamelCase     ,       __lowerCamelCase     ,       num=__lowerCamelCase     ,       dtype=np.intaa )

       SCREAMING_SNAKE_CASE__   =   videoreader.get_batch(__lowerCamelCase ).asnumpy()
       SCREAMING_SNAKE_CASE__   =   list(__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   self.image_processor(__lowerCamelCase     ,       return_tensors=self.framework )
       return model_inputs






  def 				lowercase_		(							self						:		str     ,       __lowerCamelCase						:		Dict )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   self.model(**__lowerCamelCase )
       return model_outputs







  def 				lowercase_		(							self						:		Union[str, Any]     ,       __lowerCamelCase						:		List[str]     ,       __lowerCamelCase						:		Optional[Any]=5 )		->							Union[str, Any]:
       if top_k > self.model.config.num_labels:
            SCREAMING_SNAKE_CASE__   =   self.model.config.num_labels

       if self.framework == "pt":
            SCREAMING_SNAKE_CASE__   =   model_outputs.logits.softmax(-1 )[0]
            SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   probs.topk(__lowerCamelCase )
       else:
            raise ValueError(f'''Unsupported framework: {self.framework}''' )

       SCREAMING_SNAKE_CASE__   =   scores.tolist()
       SCREAMING_SNAKE_CASE__   =   ids.tolist()
       return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(__lowerCamelCase     ,       __lowerCamelCase )]





import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


_SCREAMING_SNAKE_CASE					:  str             =      {
    '''sample_size''': 32,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 2,
    '''num_class_embeds''': 1000,
    '''block_out_channels''': [32, 64],
    '''attention_head_dim''': 8,
    '''down_block_types''': [
        '''ResnetDownsampleBlock2D''',
        '''AttnDownBlock2D''',
    ],
    '''up_block_types''': [
        '''AttnUpBlock2D''',
        '''ResnetUpsampleBlock2D''',
    ],
    '''resnet_time_scale_shift''': '''scale_shift''',
    '''upsample_type''': '''resnet''',
    '''downsample_type''': '''resnet''',
}

_SCREAMING_SNAKE_CASE					:  Dict             =      {
    '''sample_size''': 64,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 3,
    '''num_class_embeds''': 1000,
    '''block_out_channels''': [192, 192 * 2, 192 * 3, 192 * 4],
    '''attention_head_dim''': 64,
    '''down_block_types''': [
        '''ResnetDownsampleBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
    ],
    '''up_block_types''': [
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''ResnetUpsampleBlock2D''',
    ],
    '''resnet_time_scale_shift''': '''scale_shift''',
    '''upsample_type''': '''resnet''',
    '''downsample_type''': '''resnet''',
}

_SCREAMING_SNAKE_CASE					:  int             =      {
    '''sample_size''': 256,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 2,
    '''num_class_embeds''': None,
    '''block_out_channels''': [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4],
    '''attention_head_dim''': 64,
    '''down_block_types''': [
        '''ResnetDownsampleBlock2D''',
        '''ResnetDownsampleBlock2D''',
        '''ResnetDownsampleBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
    ],
    '''up_block_types''': [
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''ResnetUpsampleBlock2D''',
        '''ResnetUpsampleBlock2D''',
        '''ResnetUpsampleBlock2D''',
    ],
    '''resnet_time_scale_shift''': '''default''',
    '''upsample_type''': '''resnet''',
    '''downsample_type''': '''resnet''',
}

_SCREAMING_SNAKE_CASE					:  int             =      {
    '''num_train_timesteps''': 40,
    '''sigma_min''': 0.0_0_2,
    '''sigma_max''': 8_0.0,
}

_SCREAMING_SNAKE_CASE					:  str             =      {
    '''num_train_timesteps''': 201,
    '''sigma_min''': 0.0_0_2,
    '''sigma_max''': 8_0.0,
}

_SCREAMING_SNAKE_CASE					:  Tuple             =      {
    '''num_train_timesteps''': 151,
    '''sigma_min''': 0.0_0_2,
    '''sigma_max''': 8_0.0,
}
def     UpperCAmelCase_       (						_A    ):





     '''simple docstring'''




     if isinstance(_A				,		_A    ):
          return v
     if v.lower() in ("yes", "true", "t", "y", "1"):
          return True
     elif v.lower() in ("no", "false", "f", "n", "0"):
          return False
     else:
          raise argparse.ArgumentTypeError('''boolean value expected'''    )
def     UpperCAmelCase_       (						_A				,		_A				,		_A				,		_A				,		_A=False    ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.in_layers.0.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.in_layers.0.bias''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.in_layers.2.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.in_layers.2.bias''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.emb_layers.1.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.emb_layers.1.bias''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.out_layers.0.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.out_layers.0.bias''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.out_layers.3.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.out_layers.3.bias''']

     if has_skip:
          SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.skip_connection.weight''']
          SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.skip_connection.bias''']

     return new_checkpoint
def     UpperCAmelCase_       (						_A				,		_A				,		_A				,		_A				,		_A=None    ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.qkv.weight'''].chunk(3				,		dim=0    )
     SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.qkv.bias'''].chunk(3				,		dim=0    )

     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.norm.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.norm.bias''']

     SCREAMING_SNAKE_CASE__   =   weight_q.squeeze(-1    ).squeeze(-1    )
     SCREAMING_SNAKE_CASE__   =   bias_q.squeeze(-1    ).squeeze(-1    )
     SCREAMING_SNAKE_CASE__   =   weight_k.squeeze(-1    ).squeeze(-1    )
     SCREAMING_SNAKE_CASE__   =   bias_k.squeeze(-1    ).squeeze(-1    )
     SCREAMING_SNAKE_CASE__   =   weight_v.squeeze(-1    ).squeeze(-1    )
     SCREAMING_SNAKE_CASE__   =   bias_v.squeeze(-1    ).squeeze(-1    )

     SCREAMING_SNAKE_CASE__   =   (
         checkpoint[F'''{old_prefix}.proj_out.weight'''].squeeze(-1    ).squeeze(-1    )
     )
     SCREAMING_SNAKE_CASE__   =   checkpoint[F'''{old_prefix}.proj_out.bias'''].squeeze(-1    ).squeeze(-1    )

     return new_checkpoint
def     UpperCAmelCase_       (						_A				,		_A    ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__   =   torch.load(_A				,		map_location='''cpu'''    )
     SCREAMING_SNAKE_CASE__   =   {}

     SCREAMING_SNAKE_CASE__   =   checkpoint['''time_embed.0.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''time_embed.0.bias''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''time_embed.2.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''time_embed.2.bias''']

     if unet_config["num_class_embeds"] is not None:
          SCREAMING_SNAKE_CASE__   =   checkpoint['''label_emb.weight''']

     SCREAMING_SNAKE_CASE__   =   checkpoint['''input_blocks.0.0.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''input_blocks.0.0.bias''']

     SCREAMING_SNAKE_CASE__   =   unet_config['''down_block_types''']
     SCREAMING_SNAKE_CASE__   =   unet_config['''layers_per_block''']
     SCREAMING_SNAKE_CASE__   =   unet_config['''attention_head_dim''']
     SCREAMING_SNAKE_CASE__   =   unet_config['''block_out_channels''']
     SCREAMING_SNAKE_CASE__   =   1
     SCREAMING_SNAKE_CASE__   =   channels_list[0]

     for i, layer_type in enumerate(_A    ):
          SCREAMING_SNAKE_CASE__   =   channels_list[i]
          SCREAMING_SNAKE_CASE__   =   current_channels != prev_channels
          if layer_type == "ResnetDownsampleBlock2D":
               for j in range(_A    ):
                    SCREAMING_SNAKE_CASE__   =   F'''down_blocks.{i}.resnets.{j}'''
                    SCREAMING_SNAKE_CASE__   =   F'''input_blocks.{current_layer}.0'''
                    SCREAMING_SNAKE_CASE__   =   True if j == 0 and downsample_block_has_skip else False
                    SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A				,		has_skip=_A    )
                    current_layer += 1

          elif layer_type == "AttnDownBlock2D":
               for j in range(_A    ):
                    SCREAMING_SNAKE_CASE__   =   F'''down_blocks.{i}.resnets.{j}'''
                    SCREAMING_SNAKE_CASE__   =   F'''input_blocks.{current_layer}.0'''
                    SCREAMING_SNAKE_CASE__   =   True if j == 0 and downsample_block_has_skip else False
                    SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A				,		has_skip=_A    )
                    SCREAMING_SNAKE_CASE__   =   F'''down_blocks.{i}.attentions.{j}'''
                    SCREAMING_SNAKE_CASE__   =   F'''input_blocks.{current_layer}.1'''
                    SCREAMING_SNAKE_CASE__   =   convert_attention(
                        _A				,		_A				,		_A				,		_A				,		_A    )
                    current_layer += 1

          if i != len(_A    ) - 1:
               SCREAMING_SNAKE_CASE__   =   F'''down_blocks.{i}.downsamplers.0'''
               SCREAMING_SNAKE_CASE__   =   F'''input_blocks.{current_layer}.0'''
               SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A    )
               current_layer += 1

          SCREAMING_SNAKE_CASE__   =   current_channels

     # hardcoded the mid-block for now
     SCREAMING_SNAKE_CASE__   =   '''mid_block.resnets.0'''
     SCREAMING_SNAKE_CASE__   =   '''middle_block.0'''
     SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A    )
     SCREAMING_SNAKE_CASE__   =   '''mid_block.attentions.0'''
     SCREAMING_SNAKE_CASE__   =   '''middle_block.1'''
     SCREAMING_SNAKE_CASE__   =   convert_attention(_A				,		_A				,		_A				,		_A				,		_A    )
     SCREAMING_SNAKE_CASE__   =   '''mid_block.resnets.1'''
     SCREAMING_SNAKE_CASE__   =   '''middle_block.2'''
     SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A    )

     SCREAMING_SNAKE_CASE__   =   0
     SCREAMING_SNAKE_CASE__   =   unet_config['''up_block_types''']

     for i, layer_type in enumerate(_A    ):
          if layer_type == "ResnetUpsampleBlock2D":
               for j in range(layers_per_block + 1    ):
                    SCREAMING_SNAKE_CASE__   =   F'''up_blocks.{i}.resnets.{j}'''
                    SCREAMING_SNAKE_CASE__   =   F'''output_blocks.{current_layer}.0'''
                    SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A				,		has_skip=_A    )
                    current_layer += 1

               if i != len(_A    ) - 1:
                    SCREAMING_SNAKE_CASE__   =   F'''up_blocks.{i}.upsamplers.0'''
                    SCREAMING_SNAKE_CASE__   =   F'''output_blocks.{current_layer-1}.1'''
                    SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A    )
          elif layer_type == "AttnUpBlock2D":
               for j in range(layers_per_block + 1    ):
                    SCREAMING_SNAKE_CASE__   =   F'''up_blocks.{i}.resnets.{j}'''
                    SCREAMING_SNAKE_CASE__   =   F'''output_blocks.{current_layer}.0'''
                    SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A				,		has_skip=_A    )
                    SCREAMING_SNAKE_CASE__   =   F'''up_blocks.{i}.attentions.{j}'''
                    SCREAMING_SNAKE_CASE__   =   F'''output_blocks.{current_layer}.1'''
                    SCREAMING_SNAKE_CASE__   =   convert_attention(
                        _A				,		_A				,		_A				,		_A				,		_A    )
                    current_layer += 1

               if i != len(_A    ) - 1:
                    SCREAMING_SNAKE_CASE__   =   F'''up_blocks.{i}.upsamplers.0'''
                    SCREAMING_SNAKE_CASE__   =   F'''output_blocks.{current_layer-1}.2'''
                    SCREAMING_SNAKE_CASE__   =   convert_resnet(_A				,		_A				,		_A				,		_A    )

     SCREAMING_SNAKE_CASE__   =   checkpoint['''out.0.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''out.0.bias''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''out.2.weight''']
     SCREAMING_SNAKE_CASE__   =   checkpoint['''out.2.bias''']

     return new_checkpoint


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE					:  Optional[Any]             =      argparse.ArgumentParser()

       parser.add_argument('''--unet_path''', default=None, type=str, required=True, help='''Path to the unet.pt to convert.''')
       parser.add_argument(
           '''--dump_path''', default=None, type=str, required=True, help='''Path to output the converted UNet model.'''
       )
       parser.add_argument('''--class_cond''', default=True, type=str, help='''Whether the model is class-conditional.''')

       _SCREAMING_SNAKE_CASE					:  Tuple             =      parser.parse_args()
       _SCREAMING_SNAKE_CASE					:  List[str]             =      strabool(args.class_cond)

       _SCREAMING_SNAKE_CASE					:  int             =      os.path.basename(args.unet_path)
       print(F"Checkpoint: {ckpt_name}")

       # Get U-Net config
       if "imagenet64" in ckpt_name:
              _SCREAMING_SNAKE_CASE					:  Optional[Any]             =      IMAGENET_64_UNET_CONFIG
       elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
              _SCREAMING_SNAKE_CASE					:  int             =      LSUN_256_UNET_CONFIG
       elif "test" in ckpt_name:
              _SCREAMING_SNAKE_CASE					:  Union[str, Any]             =      TEST_UNET_CONFIG
       else:
              raise ValueError(F"Checkpoint type {ckpt_name} is not currently supported.")

       if not args.class_cond:
              _SCREAMING_SNAKE_CASE					:  Union[str, Any]             =      None

       _SCREAMING_SNAKE_CASE					:  int             =      con_pt_to_diffuser(args.unet_path, unet_config)

       _SCREAMING_SNAKE_CASE					:  Optional[int]             =      UNetaDModel(**unet_config)
       image_unet.load_state_dict(converted_unet_ckpt)

       # Get scheduler config
       if "cd" in ckpt_name or "test" in ckpt_name:
              _SCREAMING_SNAKE_CASE					:  Optional[Any]             =      CD_SCHEDULER_CONFIG
       elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
              _SCREAMING_SNAKE_CASE					:  Any             =      CT_IMAGENET_64_SCHEDULER_CONFIG
       elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
              _SCREAMING_SNAKE_CASE					:  int             =      CT_LSUN_256_SCHEDULER_CONFIG
       else:
              raise ValueError(F"Checkpoint type {ckpt_name} is not currently supported.")

       _SCREAMING_SNAKE_CASE					:  int             =      CMStochasticIterativeScheduler(**scheduler_config)

       _SCREAMING_SNAKE_CASE					:  str             =      ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
       consistency_model.save_pretrained(args.dump_path)


