| import argparse |
| import json |
| import os |
| import shutil |
| from diffusers.pipelines.stable_diffusion import safety_checker |
| import torch |
| from tempfile import TemporaryDirectory |
| from typing import List, Optional |
| from diffusers import StableDiffusionPipeline, ControlNetModel |
|
|
| from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download |
| from huggingface_hub.file_download import repo_folder_name |
|
|
|
|
| def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]: |
| info = api.model_info(model_id) |
| filenames = set(s.rfilename for s in info.siblings) |
|
|
| is_sd = "model_index.json" in filenames |
|
|
| if is_sd: |
| model = StableDiffusionPipeline.from_pretrained(model_id, from_flax=True, safety_checker=None) |
| else: |
| model = ControlNetModel.from_pretrained(model_id, from_flax=True) |
|
|
| with TemporaryDirectory() as d: |
| folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) |
| os.makedirs(folder) |
|
|
| model.save_pretrained(folder) |
| model.save_pretrained(folder, safe_serialization=True) |
|
|
| if is_sd: |
| model.to(torch_dtype=torch.float16) |
| else: |
| model.half() |
|
|
| model.save_pretrained(folder, variant="fp16") |
| model.save_pretrained(folder, safe_serialization=True, variant="fp16") |
|
|
| api.upload_folder( |
| folder_path=folder, |
| repo_id=model_id, |
| repo_type="model", |
| create_pr=True, |
| ) |
| print(model_id) |
|
|
| if __name__ == "__main__": |
| DESCRIPTION = """ |
| Simple utility tool to convert automatically some weights on the hub to `safetensors` format. |
| It is PyTorch exclusive for now. |
| It works by downloading the weights (PT), converting them locally, and uploading them back |
| as a PR on the hub. |
| """ |
| parser = argparse.ArgumentParser(description=DESCRIPTION) |
| parser.add_argument( |
| "model_id", |
| type=str, |
| help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`", |
| ) |
| args = parser.parse_args() |
| model_id = args.model_id |
| api = HfApi() |
| convert(api, model_id) |
|
|
