




import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import (
    SegformerConfig,
    SegformerForImageClassification,
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase	:		Tuple							       =   logging.get_logger(__name__)







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_=False						)   ->       Union[str, Any]:



       """simple docstring"""

       A__       =    OrderedDict()
       for key, value in state_dict.items():
              if encoder_only and not key.startswith('''head'''						):
                     A__       =    '''segformer.encoder.''' + key
              if key.startswith('''backbone'''						):
                     A__       =    key.replace('''backbone''' ,				'''segformer.encoder'''						)
              if "patch_embed" in key:
                     # replace for example patch_embed1 by patch_embeddings.0
                     A__       =    key[key.find('''patch_embed'''						) + len('''patch_embed'''						)]
                     A__       =    key.replace(f"""patch_embed{idx}""" ,				f"""patch_embeddings.{int(lowercase_						)-1}"""						)
              if "norm" in key:
                     A__       =    key.replace('''norm''' ,				'''layer_norm'''						)
              if "segformer.encoder.layer_norm" in key:
                     # replace for example layer_norm1 by layer_norm.0
                     A__       =    key[key.find('''segformer.encoder.layer_norm'''						) + len('''segformer.encoder.layer_norm'''						)]
                     A__       =    key.replace(f"""layer_norm{idx}""" ,				f"""layer_norm.{int(lowercase_						)-1}"""						)
              if "layer_norm1" in key:
                     A__       =    key.replace('''layer_norm1''' ,				'''layer_norm_1'''						)
              if "layer_norm2" in key:
                     A__       =    key.replace('''layer_norm2''' ,				'''layer_norm_2'''						)
              if "block" in key:
                     # replace for example block1 by block.0
                     A__       =    key[key.find('''block'''						) + len('''block'''						)]
                     A__       =    key.replace(f"""block{idx}""" ,				f"""block.{int(lowercase_						)-1}"""						)
              if "attn.q" in key:
                     A__       =    key.replace('''attn.q''' ,				'''attention.self.query'''						)
              if "attn.proj" in key:
                     A__       =    key.replace('''attn.proj''' ,				'''attention.output.dense'''						)
              if "attn" in key:
                     A__       =    key.replace('''attn''' ,				'''attention.self'''						)
              if "fc1" in key:
                     A__       =    key.replace('''fc1''' ,				'''dense1'''						)
              if "fc2" in key:
                     A__       =    key.replace('''fc2''' ,				'''dense2'''						)
              if "linear_pred" in key:
                     A__       =    key.replace('''linear_pred''' ,				'''classifier'''						)
              if "linear_fuse" in key:
                     A__       =    key.replace('''linear_fuse.conv''' ,				'''linear_fuse'''						)
                     A__       =    key.replace('''linear_fuse.bn''' ,				'''batch_norm'''						)
              if "linear_c" in key:
                     # replace for example linear_c4 by linear_c.3
                     A__       =    key[key.find('''linear_c'''						) + len('''linear_c'''						)]
                     A__       =    key.replace(f"""linear_c{idx}""" ,				f"""linear_c.{int(lowercase_						)-1}"""						)
              if key.startswith('''head'''						):
                     A__       =    key.replace('''head''' ,				'''classifier'''						)
              A__       =    value

       return new_state_dict







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_						)   ->       str:



       """simple docstring"""

       for i in range(config.num_encoder_blocks						):
              for j in range(config.depths[i]						):
                     # read in weights + bias of keys and values (which is a single matrix in the original implementation)
                     A__       =    state_dict.pop(f"""segformer.encoder.block.{i}.{j}.attention.self.kv.weight"""						)
                     A__       =    state_dict.pop(f"""segformer.encoder.block.{i}.{j}.attention.self.kv.bias"""						)
                     # next, add keys and values (in that order) to the state dict
                     A__       =    kv_weight[
                         : config.hidden_sizes[i], :
                     ]
                     A__       =    kv_bias[: config.hidden_sizes[i]]
                     A__       =    kv_weight[
                         config.hidden_sizes[i] :, :
                     ]
                     A__       =    kv_bias[
                         config.hidden_sizes[i] :
                     ]







def SCREAMING_SNAKE_CASE				(       )   ->       Optional[Any]:



       """simple docstring"""

       A__       =    '''http://images.cocodataset.org/val2017/000000039769.jpg'''
       A__       =    Image.open(requests.get(lowercase_ ,				stream=lowercase_						).raw						)

       return image







@torch.no_grad()
def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_ ,				lowercase_						)   ->       List[str]:



       """simple docstring"""

       A__       =    SegformerConfig()
       A__       =    False

       # set attributes based on model_name
       A__       =    '''huggingface/label-files'''
       if "segformer" in model_name:
              A__       =    model_name[len('''segformer.'''						) : len('''segformer.'''						) + 2]
              if "ade" in model_name:
                     A__       =    150
                     A__       =    '''ade20k-id2label.json'''
                     A__       =    (1, 150, 128, 128)
              elif "city" in model_name:
                     A__       =    19
                     A__       =    '''cityscapes-id2label.json'''
                     A__       =    (1, 19, 128, 128)
              else:
                     raise ValueError(f"""Model {model_name} not supported"""						)
       elif "mit" in model_name:
              A__       =    True
              A__       =    model_name[4:6]
              A__       =    1_000
              A__       =    '''imagenet-1k-id2label.json'''
              A__       =    (1, 1_000)
       else:
              raise ValueError(f"""Model {model_name} not supported"""						)

       # set config attributes
       A__       =    json.load(open(hf_hub_download(lowercase_ ,				lowercase_ ,				repo_type='''dataset'''						) ,				'''r'''						)						)
       A__       =    {int(lowercase_						): v for k, v in idalabel.items()}
       A__       =    idalabel
       A__       =    {v: k for k, v in idalabel.items()}
       if size == "b0":
              pass
       elif size == "b1":
              A__       =    [64, 128, 320, 512]
              A__       =    256
       elif size == "b2":
              A__       =    [64, 128, 320, 512]
              A__       =    768
              A__       =    [3, 4, 6, 3]
       elif size == "b3":
              A__       =    [64, 128, 320, 512]
              A__       =    768
              A__       =    [3, 4, 18, 3]
       elif size == "b4":
              A__       =    [64, 128, 320, 512]
              A__       =    768
              A__       =    [3, 8, 27, 3]
       elif size == "b5":
              A__       =    [64, 128, 320, 512]
              A__       =    768
              A__       =    [3, 6, 40, 3]
       else:
              raise ValueError(f"""Size {size} not supported"""						)

       # load image processor (only resize + normalize)
       A__       =    SegformerImageProcessor(
           image_scale=(512, 512) ,				keep_ratio=lowercase_ ,				align=lowercase_ ,				do_random_crop=lowercase_						)

       # prepare image
       A__       =    prepare_img()
       A__       =    image_processor(images=lowercase_ ,				return_tensors='''pt'''						).pixel_values

       logger.info(f"""Converting model {model_name}..."""						)

       # load original state dict
       if encoder_only:
              A__       =    torch.load(lowercase_ ,				map_location=torch.device('''cpu'''						)						)
       else:
              A__       =    torch.load(lowercase_ ,				map_location=torch.device('''cpu'''						)						)['''state_dict''']

       # rename keys
       A__       =    rename_keys(lowercase_ ,				encoder_only=lowercase_						)
       if not encoder_only:
              del state_dict["decode_head.conv_seg.weight"]
              del state_dict["decode_head.conv_seg.bias"]

       # key and value matrices need special treatment
       read_in_k_v(lowercase_ ,				lowercase_						)

       # create HuggingFace model and load state dict
       if encoder_only:
              A__       =    False
              A__       =    SegformerForImageClassification(lowercase_						)
       else:
              A__       =    SegformerForSemanticSegmentation(lowercase_						)
       model.load_state_dict(lowercase_						)
       model.eval()

       # forward pass
       A__       =    model(lowercase_						)
       A__       =    outputs.logits

       # set expected_slice based on model name
       # ADE20k checkpoints
       if model_name == "segformer.b0.512x512.ade.160k":
              A__       =    torch.tensor(
                  [
                      [[-4.63_10, -5.52_32, -6.23_56], [-5.19_21, -6.14_44, -6.59_96], [-5.44_24, -6.27_90, -6.75_74]],
                      [[-12.13_91, -13.31_22, -13.95_54], [-12.87_32, -13.93_52, -14.35_63], [-12.94_38, -13.82_26, -14.25_13]],
                      [[-12.51_34, -13.46_86, -14.49_15], [-12.86_69, -14.43_43, -14.77_58], [-13.25_23, -14.58_19, -15.06_94]],
                  ]						)
       elif model_name == "segformer.b1.512x512.ade.160k":
              A__       =    torch.tensor(
                  [
                      [[-7.58_20, -8.72_31, -8.32_15], [-8.06_00, -10.35_29, -10.03_04], [-7.52_08, -9.41_03, -9.62_39]],
                      [[-12.69_18, -13.89_94, -13.71_37], [-13.31_96, -15.75_23, -15.47_89], [-12.93_43, -14.87_57, -14.96_89]],
                      [[-11.19_11, -11.94_21, -11.32_43], [-11.33_42, -13.68_39, -13.35_81], [-10.39_09, -12.18_32, -12.48_58]],
                  ]						)
       elif model_name == "segformer.b2.512x512.ade.160k":
              A__       =    torch.tensor(
                  [
                      [[-11.81_73, -14.38_50, -16.31_28], [-14.56_48, -16.58_04, -18.65_68], [-14.72_23, -15.73_87, -18.42_18]],
                      [[-15.72_90, -17.91_71, -19.44_23], [-18.31_05, -19.94_48, -21.46_61], [-17.92_96, -18.64_97, -20.79_10]],
                      [[-15.07_83, -17.03_36, -18.27_89], [-16.87_71, -18.68_70, -20.16_12], [-16.24_54, -17.14_26, -19.50_55]],
                  ]						)
       elif model_name == "segformer.b3.512x512.ade.160k":
              A__       =    torch.tensor(
                  [
                      [[-9.08_78, -10.20_81, -10.18_91], [-9.31_44, -10.79_41, -10.98_43], [-9.22_94, -10.38_55, -10.57_04]],
                      [[-12.23_16, -13.90_68, -13.61_02], [-12.91_61, -14.37_02, -14.32_35], [-12.52_33, -13.71_74, -13.79_32]],
                      [[-14.62_75, -15.24_90, -14.97_27], [-14.34_00, -15.96_87, -16.28_27], [-14.14_84, -15.40_33, -15.89_37]],
                  ]						)
       elif model_name == "segformer.b4.512x512.ade.160k":
              A__       =    torch.tensor(
                  [
                      [[-12.31_44, -13.24_47, -14.08_02], [-13.36_14, -14.58_16, -15.61_17], [-13.33_40, -14.44_33, -16.22_19]],
                      [[-19.27_81, -20.41_28, -20.75_06], [-20.61_53, -21.65_66, -22.09_98], [-19.98_00, -21.04_30, -22.14_94]],
                      [[-18.87_39, -19.78_04, -21.18_34], [-20.12_33, -21.67_65, -23.29_44], [-20.03_15, -21.26_41, -23.69_44]],
                  ]						)
       elif model_name == "segformer.b5.640x640.ade.160k":
              A__       =    torch.tensor(
                  [
                      [[-9.55_24, -12.08_35, -11.73_48], [-10.52_29, -13.64_46, -14.56_62], [-9.58_42, -12.88_51, -13.94_14]],
                      [[-15.34_32, -17.53_23, -17.08_18], [-16.33_30, -18.92_55, -19.21_01], [-15.13_40, -17.78_48, -18.39_71]],
                      [[-12.60_72, -14.94_86, -14.66_31], [-13.76_29, -17.09_07, -17.77_45], [-12.78_99, -16.16_95, -17.16_71]],
                  ]						)
       # Cityscapes checkpoints
       elif model_name == "segformer.b0.1024x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-11.92_95, -13.40_57, -14.81_06], [-13.34_31, -14.81_79, -15.37_81], [-14.28_36, -15.59_42, -16.15_88]],
                      [[-11.49_06, -12.80_67, -13.65_64], [-13.11_89, -14.05_00, -14.15_43], [-13.87_48, -14.51_36, -14.87_89]],
                      [[0.53_74, 0.10_67, -0.47_42], [0.11_41, -0.22_55, -0.70_99], [-0.30_00, -0.59_24, -1.31_05]],
                  ]						)
       elif model_name == "segformer.b0.512x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-7.82_17, -9.87_67, -10.17_17], [-9.44_38, -10.90_58, -11.40_47], [-9.79_39, -12.34_95, -12.10_79]],
                      [[-7.15_14, -9.53_36, -10.08_60], [-9.77_76, -11.68_22, -11.84_39], [-10.14_11, -12.76_55, -12.89_72]],
                      [[0.30_21, 0.08_05, -0.23_10], [-0.03_28, -0.16_05, -0.27_14], [-0.14_08, -0.54_77, -0.69_76]],
                  ]						)
       elif model_name == "segformer.b0.640x1280.city.160k":
              A__       =    torch.tensor(
                  [
                      [
                          [-1.1_3_7_2E0_1, -1.2_7_8_7E0_1, -1.3_4_7_7E0_1],
                          [-1.2_5_3_6E0_1, -1.4_1_9_4E0_1, -1.4_4_0_9E0_1],
                          [-1.3_2_1_7E0_1, -1.4_8_8_8E0_1, -1.5_3_2_7E0_1],
                      ],
                      [
                          [-1.4_7_9_1E0_1, -1.7_1_2_2E0_1, -1.8_2_7_7E0_1],
                          [-1.7_1_6_3E0_1, -1.9_1_9_2E0_1, -1.9_5_3_3E0_1],
                          [-1.7_8_9_7E0_1, -1.9_9_9_1E0_1, -2.0_3_1_5E0_1],
                      ],
                      [
                          [7.6_7_2_3E-0_1, 4.1_9_2_1E-0_1, -7.7_8_7_8E-0_2],
                          [4.7_7_7_2E-0_1, 9.5_5_5_7E-0_3, -2.8_0_8_2E-0_1],
                          [3.6_0_3_2E-0_1, -2.4_8_2_6E-0_1, -5.1_1_6_8E-0_1],
                      ],
                  ]						)
       elif model_name == "segformer.b0.768x768.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-9.49_59, -11.30_87, -11.74_79], [-11.00_25, -12.65_40, -12.33_19], [-11.40_64, -13.04_87, -12.99_05]],
                      [[-9.89_05, -11.30_84, -12.08_54], [-11.17_26, -12.76_98, -12.95_83], [-11.59_85, -13.32_78, -14.17_74]],
                      [[0.22_13, 0.01_92, -0.24_66], [-0.17_31, -0.42_13, -0.48_74], [-0.31_26, -0.65_41, -1.13_89]],
                  ]						)
       elif model_name == "segformer.b1.1024x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-13.57_48, -13.91_11, -12.65_00], [-14.35_00, -15.36_83, -14.23_28], [-14.75_32, -16.04_24, -15.60_87]],
                      [[-17.16_51, -15.87_25, -12.96_53], [-17.25_80, -17.37_18, -14.82_23], [-16.60_58, -16.87_83, -16.74_52]],
                      [[-3.64_56, -3.02_09, -1.42_03], [-3.07_97, -3.19_59, -2.00_00], [-1.87_57, -1.92_17, -1.69_97]],
                  ]						)
       elif model_name == "segformer.b2.1024x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-16.09_76, -16.48_56, -17.39_62], [-16.62_34, -19.03_42, -19.76_85], [-16.09_00, -18.06_61, -19.11_80]],
                      [[-18.47_50, -18.84_88, -19.50_74], [-19.40_30, -22.15_70, -22.59_77], [-19.11_91, -20.84_86, -22.37_83]],
                      [[-4.51_78, -5.50_37, -6.51_09], [-5.08_84, -7.21_74, -8.03_34], [-4.41_56, -5.81_17, -7.29_70]],
                  ]						)
       elif model_name == "segformer.b3.1024x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-14.20_81, -14.47_32, -14.19_77], [-14.58_67, -16.44_23, -16.63_56], [-13.44_41, -14.96_85, -16.86_96]],
                      [[-14.45_76, -14.70_73, -15.04_51], [-15.08_16, -17.62_37, -17.98_73], [-14.42_13, -16.01_99, -18.59_92]],
                      [[-4.73_49, -4.95_88, -5.09_66], [-4.32_10, -6.93_25, -7.25_91], [-3.43_12, -4.74_84, -7.19_17]],
                  ]						)
       elif model_name == "segformer.b4.1024x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-11.77_37, -11.95_26, -11.32_73], [-13.66_92, -14.45_74, -13.88_78], [-13.89_37, -14.69_24, -15.93_45]],
                      [[-14.67_06, -14.53_30, -14.13_06], [-16.15_02, -16.81_80, -16.42_69], [-16.83_38, -17.89_39, -20.17_46]],
                      [[1.04_91, 0.82_89, 1.03_10], [1.10_44, 0.52_19, 0.80_55], [1.08_99, 0.69_26, 0.55_90]],
                  ]						)
       elif model_name == "segformer.b5.1024x1024.city.160k":
              A__       =    torch.tensor(
                  [
                      [[-12.56_41, -13.47_77, -13.06_84], [-13.95_87, -15.89_83, -16.65_57], [-13.31_09, -15.73_50, -16.31_41]],
                      [[-14.70_74, -15.43_52, -14.59_44], [-16.63_53, -18.16_63, -18.61_20], [-15.17_02, -18.03_29, -18.15_47]],
                      [[-1.79_90, -2.09_51, -1.77_84], [-2.63_97, -3.82_45, -3.96_86], [-1.52_64, -2.81_26, -2.93_16]],
                  ]						)
       else:
              A__       =    logits.argmax(-1						).item()
              print('''Predicted class:''' ,				model.config.idalabel[predicted_class_idx]						)

       # verify logits
       if not encoder_only:
              assert logits.shape == expected_shape
              assert torch.allclose(logits[0, :3, :3, :3] ,				lowercase_ ,				atol=1E-2						)

       # finally, save model and image processor
       logger.info(f"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""						)
       Path(lowercase_						).mkdir(exist_ok=lowercase_						)
       model.save_pretrained(lowercase_						)
       image_processor.save_pretrained(lowercase_						)


if __name__ == "__main__":
  _lowerCamelCase	:		Any							       =   argparse.ArgumentParser()

  parser.add_argument(
      """--model_name""",
      default="""segformer.b0.512x512.ade.160k""",
      type=str,
      help="""Name of the model you'd like to convert.""",
  )
  parser.add_argument(
      """--checkpoint_path""", default=None, type=str, help="""Path to the original PyTorch checkpoint (.pth file)."""
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
  )
  _lowerCamelCase	:		Union[str, Any]							       =   parser.parse_args()
  convert_segformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)








import argparse

import torch
from safetensors.torch import load_file

from diffusers import StableDiffusionPipeline







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_ ,				lowercase_ ,				lowercase_ ,				lowercase_						)   ->       List[Any]:



       """simple docstring"""

       A__       =    StableDiffusionPipeline.from_pretrained(lowercase_ ,				torch_dtype=torch.floataa						)

       # load LoRA weight from .safetensors
       A__       =    load_file(lowercase_						)

       A__       =    []

       # directly update weight in diffusers model
       for key in state_dict:
              # it is suggested to print out the key, it usually will be something like below
              # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

              # as we have set the alpha beforehand, so just skip
              if ".alpha" in key or key in visited:
                     continue

              if "text" in key:
                     A__       =    key.split('''.'''						)[0].split(LORA_PREFIX_TEXT_ENCODER + '''_'''						)[-1].split('''_'''						)
                     A__       =    pipeline.text_encoder
              else:
                     A__       =    key.split('''.'''						)[0].split(LORA_PREFIX_UNET + '''_'''						)[-1].split('''_'''						)
                     A__       =    pipeline.unet

              # find the target layer
              A__       =    layer_infos.pop(0						)
              while len(lowercase_						) > -1:
                     try:
                            A__       =    curr_layer.__getattr__(lowercase_						)
                            if len(lowercase_						) > 0:
                                   A__       =    layer_infos.pop(0						)
                            elif len(lowercase_						) == 0:
                                   break
                     except Exception:
                            if len(lowercase_						) > 0:
                                   temp_name += "_" + layer_infos.pop(0						)
                            else:
                                   A__       =    layer_infos.pop(0						)

              A__       =    []
              if "lora_down" in key:
                     pair_keys.append(key.replace('''lora_down''' ,				'''lora_up'''						)						)
                     pair_keys.append(lowercase_						)
              else:
                     pair_keys.append(lowercase_						)
                     pair_keys.append(key.replace('''lora_up''' ,				'''lora_down'''						)						)

              # update weight
              if len(state_dict[pair_keys[0]].shape						) == 4:
                     A__       =    state_dict[pair_keys[0]].squeeze(3						).squeeze(2						).to(torch.floataa						)
                     A__       =    state_dict[pair_keys[1]].squeeze(3						).squeeze(2						).to(torch.floataa						)
                     curr_layer.weight.data += alpha * torch.mm(lowercase_ ,				lowercase_						).unsqueeze(2						).unsqueeze(3						)
              else:
                     A__       =    state_dict[pair_keys[0]].to(torch.floataa						)
                     A__       =    state_dict[pair_keys[1]].to(torch.floataa						)
                     curr_layer.weight.data += alpha * torch.mm(lowercase_ ,				lowercase_						)

              # update visited list
              for item in pair_keys:
                     visited.append(lowercase_						)

       return pipeline


if __name__ == "__main__":
  _lowerCamelCase	:		Tuple							       =   argparse.ArgumentParser()

  parser.add_argument(
      """--base_model_path""", default=None, type=str, required=True, help="""Path to the base model in diffusers format."""
  )
  parser.add_argument(
      """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
  )
  parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")
  parser.add_argument(
      """--lora_prefix_unet""", default="""lora_unet""", type=str, help="""The prefix of UNet weight in safetensors"""
  )
  parser.add_argument(
      """--lora_prefix_text_encoder""",
      default="""lora_te""",
      type=str,
      help="""The prefix of text encoder weight in safetensors""",
  )
  parser.add_argument("""--alpha""", default=0.75, type=float, help="""The merging ratio in W = W0 + alpha * deltaW""")
  parser.add_argument(
      """--to_safetensors""", action="""store_true""", help="""Whether to store pipeline in safetensors format or not."""
  )
  parser.add_argument("""--device""", type=str, help="""Device to use (e.g. cpu, cuda:0, cuda:1, etc.)""")

  _lowerCamelCase	:		Tuple							       =   parser.parse_args()

  _lowerCamelCase	:		List[Any]							       =   args.base_model_path
  _lowerCamelCase	:		Optional[int]							       =   args.checkpoint_path
  _lowerCamelCase	:		Dict							       =   args.dump_path
  _lowerCamelCase	:		Optional[Any]							       =   args.lora_prefix_unet
  _lowerCamelCase	:		Optional[int]							       =   args.lora_prefix_text_encoder
  _lowerCamelCase	:		List[Any]							       =   args.alpha

  _lowerCamelCase	:		int							       =   convert(base_model_path, checkpoint_path, lora_prefix_unet, lora_prefix_text_encoder, alpha)

  _lowerCamelCase	:		Tuple							       =   pipe.to(args.device)
  pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)


