handler

handler.py

@@ -0,0 +1,133 @@
+from typing import Dict, List, Any
+from transformers import pipeline
+from PIL import Image
+import requests
+from transformers import AutoModelForCausalLM, LlamaTokenizer
+import torch
+from accelerate import (
+    init_empty_weights,
+    infer_auto_device_map,
+    load_checkpoint_and_dispatch,
+)
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Preload all the elements you are going to need at inference.
+
+        # self.pipeline = pipeline(
+        #     "text-generation", model="THUDM/cogvlm-chat-hf", trust_remote_code=True
+        # )
+
+        # self.model = AutoModelForCausalLM.from_pretrained(
+        #     "THUDM/cogvlm-chat-hf", trust_remote_code=True
+        # )
+
+        self.tokenizer = LlamaTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5")
+
+        # self.model = (
+        #     AutoModelForCausalLM.from_pretrained(
+        #         "THUDM/cogvlm-chat-hf",
+        #         torch_dtype=torch.bfloat16,
+        #         low_cpu_mem_usage=True,
+        #         trust_remote_code=True,
+        #     )
+        #     .to("cuda")
+        #     .eval()
+        # )
+
+        # DISTRIBUTED GPUS
+        with init_empty_weights():
+            self.model = AutoModelForCausalLM.from_pretrained(
+                "THUDM/cogvlm-chat-hf",
+                torch_dtype=torch.bfloat16,
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+            )
+
+        device_map = infer_auto_device_map(
+            self.model,
+            max_memory={
+                0: "16GiB",
+                1: "16GiB",
+                2: "16GiB",
+                3: "16GiB",
+                "cpu": "180GiB",
+            },
+            no_split_module_classes=["CogVLMDecoderLayer"],
+        )
+        self.model = load_checkpoint_and_dispatch(
+            self.model,
+            "~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/8abca878c4257412c4c38eeafaed3fe27a036730",
+            "~/.cache/huggingface/modules/transformers_modules/THUDM/cogvlm-chat-hf/8abca878c4257412c4c38eeafaed3fe27a036730",  # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
+            # "/home/ec2-user/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/8abca878c4257412c4c38eeafaed3fe27a036730",  # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
+            device_map=device_map,
+            no_split_module_classes=["CogVLMDecoderLayer"],
+        )
+        self.model = self.model.eval()
+        ## DISTRIBUTED GPUS
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+         data args:
+              inputs (:obj: `str` | `PIL.Image` | `np.array`)
+              kwargs
+        Return:
+              A :obj:`list` | `dict`: will be serialized and returned
+        """
+
+        query = data["query"]
+        img_uri = data["img_uri"]
+
+        image = Image.open(
+            requests.get(
+                img_uri,
+                stream=True,
+            ).raw
+        ).convert("RGB")
+
+        inputs = self.model.build_conversation_input_ids(
+            self.tokenizer,
+            query=query,
+            history=[],
+            images=[image],
+            template_version="vqa",
+        )  # vqa mode
+
+        inputs = {
+            "input_ids": inputs["input_ids"].unsqueeze(0).to("cuda"),
+            "token_type_ids": inputs["token_type_ids"].unsqueeze(0).to("cuda"),
+            "attention_mask": inputs["attention_mask"].unsqueeze(0).to("cuda"),
+            "images": [[inputs["images"][0].to("cuda").to(torch.bfloat16)]],
+        }
+
+        gen_kwargs = {"max_length": 2048, "do_sample": False}
+
+        with torch.no_grad():
+            outputs = self.model.generate(**inputs, **gen_kwargs)
+            outputs = outputs[:, inputs["input_ids"].shape[1] :]
+            response = self.tokenizer.decode(outputs[0])
+            return response
+
+
+# query = "How many houses are there in this cartoon?"
+# image = Image.open(
+#     requests.get(
+#         "https://github.com/THUDM/CogVLM/blob/main/examples/3.jpg?raw=true", stream=True
+#     ).raw
+# ).convert("RGB")
+# inputs = model.build_conversation_input_ids(
+#     tokenizer, query=query, history=[], images=[image], template_version="vqa"
+# )  # vqa mode
+# inputs = {
+#     "input_ids": inputs["input_ids"].unsqueeze(0).to("cuda"),
+#     "token_type_ids": inputs["token_type_ids"].unsqueeze(0).to("cuda"),
+#     "attention_mask": inputs["attention_mask"].unsqueeze(0).to("cuda"),
+#     "images": [[inputs["images"][0].to("cuda").to(torch.bfloat16)]],
+# }
+# gen_kwargs = {"max_length": 2048, "do_sample": False}
+
+# with torch.no_grad():
+#     outputs = model.generate(**inputs, **gen_kwargs)
+#     outputs = outputs[:, inputs["input_ids"].shape[1] :]
+#     print(tokenizer.decode(outputs[0]))

requirements.txt

@@ -0,0 +1,11 @@
+einops
+Pillow==10.1.0
+# torch==2.1.0
+torch==1.13.1
+# transformers==4.35.0
+accelerate==0.24.1
+sentencepiece==0.1.99
+einops==0.7.0
+# xformers==0.0.22.post7
+xformers
+triton==2.1.0