Merge pull request #32 from ackl/allow-download-models-from-gdrive

allow downloading models from public gdrive urls using gdown lib
This commit is contained in:
Cyberes 2022-12-09 21:34:29 -07:00 committed by GitHub
commit a01b0b5035
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 4 deletions

View File

@ -238,7 +238,8 @@
"\n",
"# For things in this notebook\n",
"!pip install requests\n",
" \n",
"!pip install gdown\n",
"\n",
"# Download popular custom scripts. This is basically remote code execution so be careful.\n",
"# See https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts\n",
"if download_scripts:\n",
@ -1405,14 +1406,12 @@
"model_uri = input('URI of model to download: ')\n",
"import re\n",
"import requests\n",
"import gdown\n",
"\n",
"def pretty_exit(str):\n",
" print(str)\n",
" raise\n",
"\n",
"if 'drive.google.com' in model_uri:\n",
" pretty_exit(\"Can't parse Google Drive links. Download the file yourself then upload it to the machine.\")\n",
"\n",
"magnet_match = re.search(r'magnet:\\?xt=urn:btih:[A-Za-z0-9&=%.]*', model_uri)\n",
"web_match = re.search(r'(https?:\\/\\/(?:www\\.|(?!www))[a-zA-Z0-9][a-zA-Z0-9-]+[a-zA-Z0-9]\\.[^\\s]{2,}|www\\.[a-zA-Z0-9][a-zA-Z0-9-]+[a-zA-Z0-9]\\.[^\\s]{2,}|https?:\\/\\/(?:www\\.|(?!www))[a-zA-Z0-9]+\\.[^\\s]{2,}|www\\.[a-zA-Z0-9]+\\.[^\\s]{2,})', model_uri)\n",
"web_dl_file = None\n",
@ -1433,13 +1432,19 @@
" web_dl_file = web_match[0].replace('/blob/', '/resolve/')\n",
" else:\n",
" web_dl_file = web_match[0]\n",
"elif 'drive.google.com' in model_uri:\n",
" (gdrive_file_id, _) = gdown.parse_url.parse_url(web_match[0])\n",
"elif web_match:\n",
" response = requests.head(web_match[0], allow_redirects=True)\n",
" pretty_exit(f'Wrong content-type: {response.headers[\"content-type\"].split(\";\")[0]}') if 'octet-stream' not in response.headers['content-type'] else None\n",
" web_dl_file = web_match[0]\n",
"\n",
"if web_dl_file is not None:\n",
" %cd \"{model_storage_dir}\"\n",
" !wget \"{web_dl_file}\" \n",
"elif gdrive_file_id is not None:\n",
" %cd \"{model_storage_dir}\"\n",
" gdown.download(f\"https://drive.google.com/uc?id={gdrive_file_id}&confirm=t\")\n",
"else:\n",
" print('Could not parse your URI')"
]