faisalhr1997/codeformer

1import math2import os3import requests4from torch.hub import download_url_to_file, get_dir5from tqdm import tqdm6from urllib.parse import urlparse7 8from .misc import sizeof_fmt9 10 11def download_file_from_google_drive(file_id, save_path):12    """Download files from google drive.13    Ref:14    https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive  # noqa E50115    Args:16        file_id (str): File id.17        save_path (str): Save path.18    """19 20    session = requests.Session()21    URL = 'https://docs.google.com/uc?export=download'22    params = {'id': file_id}23 24    response = session.get(URL, params=params, stream=True)25    token = get_confirm_token(response)26    if token:27        params['confirm'] = token28        response = session.get(URL, params=params, stream=True)29 30    # get file size31    response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})32    print(response_file_size)33    if 'Content-Range' in response_file_size.headers:34        file_size = int(response_file_size.headers['Content-Range'].split('/')[1])35    else:36        file_size = None37 38    save_response_content(response, save_path, file_size)39 40 41def get_confirm_token(response):42    for key, value in response.cookies.items():43        if key.startswith('download_warning'):44            return value45    return None46 47 48def save_response_content(response, destination, file_size=None, chunk_size=32768):49    if file_size is not None:50        pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')51 52        readable_file_size = sizeof_fmt(file_size)53    else:54        pbar = None55 56    with open(destination, 'wb') as f:57        downloaded_size = 058        for chunk in response.iter_content(chunk_size):59            downloaded_size += chunk_size60            if pbar is not None:61                pbar.update(1)62                pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')63            if chunk:  # filter out keep-alive new chunks64                f.write(chunk)65        if pbar is not None:66            pbar.close()67 68 69def load_file_from_url(url, model_dir=None, progress=True, file_name=None):70    """Load file form http url, will download models if necessary.71    Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py72    Args:73        url (str): URL to be downloaded.74        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.75            Default: None.76        progress (bool): Whether to show the download progress. Default: True.77        file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.78    Returns:79        str: The path to the downloaded file.80    """81    if model_dir is None:  # use the pytorch hub_dir82        hub_dir = get_dir()83        model_dir = os.path.join(hub_dir, 'checkpoints')84 85    os.makedirs(model_dir, exist_ok=True)86 87    parts = urlparse(url)88    filename = os.path.basename(parts.path)89    if file_name is not None:90        filename = file_name91    cached_file = os.path.abspath(os.path.join(model_dir, filename))92    if not os.path.exists(cached_file):93        print(f'Downloading: "{url}" to {cached_file}\n')94        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)95    return cached_file