Added dependencies

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
+import torch
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM, AutoConfig
-import json
 
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -9,11 +9,12 @@ subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENT
 # Function to load the model and processor
 @st.cache_resource
 def load_model_and_processor():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     config = AutoConfig.from_pretrained("microsoft/Florence-2-base-ft", trust_remote_code=True)
     config.vision_config.model_type = "davit"
-    model = AutoModelForCausalLM.from_pretrained("sujet-ai/Lutece-Vision-Base", config=config, trust_remote_code=True).eval()
+    model = AutoModelForCausalLM.from_pretrained("sujet-ai/Lutece-Vision-Base", config=config, trust_remote_code=True).to(device).eval()
     processor = AutoProcessor.from_pretrained("sujet-ai/Lutece-Vision-Base", config=config, trust_remote_code=True)
-    return model, processor
+    return model, processor, device
 
 # Function to generate answer
 def generate_answer(model, processor, device, image, prompt):
@@ -53,7 +54,7 @@ def main():
     st.sidebar.markdown("Our website : [sujet.ai](https://sujet.ai)")
 
     # Load model and processor
-    model, processor = load_model_and_processor()
+    model, processor, device = load_model_and_processor()
 
     # File uploader for document
     uploaded_file = st.file_uploader("📄 Upload a financial document", type=["png", "jpg", "jpeg"])
@@ -67,7 +68,7 @@ def main():
 
         if st.button("🔍 Generate Answer"):
             with st.spinner("Generating answer..."):
-                answer = generate_answer(model, processor, image, question)
+                answer = generate_answer(model, processor, device, image, question)
             st.success(f"## 💡 {answer}")
 
     # # Model configuration viewer