hugging face

Modules/PoseEstimation/pose_agent.py

@@ -3,6 +3,7 @@ from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.messages import HumanMessage
 from langchain_mistralai.chat_models import ChatMistralAI
+import torch
 import os
 import sys
 import json
@@ -11,6 +12,7 @@ from Modules.PoseEstimation.pose_estimator import model, compute_right_knee_angl
 
 # If api_key is not passed, default behavior is to use the `MISTRAL_API_KEY` environment variable.
 llm = ChatMistralAI(model='mistral-large-latest', api_key="i5jSJkCFNGKfgIztloxTMjfckiFbYBj4")
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 @tool
 def get_keypoints_from_keypoints(video_path: str) -> str:
@@ -27,7 +29,7 @@ def get_keypoints_from_keypoints(video_path: str) -> str:
     save_folder='tmp'
     os.makedirs(save_folder, exist_ok=True)
     keypoints = []
-    results = model(video_path, save=True, show_conf=False, show_boxes=False)
+    results = model(video_path, save=True, show_conf=False, show_boxes=False, device=device)
     for (i, frame) in enumerate(results):
         frame_dict = {}
         frame_dict['frame'] = i

app.py

@@ -52,7 +52,7 @@ with col1:
         with st.chat_message("user"):
             st.markdown(prompt)
 
-        with st.chat_message("assistant"):
+        with st.chat_message("assistant", avatar="data/AI_Bro.png"):
             # Build answer from LLM
             direction = router_chain.invoke({"question":prompt})