Initial import

.gitignore

@@ -0,0 +1 @@
+venv

handler.py

@@ -0,0 +1,25 @@
+from typing import Dict, List, Any
+import torch
+from transformers import BitsAndBytesConfig, pipeline
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16
+        )
+        # self.pipeline = pipeline("image-to-text", model="llava-hf/llava-1.5-7b-hf",
+        #                          model_kwargs={"quantization_config": quantization_config})
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        print(data)
+        inputs = data.pop("inputs", data)
+        image = data.pop("image", None)
+        prompt = data.pop("prompt", None)
+        # outputs = self.pipeline(image, prompt=prompt, generate_kwargs={
+        # "max_new_tokens": 200})
+        # return {"caption": outputs[0]["generated_text"]}
+        return {"image": image, "prompt": prompt}

requirements.txt

@@ -0,0 +1,2 @@
+transformers==4.38.2
+bitsandbytes