# @title ## **Custom Download Corner**
import os
import shutil

%store -r

# @markdown Fill in the URL fields with the links to the files you want to download. Separate multiple URLs with a comma.
# @markdown Example: `url1, url2, url3`
os.chdir(root_dir)

custom_download_list = []

custom_model_url = "https://civitai.com/api/download/models/7425,https://civitai.com/api/download/models/11745,https://civitai.com/api/download/models/40248"  # @param {'type': 'string'}
custom_vae_url = "https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/VAEs/orangemix.vae.pt,https://huggingface.co/hakurei/waifu-diffusion-v1-4/blob/main/vae/kl-f8-anime2.ckpt"  # @param {'type': 'string'}
custom_embedding_url = "https://huggingface.co/yesyeahvh/bad-hands-5/blob/main/bad-hands-5.pt,https://huggingface.co/datasets/Nerfgun3/bad_prompt/blob/main/bad_prompt_version2.pt,https://huggingface.co/embed/EasyNegative/blob/main/EasyNegative.safetensors,https://huggingface.co/nick-x-hacker/bad-artist/resolve/main/bad-artist-anime.pt,  https://huggingface.co/nick-x-hacker/bad-artist/resolve/main/bad-artist.pt,  "  # @param {'type': 'string'}
custom_LoRA_url = "https://huggingface.co/suppcuti/riven/blob/main/Sigma99-Magallan.safetensors,https://huggingface.co/suppcuti/riven/blob/main/Kindred.safetensors,https://huggingface.co/suppcuti/riven/blob/main/riven-v1.2.safetensors"  # @param {'type': 'string'}
custom_hypernetwork_url = ""  # @param {'type': 'string'}
custom_control_url = ""  # @param {'type': 'string'}
custom_extensions_url = ""  # @param {'type': 'string'}
custom_components_url = ""  # @param {'type': 'string'}
custom_upscaler_url = ""  # @param {'type': 'string'}


def extract(url, dst):
    if not url.startswith("/content/"):
        filename = os.path.basename(url)
        zipfile = os.path.join(dst, filename)
    else:
        zipfile = url

    if url.endswith(".zip"):
        if os.path.exists(zipfile):
            !unzip -j -o {zipfile} -d "{dst}"
            os.remove(zipfile)
    elif url.endswith(".tar.lz4"):
        if os.path.exists(zipfile):
            !tar -xI lz4 -f {zipfile} --directory={dst}
            os.remove(zipfile)
    else:
        pass


def download(url_list, dst_dir, download_list, is_extensions):
    downloaded = []
    for url in tqdm(url_list, desc="Downloading Custom URLs"):
        if url:
            url = url.strip()
            download_list.append(url)
            basename = os.path.basename(url)

            with capture.capture_output() as cap:
                if is_extensions:
                    os.chdir(extensions_dir)
                    !git clone {url}
                elif url.startswith("/content/drive/MyDrive/"):
                    shutil.copy(url, dst_dir)
                elif url.startswith("https://drive.google.com"):
                    if "folders" in url:
                        !gdown --folder "{url}" -O {dst_dir} --fuzzy -c
                    else:
                        !gdown "{url}" -O {dst_dir} --fuzzy -c
                elif url.startswith("https://huggingface.co/"):
                    if "/blob/" in url:
                        url = url.replace("/blob/", "/resolve/")
                    hf_token = "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"
                    user_header = f'"Authorization: Bearer {hf_token}"'
                    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {dst_dir} -o {basename} {url}
                else:
                    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {dst_dir} {url}

                extract(url, dst_dir)

            del cap
            downloaded.append((basename, dst_dir))


custom_download_list = []

custom_dirs = {
    custom_model_url: models_dir,
    custom_vae_url: vaes_dir,
    custom_embedding_url: embeddings_dir,
    custom_LoRA_url: lora_dir,
    custom_hypernetwork_url: hypernetworks_dir,
    custom_control_url: control_dir,
    custom_extensions_url: extensions_dir,
    custom_components_url: components_dir,
    custom_upscaler_url: esrgan_dir,
}

print(f"Downloading...")
start_time = time.time()

for custom_url in [
    custom_model_url,
    custom_vae_url,
    custom_embedding_url,
    custom_LoRA_url,
    custom_hypernetwork_url,
    custom_control_url,
    custom_extensions_url,
    custom_components_url,
    custom_upscaler_url,
]:
    if custom_url:
        urls = custom_url.split(",")
        download(
            urls,
            custom_dirs[custom_url],
            custom_download_list,
            custom_url == custom_extensions_url,
        )

end_time = time.time()
elapsed_time = int(end_time - start_time)

if elapsed_time < 60:
    print(f"\nDownload completed. Took {elapsed_time} sec")
else:
    mins, secs = divmod(elapsed_time, 60)
    print(f"\nDownload completed. Took {mins} mins {secs} sec")

print("All is done! Go to the next step")