gpu support

handler.py

@@ -11,7 +11,7 @@ class EndpointHandler:
         self.processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
 
     
-        device = 'cpu' if torch.cuda.is_available() else 'cpu'
+        device = 'gpu' if torch.cuda.is_available() else 'cpu'
 
         model = LlavaNextForConditionalGeneration.from_pretrained(
             "llava-hf/llava-v1.6-mistral-7b-hf", 

test_handler_locally.py

@@ -0,0 +1,16 @@
+from handler import EndpointHandler
+
+# init handler
+my_handler = EndpointHandler(path=".")
+
+# prepare sample payload
+prompt = "Can you describe this picture focusing on specifics visual artifacts and ambiance (objects, colors, person, athmosphere..). Please stay concise only output keywords and concepts detected."
+files = [
+    {"path": "https://media.rolex.com/image/upload/q_auto/f_auto/c_limit,w_2440/v1708384234/rolexcom/about-rolex/hub/about-rolex-hub-cover-aca202310cw-0002-portrait"}
+]
+
+# test the handler
+results=my_handler({
+    'prompt': prompt,
+    'files': files
+})