





'''simple docstring'''

import argparse
import json

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

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling





def lowerCamelCase						(     __lowerCamelCase :    Optional[int]    )						->Dict:
    _SCREAMING_SNAKE_CASE									=			[2, 2, 6, 2] if """tiny""" in model_name else [2, 2, 18, 2]
    _SCREAMING_SNAKE_CASE									=			True if """large""" in model_name or """huge""" in model_name else False
    _SCREAMING_SNAKE_CASE									=			True if """large""" in model_name or """huge""" in model_name else False
    _SCREAMING_SNAKE_CASE									=			True if """large""" in model_name or """huge""" in model_name else False

    if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
        if "fl3" in model_name:
            _SCREAMING_SNAKE_CASE									=			[3, 3, 3, 3]
            _SCREAMING_SNAKE_CASE									=			[5, 5, 5, 5]
        elif "fl4" in model_name:
            _SCREAMING_SNAKE_CASE									=			[4, 4, 4, 4]
            _SCREAMING_SNAKE_CASE									=			[3, 3, 3, 3]

    if "tiny" in model_name or "small" in model_name or "base" in model_name:
        _SCREAMING_SNAKE_CASE									=			[3, 3, 3, 3]
        if "lrf" in model_name:
            _SCREAMING_SNAKE_CASE									=			[3, 3, 3, 3]
        else:
            _SCREAMING_SNAKE_CASE									=			[2, 2, 2, 2]

    if "tiny" in model_name:
        _SCREAMING_SNAKE_CASE									=			96
    elif "small" in model_name:
        _SCREAMING_SNAKE_CASE									=			96
    elif "base" in model_name:
        _SCREAMING_SNAKE_CASE									=			128
    elif "large" in model_name:
        _SCREAMING_SNAKE_CASE									=			192
    elif "xlarge" in model_name:
        _SCREAMING_SNAKE_CASE									=			256
    elif "huge" in model_name:
        _SCREAMING_SNAKE_CASE									=			352

    # set label information
    _SCREAMING_SNAKE_CASE									=			"""huggingface/label-files"""
    if "large" in model_name or "huge" in model_name:
        _SCREAMING_SNAKE_CASE									=			"""imagenet-22k-id2label.json"""
    else:
        _SCREAMING_SNAKE_CASE									=			"""imagenet-1k-id2label.json"""

    _SCREAMING_SNAKE_CASE									=			json.load(open(hf_hub_download(__lowerCamelCase					,       __lowerCamelCase					,       repo_type="""dataset"""    )					,       """r"""    )    )
    _SCREAMING_SNAKE_CASE									=			{int(__lowerCamelCase    ): v for k, v in idalabel.items()}
    _SCREAMING_SNAKE_CASE									=			{v: k for k, v in idalabel.items()}

    _SCREAMING_SNAKE_CASE									=			FocalNetConfig(
        embed_dim=__lowerCamelCase					,       depths=__lowerCamelCase					,       focal_levels=__lowerCamelCase					,       focal_windows=__lowerCamelCase					,       use_conv_embed=__lowerCamelCase					,       idalabel=__lowerCamelCase					,       labelaid=__lowerCamelCase					,       use_post_layernorm=__lowerCamelCase					,       use_layerscale=__lowerCamelCase					,       )

    return config





def lowerCamelCase						(     __lowerCamelCase :    Optional[Any]    )						->Optional[Any]:
    if "patch_embed.proj" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""patch_embed.proj"""					,       """embeddings.patch_embeddings.projection"""    )
    if "patch_embed.norm" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""patch_embed.norm"""					,       """embeddings.norm"""    )
    if "layers" in name:
        _SCREAMING_SNAKE_CASE									=			"""encoder.""" + name
    if "encoder.layers" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""encoder.layers"""					,       """encoder.stages"""    )
    if "downsample.proj" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""downsample.proj"""					,       """downsample.projection"""    )
    if "blocks" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""blocks"""					,       """layers"""    )
    if "modulation.f.weight" in name or "modulation.f.bias" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""modulation.f"""					,       """modulation.projection_in"""    )
    if "modulation.h.weight" in name or "modulation.h.bias" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""modulation.h"""					,       """modulation.projection_context"""    )
    if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""modulation.proj"""					,       """modulation.projection_out"""    )

    if name == "norm.weight":
        _SCREAMING_SNAKE_CASE									=			"""layernorm.weight"""
    if name == "norm.bias":
        _SCREAMING_SNAKE_CASE									=			"""layernorm.bias"""

    if "head" in name:
        _SCREAMING_SNAKE_CASE									=			name.replace("""head"""					,       """classifier"""    )
    else:
        _SCREAMING_SNAKE_CASE									=			"""focalnet.""" + name

    return name





def lowerCamelCase						(     __lowerCamelCase :    Tuple					,       __lowerCamelCase :    Dict					,       __lowerCamelCase :    Optional[int]=False    )						->List[Any]:
    # fmt: off
    _SCREAMING_SNAKE_CASE									=			{
        """focalnet-tiny""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth""",
        """focalnet-tiny-lrf""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth""",
        """focalnet-small""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth""",
        """focalnet-small-lrf""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth""",
        """focalnet-base""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth""",
        """focalnet-base-lrf""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth""",
        """focalnet-large-lrf-fl3""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth""",
        """focalnet-large-lrf-fl4""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth""",
        """focalnet-xlarge-lrf-fl3""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth""",
        """focalnet-xlarge-lrf-fl4""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth""",
    }
    # fmt: on

    _SCREAMING_SNAKE_CASE									=			model_name_to_url[model_name]
    print("""Checkpoint URL: """					,       __lowerCamelCase    )
    _SCREAMING_SNAKE_CASE									=			torch.hub.load_state_dict_from_url(__lowerCamelCase					,       map_location="""cpu"""    )["""model"""]

    # rename keys
    for key in state_dict.copy().keys():
        _SCREAMING_SNAKE_CASE									=			state_dict.pop(__lowerCamelCase    )
        _SCREAMING_SNAKE_CASE									=			val

    _SCREAMING_SNAKE_CASE									=			get_focalnet_config(__lowerCamelCase    )
    _SCREAMING_SNAKE_CASE									=			FocalNetForImageClassification(__lowerCamelCase    )
    model.eval()

    # load state dict
    model.load_state_dict(__lowerCamelCase    )

    # verify conversion
    _SCREAMING_SNAKE_CASE									=			"""http://images.cocodataset.org/val2017/000000039769.jpg"""

    _SCREAMING_SNAKE_CASE									=			BitImageProcessor(
        do_resize=__lowerCamelCase					,       size={"""shortest_edge""": 256}					,       resample=PILImageResampling.BILINEAR					,       do_center_crop=__lowerCamelCase					,       crop_size=224					,       do_normalize=__lowerCamelCase					,       image_mean=__lowerCamelCase					,       image_std=__lowerCamelCase					,       )
    _SCREAMING_SNAKE_CASE									=			Image.open(requests.get(__lowerCamelCase					,       stream=__lowerCamelCase    ).raw    )
    _SCREAMING_SNAKE_CASE									=			processor(images=__lowerCamelCase					,       return_tensors="""pt"""    )

    _SCREAMING_SNAKE_CASE									=			transforms.Compose(
        [
            transforms.Resize(256    ),
            transforms.CenterCrop(224    ),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406]					,       std=[0.229, 0.224, 0.225]    ),
        ]    )

    _SCREAMING_SNAKE_CASE									=			image_transforms(__lowerCamelCase    ).unsqueeze(0    )

    # verify pixel_values
    assert torch.allclose(inputs.pixel_values					,       __lowerCamelCase					,       atol=1e-4    )

    _SCREAMING_SNAKE_CASE									=			model(**__lowerCamelCase    )

    _SCREAMING_SNAKE_CASE									=			outputs.logits.argmax(-1    ).item()
    print("""Predicted class:"""					,       model.config.idalabel[predicted_class_idx]    )

    print("""First values of logits:"""					,       outputs.logits[0, :3]    )

    if model_name == "focalnet-tiny":
        _SCREAMING_SNAKE_CASE									=			torch.tensor([0.2166, -0.4368, 0.2191]    )
    elif model_name == "focalnet-tiny-lrf":
        _SCREAMING_SNAKE_CASE									=			torch.tensor([1.1669, 0.0125, -0.1695]    )
    elif model_name == "focalnet-small":
        _SCREAMING_SNAKE_CASE									=			torch.tensor([0.4917, -0.0430, 0.1341]    )
    elif model_name == "focalnet-small-lrf":
        _SCREAMING_SNAKE_CASE									=			torch.tensor([-0.2588, -0.5342, -0.2331]    )
    elif model_name == "focalnet-base":
        _SCREAMING_SNAKE_CASE									=			torch.tensor([-0.1655, -0.4090, -0.1730]    )
    elif model_name == "focalnet-base-lrf":
        _SCREAMING_SNAKE_CASE									=			torch.tensor([0.5306, -0.0483, -0.3928]    )
    assert torch.allclose(outputs.logits[0, :3]					,       __lowerCamelCase					,       atol=1e-4    )
    print("""Looks ok!"""    )

    if pytorch_dump_folder_path is not None:
        print(F'Saving model and processor of {model_name} to {pytorch_dump_folder_path}'    )
        model.save_pretrained(__lowerCamelCase    )
        processor.save_pretrained(__lowerCamelCase    )

    if push_to_hub:
        print(F'Pushing model and processor of {model_name} to the hub...'    )
        model.push_to_hub(F'{model_name}'    )
        processor.push_to_hub(F'{model_name}'    )


if __name__ == "__main__":
      lowercase_					= argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--model_name""",
          default="""focalnet-tiny""",
          type=str,
          help="""Name of the FocalNet model you'd like to convert.""",
      )
      parser.add_argument(
          """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
      )
      parser.add_argument(
          """--push_to_hub""",
          action="""store_true""",
          help="""Whether to push the model and processor to the hub.""",
      )

      lowercase_					= parser.parse_args()
      convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''

from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
      if not (is_transformers_available() and is_torch_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from ...utils.dummy_torch_and_transformers_objects import (
          ImageTextPipelineOutput,
          UniDiffuserPipeline,
      )
else:
      from .modeling_text_decoder import UniDiffuserTextDecoder
      from .modeling_uvit import UniDiffuserModel, UTransformeraDModel
      from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline



