diffusers/tools

1import argparse2import json3import os4import shutil5from diffusers.pipelines.stable_diffusion import safety_checker6import torch7from tempfile import TemporaryDirectory8from typing import List, Optional9from diffusers import StableDiffusionPipeline, ControlNetModel10 11from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download12from huggingface_hub.file_download import repo_folder_name13 14 15def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:16    info = api.model_info(model_id)17    filenames = set(s.rfilename for s in info.siblings)18 19    is_sd = "model_index.json" in filenames20 21    if is_sd:22        model = StableDiffusionPipeline.from_pretrained(model_id, from_flax=True, safety_checker=None)23    else:24        model = ControlNetModel.from_pretrained(model_id, from_flax=True)25 26    with TemporaryDirectory() as d:27        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))28        os.makedirs(folder)29 30        model.save_pretrained(folder)31        model.save_pretrained(folder, safe_serialization=True)32 33        if is_sd:34            model.to(torch_dtype=torch.float16)35        else:36            model.half()37 38        model.save_pretrained(folder, variant="fp16")39        model.save_pretrained(folder, safe_serialization=True, variant="fp16")40 41        api.upload_folder(42            folder_path=folder,43            repo_id=model_id,44            repo_type="model",45            create_pr=True,46        )47        print(model_id)48 49if __name__ == "__main__":50    DESCRIPTION = """51    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.52    It is PyTorch exclusive for now.53    It works by downloading the weights (PT), converting them locally, and uploading them back54    as a PR on the hub.55    """56    parser = argparse.ArgumentParser(description=DESCRIPTION)57    parser.add_argument(58        "model_id",59        type=str,60        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",61    )62    args = parser.parse_args()63    model_id = args.model_id64    api = HfApi()65    convert(api, model_id)66