Update handler.py

handler.py

@@ -1,13 +1,28 @@
 from typing import Dict, List, Any
 import torch
-from transformers import LlavaNextVideoForConditionalGeneration, AutoProcessor, AutoConfig
+from transformers import LlavaNextVideoForConditionalGeneration, LlavaNextVideoProcessor
+from peft import PeftModel
 
 class EndpointHandler:
-    def __init__(self, path="/app"):
-        self.model = LlavaNextVideoForConditionalGeneration.from_pretrained(path)
+    def __init__(self):
+        self.base_model_name = "llava-hf/LLaVA-NeXT-Video-7B-hf"
+        self.adapter_model_name = "EnariGmbH/surftown-1.0"
 
-        # Load the processor from the configuration files
-        self.processor = AutoProcessor.from_pretrained(path)
+        # Load the base model
+        self.model = LlavaNextVideoForConditionalGeneration.from_pretrained(
+            self.base_model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+
+        # Load the fine-tuned adapter model into the base model
+        self.model = PeftModel.from_pretrained(self.model, self.adapter_model_name)
+
+        # Merge the adapter weights into the base model and unload the adapter
+        self.model = self.model.merge_and_unload()
+
+        # # Optionally, load and save the processor (if needed)
+        self.processor = LlavaNextVideoProcessor.from_pretrained(self.adapter_model_name)
 
         # Ensure the model is in evaluation mode
         self.model.eval()