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
@@ -211,6 +211,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":

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.