From be3d3cb8cc5c0e004d41c35fd8ba60b1fc3fc7b0 Mon Sep 17 00:00:00 2001 From: Drake Panzer <64224601+Cyberes@users.noreply.github.com> Date: Tue, 3 Jan 2023 12:44:39 -0700 Subject: [PATCH] clean exit from automated model downloader --- StableDiffusionUI_Voldemort_paperspace.ipynb | 56 +++++++++----------- 1 file changed, 26 insertions(+), 30 deletions(-) diff --git a/StableDiffusionUI_Voldemort_paperspace.ipynb b/StableDiffusionUI_Voldemort_paperspace.ipynb index 8711d76..542ee8c 100644 --- a/StableDiffusionUI_Voldemort_paperspace.ipynb +++ b/StableDiffusionUI_Voldemort_paperspace.ipynb @@ -147,7 +147,8 @@ "outputs": [], "source": [ "import os\n", - "# You'll see this little code block at the beginning of every cell. It makes sure you have ran the first block that defines your settings.\n", + "# You'll see this little code block at the beginning of every cell.\n", + "# It makes sure you have ran the first block that defines your settings.\n", "try:\n", " %store -r symlink_to_notebooks model_storage_dir repo_storage_dir\n", " test = [symlink_to_notebooks, model_storage_dir, repo_storage_dir]\n", @@ -1168,7 +1169,6 @@ { "cell_type": "markdown", "metadata": { - "jp-MarkdownHeadingCollapsed": true, "tags": [] }, "source": [ @@ -1191,8 +1191,8 @@ }, "outputs": [], "source": [ - "%store -r model_storage_dir repo_storage_dir\n", "try:\n", + " %store -r model_storage_dir repo_storage_dir\n", " test = [model_storage_dir, repo_storage_dir]\n", "except NameError as e:\n", " print(\"There is an issue with your variables.\")\n", @@ -1424,51 +1424,47 @@ "import re\n", "import requests\n", "import gdown\n", + " \n", + "def dl_web_file(web_dl_file):\n", + " %cd \"{model_storage_dir}\"\n", + " # We're going to use aria2 to split the download into threads which will allow us to download\n", + " # the file very fast even if the site serves the file slow.\n", + " !if [ $(dpkg-query -W -f='${Status}' aria2 2>/dev/null | grep -c \"ok installed\") = 0 ]; then sudo apt update && sudo apt install -y aria2; fi\n", + " !aria2c --file-allocation=none -c -x 16 -s 16 --summary-interval=0 \"{web_dl_file}\" \n", "\n", - "def pretty_exit(str):\n", - " print(str)\n", - " raise\n", - "\n", - "magnet_match = re.search(r'magnet:\\?xt=urn:btih:[A-Za-z0-9&=%.]*', model_uri)\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", - "gdrive_file_id = None\n", "\n", "if magnet_match:\n", - " !apt update\n", - " !apt install -y aria2\n", + " !if [ $(dpkg-query -W -f='${Status}' aria2 2>/dev/null | grep -c \"ok installed\") = 0 ]; then sudo apt update && sudo apt install -y aria2; fi\n", " %cd \"{model_storage_dir}\"\n", " bash_var = magnet_match[0]\n", " !aria2c --seed-time=0 --max-overall-upload-limit=1K --bt-max-peers=120 --summary-interval=0 --file-allocation=none \"{bash_var}\"\n", - " exit\n", "elif 'https://huggingface.co/' in model_uri:\n", " response = requests.head(web_match[0], allow_redirects=True)\n", " if 'octet-stream' not in response.headers['content-type']:\n", " response = requests.head(web_match[0].replace('/blob/', '/resolve/'), allow_redirects=True)\n", " if 'octet-stream' not in response.headers['content-type']:\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", + " print(f'Wrong content-type: {response.headers[\"content-type\"].split(\";\")[0]}')\n", + " # clean exit here\n", " else:\n", - " web_dl_file = web_match[0].replace('/blob/', '/resolve/')\n", + " dl_web_file(web_match[0].replace('/blob/', '/resolve/'))\n", " else:\n", - " web_dl_file = web_match[0]\n", + " dl_web_file(web_match[0])\n", "elif 'https://drive.google.com' in model_uri:\n", - " (gdrive_file_id, _) = gdown.parse_url.parse_url(web_match[0])\n", + " gdrive_file_id, _ = gdown.parse_url.parse_url(web_match[0])\n", + " %cd \"{model_storage_dir}\"\n", + " gdown.download(f\"https://drive.google.com/uc?id={gdrive_file_id}&confirm=t\") \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", - " # We're going to use aria2 to split the download into threads which will allow us to download\n", - " # the file very fast even if the connection is slow.\n", - " !if [ $(dpkg-query -W -f='${Status}' aria2 2>/dev/null | grep -c \"ok installed\") = 0 ]; then sudo apt update && sudo apt install -y aria2; fi\n", - " !aria2c --file-allocation=none -c -x 16 -s 16 --summary-interval=0 \"{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", + " if 'octet-stream' not in response.headers['content-type']:\n", + " print(f'Wrong content-type: {response.headers[\"content-type\"].split(\";\")[0]}')\n", + " # clean exit here\n", + " else:\n", + " dl_web_file(web_match[0])\n", "else:\n", - " print('Could not parse your URI.')" + " print('Could not parse your URI.')\n", + " # clean exit here" ] } ],