From 994502d41b8ec876730b7b8b0dc61c1d5e5854de Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sun, 22 Oct 2023 08:57:16 -0700
Subject: [PATCH] Colab fixes

---
 Colab-TextGen-GPU.ipynb | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/Colab-TextGen-GPU.ipynb b/Colab-TextGen-GPU.ipynb
index c4fba4e8..d5525416 100644
--- a/Colab-TextGen-GPU.ipynb
+++ b/Colab-TextGen-GPU.ipynb
@@ -78,7 +78,7 @@
         "  !pip install -r extensions/api/requirements.txt --upgrade\n",
         "  !pip install -r temp_requirements.txt --upgrade\n",
         "\n",
-        "  print(\"\\033[1;32;1m\\n --> If you see a warning about \\\"previously imported packages\\\", just ignore it.\\n\\033[0;37;0m\")",
+        "  print(\"\\033[1;32;1m\\n --> If you see a warning about \\\"previously imported packages\\\", just ignore it.\\033[0;37;0m\")\n",
         "  print(\"\\033[1;32;1m\\n --> There is no need to restart the runtime.\\n\\033[0;37;0m\")\n",
         "\n",
         "  try:\n",
@@ -98,21 +98,27 @@
         "      command_line_flags += f\" {param}\"\n",
         "\n",
         "model_url = model_url.strip()\n",
-        "if not model_url.startswith('http'):\n",
-        "  model_url = 'https://huggingface.co/' + model_url\n",
+        "if model_url != \"\":\n",
+        "    if not model_url.startswith('http'):\n",
+        "        model_url = 'https://huggingface.co/' + model_url\n",
         "\n",
-        "# Download the model\n",
-        "url_parts = model_url.strip('/').strip().split('/')\n",
-        "output_folder = f\"{url_parts[-2]}_{url_parts[-1]}\"\n",
-        "branch = branch.strip('\"\\' ')\n",
-        "if branch.strip() != '':\n",
-        "  output_folder += f\"_{branch}\"\n",
-        "  !python download-model.py {model_url} --branch {branch}\n",
+        "    # Download the model\n",
+        "    url_parts = model_url.strip('/').strip().split('/')\n",
+        "    output_folder = f\"{url_parts[-2]}_{url_parts[-1]}\"\n",
+        "    branch = branch.strip('\"\\' ')\n",
+        "    if branch.strip() != '':\n",
+        "        output_folder += f\"_{branch}\"\n",
+        "        !python download-model.py {model_url} --branch {branch}\n",
+        "    else:\n",
+        "        !python download-model.py {model_url}\n",
         "else:\n",
-        "  !python download-model.py {model_url}\n",
+        "    output_folder = \"\"\n",
         "\n",
         "# Start the web UI\n",
-        "cmd = f\"python server.py --share --model {output_folder} {command_line_flags}\"\n",
+        "cmd = f\"python server.py --share\"\n",
+        "if output_folder != \"\":\n",
+        "    cmd += f\" --model {output_folder}\"\n",
+        "cmd += f\" {command_line_flags}\"\n",
         "print(cmd)\n",
         "!$cmd"
       ],