





"""simple docstring"""




import math
from typing import Optional

import numpy as np
import torch
from torch import nn

from .activations import get_activation





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase = False					,   _UpperCAmelCase = 1					,   _UpperCAmelCase = 1					,   _UpperCAmelCase = 10000					,   ):






  """simple docstring"""





  assert len(timesteps.shape   ) == 1, "Timesteps should be a 1d-array"

  A_    :			Dict									=    embedding_dim // 2
  A_    :			Optional[Any]									=    -math.log(_UpperCAmelCase   ) * torch.arange(
      start=0					,   end=_UpperCAmelCase					,   dtype=torch.floataa					,   device=timesteps.device   )
  A_    :			List[str]									=    exponent / (half_dim - downscale_freq_shift)

  A_    :			str									=    torch.exp(_UpperCAmelCase   )
  A_    :			str									=    timesteps[:, None].float() * emb[None, :]

  # scale embeddings
  A_    :			Union[str, Any]									=    scale * emb

  # concat sine and cosine embeddings
  A_    :			List[Any]									=    torch.cat([torch.sin(_UpperCAmelCase   ), torch.cos(_UpperCAmelCase   )]					,   dim=-1   )

  # flip sine and cosine embeddings
  if flip_sin_to_cos:
    A_    :			Any									=    torch.cat([emb[:, half_dim:], emb[:, :half_dim]]					,   dim=-1   )

  # zero pad
  if embedding_dim % 2 == 1:
    A_    :			str									=    torch.nn.functional.pad(_UpperCAmelCase					,   (0, 1, 0, 0)   )
  return emb





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase=False					,   _UpperCAmelCase=0   ):






  """simple docstring"""





  A_    :			Union[str, Any]									=    np.arange(_UpperCAmelCase					,   dtype=np.floataa   )
  A_    :			Any									=    np.arange(_UpperCAmelCase					,   dtype=np.floataa   )
  A_    :			Tuple									=    np.meshgrid(_UpperCAmelCase					,   _UpperCAmelCase   )  # here w goes first
  A_    :			List[Any]									=    np.stack(_UpperCAmelCase					,   axis=0   )

  A_    :			Optional[int]									=    grid.reshape([2, 1, grid_size, grid_size]   )
  A_    :			Dict									=    get_ad_sincos_pos_embed_from_grid(_UpperCAmelCase					,   _UpperCAmelCase   )
  if cls_token and extra_tokens > 0:
    A_    :			List[str]									=    np.concatenate([np.zeros([extra_tokens, embed_dim]   ), pos_embed]					,   axis=0   )
  return pos_embed





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase   ):






  """simple docstring"""





  if embed_dim % 2 != 0:
    raise ValueError('embed_dim must be divisible by 2'   )

  # use half of dimensions to encode grid_h
  A_    :			Union[str, Any]									=    get_ad_sincos_pos_embed_from_grid(embed_dim // 2					,   grid[0]   )  # (H*W, D/2)
  A_    :			str									=    get_ad_sincos_pos_embed_from_grid(embed_dim // 2					,   grid[1]   )  # (H*W, D/2)

  A_    :			List[Any]									=    np.concatenate([emb_h, emb_w]					,   axis=1   )  # (H*W, D)
  return emb





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase   ):






  """simple docstring"""





  if embed_dim % 2 != 0:
    raise ValueError('embed_dim must be divisible by 2'   )

  A_    :			Tuple									=    np.arange(embed_dim // 2					,   dtype=np.floataa   )
  omega /= embed_dim / 2.0
  A_    :			Optional[Any]									=    1.0 / 10000**omega  # (D/2,)

  A_    :			Tuple									=    pos.reshape(-1   )  # (M,)
  A_    :			List[Any]									=    np.einsum('m,d->md'					,   _UpperCAmelCase					,   _UpperCAmelCase   )  # (M, D/2), outer product

  A_    :			str									=    np.sin(_UpperCAmelCase   )  # (M, D/2)
  A_    :			Union[str, Any]									=    np.cos(_UpperCAmelCase   )  # (M, D/2)

  A_    :			int									=    np.concatenate([emb_sin, emb_cos]					,   axis=1   )  # (M, D)
  return emb

class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_=2_2_4   , snake_case_=2_2_4   , snake_case_=1_6   , snake_case_=3   , snake_case_=7_6_8   , snake_case_=False   , snake_case_=True   , snake_case_=True   , ):

        """simple docstring"""

        super().__init__()

        A_    :			List[Any]									=    (height // patch_size) * (width // patch_size)
        A_    :			Optional[Any]									=    flatten
        A_    :			Dict									=    layer_norm

        A_    :			str									=    nn.Convad(
            snake_case_   , snake_case_   , kernel_size=(patch_size, patch_size)   , stride=snake_case_   , bias=snake_case_    )
        if layer_norm:
          A_    :			str									=    nn.LayerNorm(snake_case_   , elementwise_affine=snake_case_   , eps=1E-6    )
        else:
          A_    :			int									=    None

        A_    :			List[str]									=    get_ad_sincos_pos_embed(snake_case_   , int(num_patches**0.5    )    )
        self.register_buffer('pos_embed'   , torch.from_numpy(snake_case_    ).float().unsqueeze(0    )   , persistent=snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			List[str]									=    self.proj(snake_case_    )
        if self.flatten:
          A_    :			Dict									=    latent.flatten(2    ).transpose(1   , 2    )  # BCHW -> BNC
        if self.layer_norm:
          A_    :			List[Any]									=    self.norm(snake_case_    )
        return latent + self.pos_embed


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_ = "silu"   , snake_case_ = None   , snake_case_ = None   , snake_case_=None   , ):

        """simple docstring"""

        super().__init__()

        A_    :			str									=    nn.Linear(snake_case_   , snake_case_    )

        if cond_proj_dim is not None:
          A_    :			Tuple									=    nn.Linear(snake_case_   , snake_case_   , bias=snake_case_    )
        else:
          A_    :			Any									=    None

        A_    :			List[Any]									=    get_activation(snake_case_    )

        if out_dim is not None:
          A_    :			Dict									=    out_dim
        else:
          A_    :			Union[str, Any]									=    time_embed_dim
        A_    :			Dict									=    nn.Linear(snake_case_   , snake_case_    )

        if post_act_fn is None:
          A_    :			Union[str, Any]									=    None
        else:
          A_    :			List[str]									=    get_activation(snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_=None    ):

        """simple docstring"""

        if condition is not None:
          A_    :			List[Any]									=    sample + self.cond_proj(snake_case_    )
        A_    :			Optional[int]									=    self.linear_a(snake_case_    )

        if self.act is not None:
          A_    :			Dict									=    self.act(snake_case_    )

        A_    :			Optional[int]									=    self.linear_a(snake_case_    )

        if self.post_act is not None:
          A_    :			Optional[int]									=    self.post_act(snake_case_    )
        return sample


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_    ):

        """simple docstring"""

        super().__init__()
        A_    :			List[Any]									=    num_channels
        A_    :			Dict									=    flip_sin_to_cos
        A_    :			Optional[Any]									=    downscale_freq_shift

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    get_timestep_embedding(
            snake_case_   , self.num_channels   , flip_sin_to_cos=self.flip_sin_to_cos   , downscale_freq_shift=self.downscale_freq_shift   , )
        return t_emb


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_ = 2_5_6   , snake_case_ = 1.0   , snake_case_=True   , snake_case_=True   , snake_case_=False    ):

        """simple docstring"""

        super().__init__()
        A_    :			List[Any]									=    nn.Parameter(torch.randn(snake_case_    ) * scale   , requires_grad=snake_case_    )
        A_    :			List[str]									=    log
        A_    :			Optional[Any]									=    flip_sin_to_cos

        if set_W_to_weight:
          # to delete later
          A_    :			List[Any]									=    nn.Parameter(torch.randn(snake_case_    ) * scale   , requires_grad=snake_case_    )

          A_    :			Any									=    self.W

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        if self.log:
          A_    :			List[Any]									=    torch.log(snake_case_    )

        A_    :			List[Any]									=    x[:, None] * self.weight[None, :] * 2 * np.pi

        if self.flip_sin_to_cos:
          A_    :			List[str]									=    torch.cat([torch.cos(snake_case_    ), torch.sin(snake_case_    )]   , dim=-1    )
        else:
          A_    :			Optional[Any]									=    torch.cat([torch.sin(snake_case_    ), torch.cos(snake_case_    )]   , dim=-1    )
        return out


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_   , snake_case_   , ):

        """simple docstring"""

        super().__init__()

        A_    :			Optional[int]									=    height
        A_    :			Union[str, Any]									=    width
        A_    :			Optional[Any]									=    num_embed
        A_    :			Dict									=    embed_dim

        A_    :			Union[str, Any]									=    nn.Embedding(self.num_embed   , snake_case_    )
        A_    :			str									=    nn.Embedding(self.height   , snake_case_    )
        A_    :			Tuple									=    nn.Embedding(self.width   , snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			List[str]									=    self.emb(snake_case_    )

        A_    :			Optional[Any]									=    self.height_emb(torch.arange(self.height   , device=index.device    ).view(1   , self.height    )    )

        # 1 x H x D -> 1 x H x 1 x D
        A_    :			str									=    height_emb.unsqueeze(2    )

        A_    :			Optional[Any]									=    self.width_emb(torch.arange(self.width   , device=index.device    ).view(1   , self.width    )    )

        # 1 x W x D -> 1 x 1 x W x D
        A_    :			List[str]									=    width_emb.unsqueeze(1    )

        A_    :			List[Any]									=    height_emb + width_emb

        # 1 x H x W x D -> 1 x L xD
        A_    :			Optional[int]									=    pos_emb.view(1   , self.height * self.width   , -1    )

        A_    :			Optional[int]									=    emb + pos_emb[:, : emb.shape[1], :]

        return emb


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_    ):

        """simple docstring"""

        super().__init__()
        A_    :			List[Any]									=    dropout_prob > 0
        A_    :			List[str]									=    nn.Embedding(num_classes + use_cfg_embedding   , snake_case_    )
        A_    :			int									=    num_classes
        A_    :			Optional[int]									=    dropout_prob

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_=None    ):

        """simple docstring"""

        if force_drop_ids is None:
          A_    :			Optional[Any]									=    torch.rand(labels.shape[0]   , device=labels.device    ) < self.dropout_prob
        else:
          A_    :			Optional[Any]									=    torch.tensor(force_drop_ids == 1    )
        A_    :			List[Any]									=    torch.where(snake_case_   , self.num_classes   , snake_case_    )
        return labels

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_=None    ):

        """simple docstring"""

        A_    :			Optional[int]									=    self.dropout_prob > 0
        if (self.training and use_dropout) or (force_drop_ids is not None):
          A_    :			Optional[int]									=    self.token_drop(snake_case_   , snake_case_    )
        A_    :			Optional[Any]									=    self.embedding_table(snake_case_    )
        return embeddings


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_ = 1_0_2_4   , snake_case_ = 7_6_8   , snake_case_ = 7_6_8   , snake_case_ = 1_0   , ):

        """simple docstring"""

        super().__init__()

        A_    :			Optional[int]									=    num_image_text_embeds
        A_    :			str									=    nn.Linear(snake_case_   , self.num_image_text_embeds * cross_attention_dim    )
        A_    :			Optional[Any]									=    nn.Linear(snake_case_   , snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			int									=    text_embeds.shape[0]

        # image
        A_    :			int									=    self.image_embeds(snake_case_    )
        A_    :			Tuple									=    image_text_embeds.reshape(snake_case_   , self.num_image_text_embeds   , -1    )

        # text
        A_    :			Optional[int]									=    self.text_proj(snake_case_    )

        return torch.cat([image_text_embeds, text_embeds]   , dim=1    )


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_ = 7_6_8   , snake_case_ = 7_6_8   , snake_case_ = 3_2   , ):

        """simple docstring"""

        super().__init__()

        A_    :			Any									=    num_image_text_embeds
        A_    :			Dict									=    nn.Linear(snake_case_   , self.num_image_text_embeds * cross_attention_dim    )
        A_    :			Union[str, Any]									=    nn.LayerNorm(snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			List[str]									=    image_embeds.shape[0]

        # image
        A_    :			str									=    self.image_embeds(snake_case_    )
        A_    :			List[str]									=    image_embeds.reshape(snake_case_   , self.num_image_text_embeds   , -1    )
        A_    :			Dict									=    self.norm(snake_case_    )
        return image_embeds


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_=0.1    ):

        """simple docstring"""

        super().__init__()

        A_    :			Union[str, Any]									=    Timesteps(num_channels=2_5_6   , flip_sin_to_cos=snake_case_   , downscale_freq_shift=1    )
        A_    :			int									=    TimestepEmbedding(in_channels=2_5_6   , time_embed_dim=snake_case_    )
        A_    :			Union[str, Any]									=    LabelEmbedding(snake_case_   , snake_case_   , snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_=None    ):

        """simple docstring"""

        A_    :			str									=    self.time_proj(snake_case_    )
        A_    :			str									=    self.timestep_embedder(timesteps_proj.to(dtype=snake_case_    )    )  # (N, D)

        A_    :			List[Any]									=    self.class_embedder(snake_case_    )  # (N, D)

        A_    :			List[str]									=    timesteps_emb + class_labels  # (N, D)

        return conditioning


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_ = 6_4    ):

        """simple docstring"""

        super().__init__()
        A_    :			str									=    nn.LayerNorm(snake_case_    )
        A_    :			Tuple									=    AttentionPooling(snake_case_   , snake_case_    )
        A_    :			Dict									=    nn.Linear(snake_case_   , snake_case_    )
        A_    :			Optional[int]									=    nn.LayerNorm(snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			int									=    self.norma(snake_case_    )
        A_    :			List[Any]									=    self.pool(snake_case_    )
        A_    :			List[Any]									=    self.proj(snake_case_    )
        A_    :			List[Any]									=    self.norma(snake_case_    )
        return hidden_states


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_ = 7_6_8   , snake_case_ = 7_6_8   , snake_case_ = 1_5_3_6    ):

        """simple docstring"""

        super().__init__()
        A_    :			Optional[Any]									=    nn.Linear(snake_case_   , snake_case_    )
        A_    :			Union[str, Any]									=    nn.LayerNorm(snake_case_    )
        A_    :			List[str]									=    nn.Linear(snake_case_   , snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			Any									=    self.text_proj(snake_case_    )
        A_    :			Any									=    self.text_norm(snake_case_    )

        # image
        A_    :			Tuple									=    self.image_proj(snake_case_    )

        return time_image_embeds + time_text_embeds


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_ = 7_6_8   , snake_case_ = 1_5_3_6    ):

        """simple docstring"""

        super().__init__()
        A_    :			List[Any]									=    nn.Linear(snake_case_   , snake_case_    )
        A_    :			Union[str, Any]									=    nn.LayerNorm(snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    self.image_proj(snake_case_    )
        A_    :			Any									=    self.image_norm(snake_case_    )
        return time_image_embeds


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_ = 7_6_8   , snake_case_ = 1_5_3_6    ):

        """simple docstring"""

        super().__init__()
        A_    :			Optional[Any]									=    nn.Linear(snake_case_   , snake_case_    )
        A_    :			Optional[Any]									=    nn.LayerNorm(snake_case_    )
        A_    :			Any									=    nn.Sequential(
            nn.Convad(3   , 1_6   , 3   , padding=1    )   , nn.SiLU()   , nn.Convad(1_6   , 1_6   , 3   , padding=1    )   , nn.SiLU()   , nn.Convad(1_6   , 3_2   , 3   , padding=1   , stride=2    )   , nn.SiLU()   , nn.Convad(3_2   , 3_2   , 3   , padding=1    )   , nn.SiLU()   , nn.Convad(3_2   , 9_6   , 3   , padding=1   , stride=2    )   , nn.SiLU()   , nn.Convad(9_6   , 9_6   , 3   , padding=1    )   , nn.SiLU()   , nn.Convad(9_6   , 2_5_6   , 3   , padding=1   , stride=2    )   , nn.SiLU()   , nn.Convad(2_5_6   , 4   , 3   , padding=1    )   , )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			int									=    self.image_proj(snake_case_    )
        A_    :			int									=    self.image_norm(snake_case_    )
        A_    :			int									=    self.input_hint_block(snake_case_    )
        return time_image_embeds, hint


class 					_UpperCAmelCase    (  nn.Module						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_=None    ):

        """simple docstring"""

        super().__init__()
        A_    :			Dict									=    dtype
        A_    :			List[str]									=    nn.Parameter(torch.randn(1   , snake_case_    ) / embed_dim**0.5    )
        A_    :			List[str]									=    nn.Linear(snake_case_   , snake_case_   , dtype=self.dtype    )
        A_    :			Optional[int]									=    nn.Linear(snake_case_   , snake_case_   , dtype=self.dtype    )
        A_    :			str									=    nn.Linear(snake_case_   , snake_case_   , dtype=self.dtype    )
        A_    :			str									=    num_heads
        A_    :			int									=    embed_dim // self.num_heads

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_		,			A_		,			A_    :			Tuple									=    x.size()

        def shape(snake_case_    ):
          # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
          A_    :			str									=    x.view(snake_case_   , -1   , self.num_heads   , self.dim_per_head    )
          # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
          A_    :			Optional[int]									=    x.transpose(1   , 2    )
          # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
          A_    :			List[Any]									=    x.reshape(bs * self.num_heads   , -1   , self.dim_per_head    )
          # (bs*n_heads, length, dim_per_head) --> (bs*n_heads, dim_per_head, length)
          A_    :			List[Any]									=    x.transpose(1   , 2    )
          return x

        A_    :			int									=    x.mean(dim=1   , keepdim=snake_case_    ) + self.positional_embedding.to(x.dtype    )
        A_    :			Optional[Any]									=    torch.cat([class_token, x]   , dim=1    )  # (bs, length+1, width)

        # (bs*n_heads, class_token_length, dim_per_head)
        A_    :			Any									=    shape(self.q_proj(snake_case_    )    )
        # (bs*n_heads, length+class_token_length, dim_per_head)
        A_    :			Tuple									=    shape(self.k_proj(snake_case_    )    )
        A_    :			List[str]									=    shape(self.v_proj(snake_case_    )    )

        # (bs*n_heads, class_token_length, length+class_token_length):
        A_    :			str									=    1 / math.sqrt(math.sqrt(self.dim_per_head    )    )
        A_    :			str									=    torch.einsum('bct,bcs->bts'   , q * scale   , k * scale    )  # More stable with f16 than dividing afterwards
        A_    :			Tuple									=    torch.softmax(weight.float()   , dim=-1    ).type(weight.dtype    )

        # (bs*n_heads, dim_per_head, class_token_length)
        A_    :			int									=    torch.einsum('bts,bcs->bct'   , snake_case_   , snake_case_    )

        # (bs, length+1, width)
        A_    :			int									=    a.reshape(snake_case_   , -1   , 1    ).transpose(1   , 2    )

        return a[:, 0, :]  # cls_token






'''simple docstring'''







import argparse
import json
import os

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
snake_case_  :							Optional[int]            =	logging.get_logger(__name__)





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[int] )			->						Dict:
						UpperCAmelCase_      :   int      =	ConvNextVaConfig()

						if "atto" in checkpoint_url:
												UpperCAmelCase_      :   Tuple      =	[2, 2, 6, 2]
												UpperCAmelCase_      :   Tuple      =	[40, 80, 160, 320]
						if "femto" in checkpoint_url:
												UpperCAmelCase_      :   List[Any]      =	[2, 2, 6, 2]
												UpperCAmelCase_      :   Any      =	[48, 96, 192, 384]
						if "pico" in checkpoint_url:
												UpperCAmelCase_      :   Union[str, Any]      =	[2, 2, 6, 2]
												UpperCAmelCase_      :   int      =	[64, 128, 256, 512]
						if "nano" in checkpoint_url:
												UpperCAmelCase_      :   str      =	[2, 2, 8, 2]
												UpperCAmelCase_      :   List[str]      =	[80, 160, 320, 640]
						if "tiny" in checkpoint_url:
												UpperCAmelCase_      :   Optional[Any]      =	[3, 3, 9, 3]
												UpperCAmelCase_      :   Any      =	[96, 192, 384, 768]
						if "base" in checkpoint_url:
												UpperCAmelCase_      :   Union[str, Any]      =	[3, 3, 27, 3]
												UpperCAmelCase_      :   Optional[int]      =	[128, 256, 512, 1024]
						if "large" in checkpoint_url:
												UpperCAmelCase_      :   int      =	[3, 3, 27, 3]
												UpperCAmelCase_      :   List[Any]      =	[192, 384, 768, 1536]
						if "huge" in checkpoint_url:
												UpperCAmelCase_      :   Tuple      =	[3, 3, 27, 3]
												UpperCAmelCase_      :   Dict      =	[352, 704, 1408, 2816]

						UpperCAmelCase_      :   Union[str, Any]      =	1000
						UpperCAmelCase_      :   Optional[Any]      =	'''imagenet-1k-id2label.json'''
						UpperCAmelCase_      :   int      =	(1, 1000)

						UpperCAmelCase_      :   Any      =	'''huggingface/label-files'''
						UpperCAmelCase_      :   List[str]      =	num_labels
						UpperCAmelCase_      :   Optional[int]      =	json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      repo_type='''dataset''' ),      '''r''' ) )
						UpperCAmelCase_      :   Tuple      =	{int(SCREAMING_SNAKE_CASE__ ): v for k, v in idalabel.items()}

						UpperCAmelCase_      :   Union[str, Any]      =	idalabel
						UpperCAmelCase_      :   List[Any]      =	{v: k for k, v in idalabel.items()}
						UpperCAmelCase_      :   str      =	hidden_sizes
						UpperCAmelCase_      :   Dict      =	depths

						return config, expected_shape





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Dict )			->						Optional[Any]:
						if "downsample_layers.0.0" in name:
												UpperCAmelCase_      :   str      =	name.replace('''downsample_layers.0.0''',      '''embeddings.patch_embeddings''' )
						if "downsample_layers.0.1" in name:
												UpperCAmelCase_      :   Optional[Any]      =	name.replace('''downsample_layers.0.1''',      '''embeddings.norm''' )  # we rename to layernorm later on
						if "downsample_layers.1.0" in name:
												UpperCAmelCase_      :   Dict      =	name.replace('''downsample_layers.1.0''',      '''stages.1.downsampling_layer.0''' )
						if "downsample_layers.1.1" in name:
												UpperCAmelCase_      :   Any      =	name.replace('''downsample_layers.1.1''',      '''stages.1.downsampling_layer.1''' )
						if "downsample_layers.2.0" in name:
												UpperCAmelCase_      :   Optional[int]      =	name.replace('''downsample_layers.2.0''',      '''stages.2.downsampling_layer.0''' )
						if "downsample_layers.2.1" in name:
												UpperCAmelCase_      :   List[Any]      =	name.replace('''downsample_layers.2.1''',      '''stages.2.downsampling_layer.1''' )
						if "downsample_layers.3.0" in name:
												UpperCAmelCase_      :   Optional[int]      =	name.replace('''downsample_layers.3.0''',      '''stages.3.downsampling_layer.0''' )
						if "downsample_layers.3.1" in name:
												UpperCAmelCase_      :   Optional[int]      =	name.replace('''downsample_layers.3.1''',      '''stages.3.downsampling_layer.1''' )
						if "stages" in name and "downsampling_layer" not in name:
												# stages.0.0. for instance should be renamed to stages.0.layers.0.
												UpperCAmelCase_      :   Any      =	name[: len('''stages.0''' )] + '''.layers''' + name[len('''stages.0''' ) :]
						if "gamma" in name:
												UpperCAmelCase_      :   Optional[Any]      =	name.replace('''gamma''',      '''weight''' )
						if "beta" in name:
												UpperCAmelCase_      :   Dict      =	name.replace('''beta''',      '''bias''' )
						if "stages" in name:
												UpperCAmelCase_      :   str      =	name.replace('''stages''',      '''encoder.stages''' )
						if "norm" in name:
												UpperCAmelCase_      :   Union[str, Any]      =	name.replace('''norm''',      '''layernorm''' )
						if "head" in name:
												UpperCAmelCase_      :   Union[str, Any]      =	name.replace('''head''',      '''classifier''' )

						return name





def 							lowerCamelCase_   (	)			->						Union[str, Any]:
						UpperCAmelCase_      :   int      =	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
						UpperCAmelCase_      :   Optional[int]      =	Image.open(requests.get(SCREAMING_SNAKE_CASE__,      stream=SCREAMING_SNAKE_CASE__ ).raw )
						return im





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Optional[int] )			->						int:
						if "224" in checkpoint_url:
												UpperCAmelCase_      :   Optional[Any]      =	224
												UpperCAmelCase_      :   Any      =	224 / 256
						elif "384" in checkpoint_url:
												UpperCAmelCase_      :   int      =	384
												UpperCAmelCase_      :   Optional[Any]      =	None
						else:
												UpperCAmelCase_      :   List[Any]      =	512
												UpperCAmelCase_      :   Union[str, Any]      =	None

						return ConvNextImageProcessor(
						    size=SCREAMING_SNAKE_CASE__,      crop_pct=SCREAMING_SNAKE_CASE__,      image_mean=[0.4_85, 0.4_56, 0.4_06],      image_std=[0.2_29, 0.2_24, 0.2_25],      resample=PILImageResampling.BICUBIC,      )





@torch.no_grad()
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      List[Any],      SCREAMING_SNAKE_CASE__   :      str )			->						Optional[Any]:
						print('''Downloading original model from checkpoint...''' )
						# define ConvNeXTV2 configuration based on URL
						UpperCAmelCase_    ,							UpperCAmelCase_      :   Tuple      =	get_convnextva_config(SCREAMING_SNAKE_CASE__ )
						# load original state_dict from URL
						UpperCAmelCase_      :   int      =	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE__ )['''model''']

						print('''Converting model parameters...''' )
						# rename keys
						for key in state_dict.copy().keys():
												UpperCAmelCase_      :   Optional[Any]      =	state_dict.pop(SCREAMING_SNAKE_CASE__ )
												UpperCAmelCase_      :   List[str]      =	val
						# add prefix to all keys expect classifier head
						for key in state_dict.copy().keys():
												UpperCAmelCase_      :   List[str]      =	state_dict.pop(SCREAMING_SNAKE_CASE__ )
												if not key.startswith('''classifier''' ):
																		UpperCAmelCase_      :   Optional[Any]      =	'''convnextv2.''' + key
												UpperCAmelCase_      :   Union[str, Any]      =	val

						# load HuggingFace model
						UpperCAmelCase_      :   Optional[int]      =	ConvNextVaForImageClassification(SCREAMING_SNAKE_CASE__ )
						model.load_state_dict(SCREAMING_SNAKE_CASE__ )
						model.eval()

						# Check outputs on an image, prepared by ConvNextImageProcessor
						UpperCAmelCase_      :   Any      =	convert_preprocessor(SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   Tuple      =	preprocessor(images=prepare_img(),      return_tensors='''pt''' )
						UpperCAmelCase_      :   Optional[Any]      =	model(**SCREAMING_SNAKE_CASE__ ).logits

						# note: the logits below were obtained without center cropping
						if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[Any]      =	torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
												UpperCAmelCase_      :   List[Any]      =	torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
												UpperCAmelCase_      :   Dict      =	torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
												UpperCAmelCase_      :   Tuple      =	torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[Any]      =	torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
												UpperCAmelCase_      :   int      =	torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
												UpperCAmelCase_      :   List[Any]      =	torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
												UpperCAmelCase_      :   Dict      =	torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
												UpperCAmelCase_      :   Dict      =	torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
												UpperCAmelCase_      :   Optional[int]      =	torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
												UpperCAmelCase_      :   Any      =	torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
												UpperCAmelCase_      :   Optional[Any]      =	torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
												UpperCAmelCase_      :   List[Any]      =	torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09] )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
												UpperCAmelCase_      :   Tuple      =	torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26] )
						else:
												raise ValueError(F"""Unknown URL: {checkpoint_url}""" )

						assert torch.allclose(logits[0, :5],      SCREAMING_SNAKE_CASE__,      atol=1E-3 )
						assert logits.shape == expected_shape
						print('''Model outputs match the original results!''' )

						if save_model:
												print('''Saving model to local...''' )
												# Create folder to save model
												if not os.path.isdir(SCREAMING_SNAKE_CASE__ ):
																		os.mkdir(SCREAMING_SNAKE_CASE__ )

												model.save_pretrained(SCREAMING_SNAKE_CASE__ )
												preprocessor.save_pretrained(SCREAMING_SNAKE_CASE__ )

						UpperCAmelCase_      :   Dict      =	'''convnextv2'''
						if "atto" in checkpoint_url:
												model_name += "-atto"
						if "femto" in checkpoint_url:
												model_name += "-femto"
						if "pico" in checkpoint_url:
												model_name += "-pico"
						if "nano" in checkpoint_url:
												model_name += "-nano"
						elif "tiny" in checkpoint_url:
												model_name += "-tiny"
						elif "base" in checkpoint_url:
												model_name += "-base"
						elif "large" in checkpoint_url:
												model_name += "-large"
						elif "huge" in checkpoint_url:
												model_name += "-huge"
						if "22k" in checkpoint_url and "1k" not in checkpoint_url:
												model_name += "-22k"
						elif "22k" in checkpoint_url and "1k" in checkpoint_url:
												model_name += "-22k-1k"
						elif "1k" in checkpoint_url:
												model_name += "-1k"
						if "224" in checkpoint_url:
												model_name += "-224"
						elif "384" in checkpoint_url:
												model_name += "-384"
						elif "512" in checkpoint_url:
												model_name += "-512"

						if push_to_hub:
												print(F"""Pushing {model_name} to the hub...""" )
												model.push_to_hub(SCREAMING_SNAKE_CASE__ )
												preprocessor.push_to_hub(SCREAMING_SNAKE_CASE__ )


if __name__ == "__main__":
					snake_case_  :							Tuple            =	argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    "--checkpoint_url",
					    default="https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt",
					    type=str,
					    help="URL of the original ConvNeXTV2 checkpoint you'd like to convert.",
					)
					parser.add_argument(
					    "--pytorch_dump_folder_path",
					    default="model",
					    type=str,
					    help="Path to the output PyTorch model directory.",
					)
					parser.add_argument("--save_model", action="store_true", help="Save model to local")
					parser.add_argument("--push_to_hub", action="store_true", help="Push model and image preprocessor to the hub")

					snake_case_  :							Optional[int]            =	parser.parse_args()
					convert_convnextva_checkpoint(
					    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
					)





