diffusers/tools

1#!/usr/bin/env python32import argparse3from huggingface_hub import HfApi4 5 6def main(api, model_id):7    info = api.list_repo_refs(model_id)8    branches = set([b.name for b in info.branches]) - set(["main"])9 10    return list(branches)11 12 13if __name__ == "__main__":14    DESCRIPTION = """15    Simple utility to get all branches from a repo16    """17    parser = argparse.ArgumentParser(description=DESCRIPTION)18    parser.add_argument(19        "--model_id",20        type=str,21        help="The name of the model on the hub to retrieve the branches from. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",22    )23 24    args = parser.parse_args()25    model_id = args.model_id26    api = HfApi()27    branches = main(api, model_id)28 29    if "non-ema" in branches:30        print(model_id)31#32#    if len(branches) > 0:33#        print(f"{model_id}: {branches}")34