Update handler.py

handler.py

@@ -11,32 +11,39 @@ def base64_to_numpy(base64_str, shape):
     return arr.reshape(shape)
 
 class EndpointHandler:
-    def __init__(self):
-        self.base_model_name = "llava-hf/LLaVA-NeXT-Video-7B-hf"  # Replace with the original base model ID
-        self.adapter_model_name = "EnariGmbH/surftown-1.0"  # Your fine-tuned adapter model ID
+    def __init__(self, model_dir: str = None):
+        self.base_model_name = "llava-hf/LLaVA-NeXT-Video-7B-hf"
+        self.adapter_model_name = "EnariGmbH/surftown-1.0"
 
         # Load the base model
+        print("Loading base model:", self.base_model_name)
         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
+        print("Base model successfully loaded.")
+
+        # Load the adapter model into the base model
+        print("Loading adapter model:", self.adapter_model_name)
+        try:
+            self.model = PeftModel.from_pretrained(self.model, self.adapter_model_name)
+            print("Adapter model successfully loaded.")
+        except Exception as e:
+            print(f"Failed to load adapter model: {e}")
+            raise e
+
+        # Merge the adapter weights into the base model
         self.model = self.model.merge_and_unload()
+        print("Adapter model merged and unloaded.")
 
-        # # Save the full model
-        # model.save_pretrained("surftown_fine_tuned_prompt_0")
-
-        # # Optionally, load and save the processor (if needed)
+        # Load processor
         self.processor = LlavaNextVideoProcessor.from_pretrained(self.adapter_model_name)
+        print("Processor loaded.")
 
-        # Ensure the model is in evaluation mode
         self.model.eval()
 
+
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Args: