utils.py

from __future__ import annotations

import gradio as gr
from huggingface_hub import whoami

def get_profile(profile: gr.OAuthProfile | None) -> str:
    if profile is None:
        return "Anonymous"
    return profile.username

def get_organizations(oauth_token: gr.OAuthToken | None) -> str:
    if oauth_token is None:
        return "No Organization"
    org_names = [org["name"] for org in whoami(oauth_token.token)["orgs"]]
    return org_names

def get_profile_and_organizations(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> tuple[str, str]:
    if profile is None:
        output_profile = "Anonymous"
    else:
        output_profile = profile.username
    
    if oauth_token is None:
        output_org = "No Organization"
    else:
        output_org = [org["name"] for org in whoami(oauth_token.token)["orgs"]]
    
    return output_profile, output_org

def download_with_restart(snapshot_download_func, repo_id, local_dir, repo_type, token, restart_func):
    try:
        print(local_dir)
        snapshot_download_func(
            repo_id=repo_id,
            local_dir=local_dir,
            repo_type=repo_type,
            tqdm_class=None,
            etag_timeout=30,
            token=token
        )
    except Exception:
        restart_func()