Update DocBrain.py

DocBrain.py

@@ -1,6 +1,6 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForVision2Seq
 import torch
+from transformers import AutoProcessor, AutoModelForVision2Seq
 from PaitentVoiceToText import record_and_transcribe  # Your STT function
 from DocVoice import text_to_speech  # Your TTS function
 
@@ -8,15 +8,15 @@ from DocVoice import text_to_speech  # Your TTS function
 # 1️⃣ Load Model & Processor
 # -------------------
 def load_model():
-    local_dir = r"C:\Users\JAY\Downloads\model\CHATDOCMODEL"
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if device == "cuda" else torch.float32
 
-    processor = AutoProcessor.from_pretrained(local_dir, trust_remote_code=True)
+    # Load directly from Hugging Face
+    processor = AutoProcessor.from_pretrained("Muhammadidrees/RaiyaChatDoc", trust_remote_code=True)
     model = AutoModelForVision2Seq.from_pretrained(
-        local_dir,
-        dtype=dtype,
-        device_map=None
+        "Muhammadidrees/RaiyaChatDoc",
+        torch_dtype=dtype,
+        device_map="auto"  # automatically assigns to GPU if available
     )
     model.to(device)
     return processor, model, device
@@ -181,7 +181,6 @@ with gr.Blocks(title="ChatDOC", theme=gr.themes.Soft()) as demo:
 
     def mic_submit(history, question_count, assistant_responses):
         user_text = record_and_transcribe(duration=5)
-        # Show user message immediately
         history.append([user_text, None])
         history, updated_history, question_count = process_message(user_text, history, question_count)
         assistant_responses = update_assistant_responses(history, assistant_responses)