get_pretrained_weight.py

from os import makedirs
from shutil import copyfileobj
import requests


def confirm(response):
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            return value

    return None


def download(id, path):
    base_url = 'https://drive.google.com/uc?export=download'

    session = requests.Session()
    res = session.get(base_url, params={'id': id}, stream=True)

    token = confirm(res)
    if token:
        res = session.get(base_url, params={'id': id, 'confirm': token}, stream=True)

    with open(path, 'wb') as f:
        MAX_ITER = 10000
        for content in res.iter_content(MAX_ITER):
            if content:
                f.write(content)

if __name__ == '__main__':
    dis_id = '1IiPpRQc90rK4vycHnUQIS4rtoKehzB5o'
    gen_id = '1Epo89V02NXmyVrzTKtPlWrrHNN-qcOD4'

    dest_dir = 'pretrained'
    makedirs(dest_dir, exist_ok=True)

    dis_path = dest_dir + '/dis.npz'
    gen_path = dest_dir + '/gen.npz'


    print('[Message] Downloading pre trained weight (usually it may take few minutes)')

    #Discriminator
    print('[Message] Now downloading Discriminator\'s weight... (about 41MB)')
    download(dis_id, dis_path)

    #Generator
    print('[Message] Now downloading Generator\'s weight... (about 76MB)')
    download(gen_id, gen_path)

    print('[Message] Successfully')