



'''simple docstring'''






import argparse
import json

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

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation
def       lowerCamelCase__       (						_A				):
  a						:					List[Any]        =     384
  if "tiny" in model_name:
    a						:					List[str]        =     [3, 3, 9, 3]
    a						:					Optional[int]        =     [96, 192, 384, 768]
  if "small" in model_name:
    a						:					int        =     [3, 3, 27, 3]
    a						:					Dict        =     [96, 192, 384, 768]
  if "base" in model_name:
    a						:					Union[str, Any]        =     [3, 3, 27, 3]
    a						:					str        =     [128, 256, 512, 1024]
    a						:					Any        =     512
  if "large" in model_name:
    a						:					Optional[Any]        =     [3, 3, 27, 3]
    a						:					Optional[Any]        =     [192, 384, 768, 1536]
    a						:					Optional[Any]        =     768
  if "xlarge" in model_name:
    a						:					str        =     [3, 3, 27, 3]
    a						:					Optional[Any]        =     [256, 512, 1024, 2048]
    a						:					List[str]        =     1024

  # set label information
  a						:					List[str]        =     150
  a						:					Union[str, Any]        =     'huggingface/label-files'
  a						:					int        =     'ade20k-id2label.json'
  a						:					Optional[int]        =     json.load(open(hf_hub_download(__lowercase ,			__lowercase ,			repo_type='dataset'				) ,			'r'				)				)
  a						:					Tuple        =     {int(__lowercase				): v for k, v in idalabel.items()}
  a						:					List[str]        =     {v: k for k, v in idalabel.items()}

  a						:					List[Any]        =     ConvNextConfig(
      depths=__lowercase ,			hidden_sizes=__lowercase ,			out_features=['stage1', 'stage2', 'stage3', 'stage4']				)
  a						:					List[str]        =     UperNetConfig(
      backbone_config=__lowercase ,			auxiliary_in_channels=__lowercase ,			num_labels=__lowercase ,			idalabel=__lowercase ,			labelaid=__lowercase ,			)

  return config
def       lowerCamelCase__       (						_A				):
  a						:					Optional[Any]        =     []

  # fmt: off
  # stem
  rename_keys.append(('backbone.downsample_layers.0.0.weight', 'backbone.embeddings.patch_embeddings.weight')				)
  rename_keys.append(('backbone.downsample_layers.0.0.bias', 'backbone.embeddings.patch_embeddings.bias')				)
  rename_keys.append(('backbone.downsample_layers.0.1.weight', 'backbone.embeddings.layernorm.weight')				)
  rename_keys.append(('backbone.downsample_layers.0.1.bias', 'backbone.embeddings.layernorm.bias')				)
  # stages
  for i in range(len(config.backbone_config.depths				)				):
    for j in range(config.backbone_config.depths[i]				):
      rename_keys.append((f"""backbone.stages.{i}.{j}.gamma""", f"""backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.depthwise_conv.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.dwconv.weight""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.depthwise_conv.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.dwconv.bias""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.norm.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.layernorm.weight""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.norm.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.layernorm.bias""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv1.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv1.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv2.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight""")				)
      rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv2.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias""")				)
    if i > 0:
      rename_keys.append((f"""backbone.downsample_layers.{i}.0.weight""", f"""backbone.encoder.stages.{i}.downsampling_layer.0.weight""")				)
      rename_keys.append((f"""backbone.downsample_layers.{i}.0.bias""", f"""backbone.encoder.stages.{i}.downsampling_layer.0.bias""")				)
      rename_keys.append((f"""backbone.downsample_layers.{i}.1.weight""", f"""backbone.encoder.stages.{i}.downsampling_layer.1.weight""")				)
      rename_keys.append((f"""backbone.downsample_layers.{i}.1.bias""", f"""backbone.encoder.stages.{i}.downsampling_layer.1.bias""")				)

    rename_keys.append((f"""backbone.norm{i}.weight""", f"""backbone.hidden_states_norms.stage{i+1}.weight""")				)
    rename_keys.append((f"""backbone.norm{i}.bias""", f"""backbone.hidden_states_norms.stage{i+1}.bias""")				)

  # decode head
  rename_keys.extend(
      [
          ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
          ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
          ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
          ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
      ]				)
  # fmt: on

  return rename_keys
def       lowerCamelCase__       (						_A ,			_A ,			_A				):
  a						:					Tuple        =     dct.pop(__lowercase				)
  a						:					List[Any]        =     val
def       lowerCamelCase__       (						_A ,			_A ,			_A				):
  a						:					Tuple        =     {
      'upernet-convnext-tiny': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth',
      'upernet-convnext-small': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth',
      'upernet-convnext-base': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth',
      'upernet-convnext-large': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth',
      'upernet-convnext-xlarge': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth',
  }
  a						:					Optional[Any]        =     model_name_to_url[model_name]
  a						:					List[Any]        =     torch.hub.load_state_dict_from_url(__lowercase ,			map_location='cpu'				)['state_dict']

  a						:					Dict        =     get_upernet_config(__lowercase				)
  a						:					Dict        =     UperNetForSemanticSegmentation(__lowercase				)
  model.eval()

  # replace "bn" => "batch_norm"
  for key in state_dict.copy().keys():
    a						:					str        =     state_dict.pop(__lowercase				)
    if "bn" in key:
      a						:					Union[str, Any]        =     key.replace('bn' ,			'batch_norm'				)
    a						:					int        =     val

  # rename keys
  a						:					Optional[Any]        =     create_rename_keys(__lowercase				)
  for src, dest in rename_keys:
    rename_key(__lowercase ,			__lowercase ,			__lowercase				)

  model.load_state_dict(__lowercase				)

  # verify on image
  a						:					Optional[Any]        =     'https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'
  a						:					Any        =     Image.open(requests.get(__lowercase ,			stream=__lowercase				).raw				).convert('RGB'				)

  a						:					str        =     SegformerImageProcessor()
  a						:					Any        =     processor(__lowercase ,			return_tensors='pt'				).pixel_values

  with torch.no_grad():
    a						:					List[str]        =     model(__lowercase				)

  if model_name == "upernet-convnext-tiny":
    a						:					List[Any]        =     torch.tensor(
        [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]				)
  elif model_name == "upernet-convnext-small":
    a						:					int        =     torch.tensor(
        [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]]				)
  elif model_name == "upernet-convnext-base":
    a						:					Optional[int]        =     torch.tensor(
        [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]]				)
  elif model_name == "upernet-convnext-large":
    a						:					str        =     torch.tensor(
        [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]]				)
  elif model_name == "upernet-convnext-xlarge":
    a						:					Any        =     torch.tensor(
        [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]]				)
  print('Logits:' ,			outputs.logits[0, 0, :3, :3]				)
  assert torch.allclose(outputs.logits[0, 0, :3, :3] ,			__lowercase ,			atol=1E-4				)
  print('Looks ok!'				)

  if pytorch_dump_folder_path is not None:
    print(f"""Saving model {model_name} to {pytorch_dump_folder_path}"""				)
    model.save_pretrained(__lowercase				)
    print(f"""Saving processor to {pytorch_dump_folder_path}"""				)
    processor.save_pretrained(__lowercase				)

  if push_to_hub:
    print(f"""Pushing model and processor for {model_name} to hub"""				)
    model.push_to_hub(f"""openmmlab/{model_name}"""				)
    processor.push_to_hub(f"""openmmlab/{model_name}"""				)


if __name__ == "__main__":
     lowerCAmelCase:					Dict						       =						argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default='upernet-convnext-tiny',
         type=str,
         choices=[F"upernet-convnext-{size}" for size in ['tiny', 'small', 'base', 'large', 'xlarge']],
         help='Name of the ConvNext UperNet 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 or not to push the converted model to the 🤗 hub.'
     )

     lowerCAmelCase:					List[Any]						       =						parser.parse_args()
     convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






'''simple docstring'''
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..models.auto import AutoProcessor
from ..models.vision_encoder_decoder import VisionEncoderDecoderModel
from ..utils import is_vision_available
from .base import PipelineTool


if is_vision_available():
    from PIL import Image

class      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_	:					Tuple					="naver-clova-ix/donut-base-finetuned-docvqa"
 SCREAMING_SNAKE_CASE_	:					Dict					=(
     "This is a tool that answers a question about an document (pdf). It takes an input named `document` which "
     "should be the document containing the information, as well as a `question` that is the question about the "
     "document. It returns a text that contains the answer to the question."
 )
 SCREAMING_SNAKE_CASE_	:					List[str]					="document_qa"
 SCREAMING_SNAKE_CASE_	:					Union[str, Any]					=AutoProcessor
 SCREAMING_SNAKE_CASE_	:					Union[str, Any]					=VisionEncoderDecoderModel

 SCREAMING_SNAKE_CASE_	:					List[Any]					=["image", "text"]
 SCREAMING_SNAKE_CASE_	:					Any					=["text"]






 def __init__(       self       :							Optional[int]							,	*__A       :							List[str]							,	**__A       :							List[Any]	):
        if not is_vision_available():
               raise ValueError('Pillow must be installed to use the DocumentQuestionAnsweringTool.'	)

        super().__init__(*__A							,	**__A	)






 def        _lowerCamelCase					(       self       :							Any							,	__A       :							"Image"							,	__A       :							str	):
        __UpperCamelCase       =  '<s_docvqa><s_question>{user_input}</s_question><s_answer>'
        __UpperCamelCase       =  task_prompt.replace('{user_input}'							,	__A	)
        __UpperCamelCase       =  self.pre_processor.tokenizer(
            __A							,	add_special_tokens=__A							,	return_tensors='pt'	).input_ids
        __UpperCamelCase       =  self.pre_processor(__A							,	return_tensors='pt'	).pixel_values

        return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values}






 def        _lowerCamelCase					(       self       :							Union[str, Any]							,	__A       :							Optional[Any]	):
        return self.model.generate(
            inputs['pixel_values'].to(self.device	)							,	decoder_input_ids=inputs['decoder_input_ids'].to(self.device	)							,	max_length=self.model.decoder.config.max_position_embeddings							,	early_stopping=__A							,	pad_token_id=self.pre_processor.tokenizer.pad_token_id							,	eos_token_id=self.pre_processor.tokenizer.eos_token_id							,	use_cache=__A							,	num_beams=1							,	bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]]							,	return_dict_in_generate=__A							,	).sequences






 def        _lowerCamelCase					(       self       :							Tuple							,	__A       :							List[Any]	):
        __UpperCamelCase       =  self.pre_processor.batch_decode(__A	)[0]
        __UpperCamelCase       =  sequence.replace(self.pre_processor.tokenizer.eos_token							,	''	)
        __UpperCamelCase       =  sequence.replace(self.pre_processor.tokenizer.pad_token							,	''	)
        __UpperCamelCase       =  re.sub(R'<.*?>'							,	''							,	__A							,	count=1	).strip()  # remove first task start token
        __UpperCamelCase       =  self.pre_processor.tokenajson(__A	)

        return sequence["answer"]
