From 011e1a157c0c2c12edb1c81a6808f2dd691a9e68 Mon Sep 17 00:00:00 2001
From: Zack <zack@zackbradshaw.com>
Date: Wed, 6 Dec 2023 11:14:27 -0800
Subject: [PATCH] feat: add dowload dir

---
 .gitignore |  2 +-
 app.py     | 19 ++++++-------------
 2 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5d6957d1..03a3b533 100644
--- a/.gitignore
+++ b/.gitignore
@@ -34,7 +34,7 @@ error.txt
 
 errors.txt
 
-/models/
+models/
 
 # Distribution / packaging
 .Python
diff --git a/app.py b/app.py
index f971bbe2..aaeacf27 100644
--- a/app.py
+++ b/app.py
@@ -127,24 +127,17 @@ MAX_SLEEP_TIME = 40
 def download_model(model_url: str, memory_utilization: int):
     # Extract model name from the URL
     model_name = model_url.split('/')[-1]
-    # TODO continue debugging 
-    # response = requests.get(model_url, stream=True)
-    # total_size = int(response.headers.get('content-length', 0))
-    # block_size = 1024 #1 Kibibyte
-    # progress_bar = gr.outputs.Progress_Bar(total_size)
-    # model_data = b""
-    # for data in response.iter_content(block_size):
-        # model_data += data
-        # progress_bar.update(len(data))
-        # yield progress_bar
-    # Save the model data to a file, or load it into a model here
-    vllm_model = LLM(
+    # Download the model using VLLM
+    vllm_model = VLLM(
         model=model_url,
         trust_remote_code=True,
         gpu_memory_utilization=memory_utilization,
+        download_dir="models/"
     )
+    # Add the downloaded model to the available_models list
     available_models.append((model_name, vllm_model))
-    return gr.update(choices=available_models)
+    # Update the dropdown choices with the new available_models list
+    model_chosen.update(choices=available_models)
 
 def load_tools():
     global valid_tools_info