model_download_with_cache.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VQGAN Download and Caching Framework\n",
    "\n",
    "This is a framework for downloading and caching VQGAN models, many Jupyter Notebook platforms such as Google Colaboratory have local fast ephemeral storage along with slower persistent storage.  This framework works by downloading the files locally then copying them to persistent storage.  When the function is executed once the cache is primed the files are copied from the cache rather than downloaded again.\n",
    "\n",
    "## Business Continuity\n",
    "\n",
    "Sometimes models previously made available go missing or are temporarily unavailable, as this framework saves a copy it can be a useful business contunity mechanism that reduces the risk of interuptions to research.\n",
    "\n",
    "## Cache Invalidation\n",
    "\n",
    "This framework **does not implement** any form of cache invalidation, it is possible that the original configs and checkpoints could be updated meaning old models are cached."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import isdir, abspath, join\n",
    "from os import mkdir\n",
    "\n",
    "working_dir = abspath('./content') # working_dir must always be on a local disk\n",
    "model_dir = join(working_dir, 'models') # model_dir must always be on a local disk\n",
    "cache_dir = join(working_dir, 'cache') # cache should in general be on a persistent disk\n",
    "model_cache_dir = join(cache_dir, 'models') # models cache should in general be in a persistent disk\n",
    "\n",
    "# create directories if they don't exist.\n",
    "if not isdir(working_dir):\n",
    "    mkdir(working_dir)\n",
    "\n",
    "if not isdir(model_dir):\n",
    "    mkdir(model_dir)\n",
    "\n",
    "if not isdir(cache_dir):\n",
    "    mkdir(cache_dir)\n",
    "\n",
    "if not isdir(model_cache_dir):\n",
    "    mkdir(model_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from urllib.request import urlretrieve\n",
    "from os.path import isdir, isfile, join\n",
    "from shutil import copy\n",
    "\n",
    "if not (isdir(working_dir) and isdir(model_dir) and isdir(cache_dir) and isdir(model_cache_dir)):\n",
    "    raise RuntimeError('Directory not found, has the cell above been executed?')\n",
    "\n",
    "def download_vqgan_model(name, config_url, checkpoint_url):\n",
    "    config_file = join(model_dir, f'{name}.yaml')\n",
    "    checkpoint_file = join(model_dir, f'{name}.ckpt')\n",
    "    model_cache = join (model_cache_dir, name)\n",
    "    if not isdir(model_cache):\n",
    "        mkdir(model_cache)\n",
    "    config_cache = join(model_cache, f'{name}.yaml')\n",
    "    checkpoint_cache = join(model_cache, f'{name}.ckpt')\n",
    "\n",
    "    is_available = isfile(config_file) and isfile(checkpoint_file)\n",
    "    is_cached = isfile(config_cache) and isfile(checkpoint_cache)\n",
    "\n",
    "    if (is_available):\n",
    "        print(f'The model ({name}) is already available locally.')\n",
    "        return name\n",
    "\n",
    "    if (is_cached):\n",
    "        print(f'The model ({name}) is available in the cache, copying locally.')\n",
    "        copy(config_cache, config_file)\n",
    "        copy(checkpoint_cache, checkpoint_file)\n",
    "        return name\n",
    "\n",
    "    print(f'The model ({name}) was not found, downloading and caching.')\n",
    "    urlretrieve(config_url, config_file)\n",
    "    urlretrieve(checkpoint_url, checkpoint_file)\n",
    "    copy(config_file, config_cache)\n",
    "    copy(checkpoint_file, checkpoint_cache)\n",
    "    return name\n",
    "\n",
    "model = download_vqgan_model(\n",
    "    name='vqgan_imagenet_f16_1024',\n",
    "    config_url='https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1',\n",
    "    checkpoint_url='https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fckpts%2Flast.ckpt&dl=1')\n",
    "\n",
    "print(model)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "76e39a95be4d5898a3570c447c27edaa62298e3a74d4da47c7cca3211a7c59ee"
  },
  "kernelspec": {
   "display_name": "Python 3.9.10 64-bit (windows store)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}