Download model before running workflow

app.py

@@ -4,7 +4,7 @@ from graph import BuildGraph
 from retriever import db_dir
 from langgraph.checkpoint.memory import MemorySaver
 from dotenv import load_dotenv
-from main import openai_model, model_id
+from main import openai_model, model_id, DownloadChatModel
 from util import get_sources, get_start_end_months
 from mods.tool_calling_llm import extract_think
 import requests
@@ -82,7 +82,6 @@ def run_workflow(input, history, compute_mode, thread_id, session_hash):
         if compute_mode == "local":
             gr.Info(
                 f"Please wait for the local model to load",
-                duration=15,
                 title=f"Model loading...",
             )
         # Get the chat model and build the graph
@@ -211,6 +210,11 @@ def to_workflow(request: gr.Request, *args):
     # Add session_hash to arguments
     new_args = args + (request.session_hash,)
     if compute_mode == "local":
+        # If graph hasn't been instantiated, download model before running workflow
+        graph = graph_instances[compute_mode].get(request.session_hash)
+        if graph is None:
+            DownloadChatModel()
+        # Call the workflow function with the @spaces.GPU decorator
         for value in run_workflow_local(*new_args):
             yield value
     if compute_mode == "remote":
@@ -218,7 +222,7 @@ def to_workflow(request: gr.Request, *args):
             yield value
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=90)
 def run_workflow_local(*args):
     for value in run_workflow(*args):
         yield value
@@ -264,13 +268,11 @@ with gr.Blocks(
             "local",
             "remote",
         ],
-        value=("local" if torch.cuda.is_available() else "remote"),
+        # Default to remote because it provides a better first impression for most people
+        # value=("local" if torch.cuda.is_available() else "remote"),
+        value="remote",
         label="Compute Mode",
-        info=(
-            "NOTE: remote mode **does not** use ZeroGPU"
-            if torch.cuda.is_available()
-            else "NOTE: local mode requires GPU"
-        ),
+        info="NOTE: remote mode **does not** use ZeroGPU",
         render=False,
     )
 

main.py

@@ -5,6 +5,7 @@ from langchain_core.output_parsers import StrOutputParser
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_core.messages import ToolMessage
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from huggingface_hub import snapshot_download
 from datetime import datetime
 from dotenv import load_dotenv
 import os
@@ -128,6 +129,13 @@ def ProcessDirectory(path, compute_mode):
             print(f"Chroma: no change for {file_path}")
 
 
+def DownloadChatModel():
+    """
+    Downloads a chat model to the local Hugging Face cache.
+    """
+    snapshot_download(model_id)
+
+
 def GetChatModel(compute_mode):
     """
     Get a chat model.