Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| from time import perf_counter | |
| from typing import Any, List, Tuple, Union | |
| import numpy as np | |
| import requests | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoModelForCausalLM, LlamaTokenizer | |
| from inference.core.entities.requests.cogvlm import CogVLMInferenceRequest | |
| from inference.core.entities.responses.cogvlm import CogVLMResponse | |
| from inference.core.env import ( | |
| API_KEY, | |
| COGVLM_LOAD_4BIT, | |
| COGVLM_LOAD_8BIT, | |
| COGVLM_VERSION_ID, | |
| MODEL_CACHE_DIR, | |
| ) | |
| from inference.core.models.base import Model, PreprocessReturnMetadata | |
| from inference.core.utils.image_utils import load_image_rgb | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| class CogVLM(Model): | |
| def __init__(self, model_id=f"cogvlm/{COGVLM_VERSION_ID}", **kwargs): | |
| self.model_id = model_id | |
| self.endpoint = model_id | |
| self.api_key = API_KEY | |
| self.dataset_id, self.version_id = model_id.split("/") | |
| if COGVLM_LOAD_4BIT and COGVLM_LOAD_8BIT: | |
| raise ValueError( | |
| "Only one of environment variable `COGVLM_LOAD_4BIT` or `COGVLM_LOAD_8BIT` can be true" | |
| ) | |
| self.cache_dir = os.path.join(MODEL_CACHE_DIR, self.endpoint) | |
| with torch.inference_mode(): | |
| self.tokenizer = LlamaTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| f"THUDM/{self.version_id}", | |
| torch_dtype=torch.float16, | |
| low_cpu_mem_usage=True, | |
| trust_remote_code=True, | |
| load_in_4bit=COGVLM_LOAD_4BIT, | |
| load_in_8bit=COGVLM_LOAD_8BIT, | |
| cache_dir=self.cache_dir, | |
| ).eval() | |
| self.task_type = "lmm" | |
| def preprocess( | |
| self, image: Any, **kwargs | |
| ) -> Tuple[Image.Image, PreprocessReturnMetadata]: | |
| pil_image = Image.fromarray(load_image_rgb(image)) | |
| return pil_image, PreprocessReturnMetadata({}) | |
| def postprocess( | |
| self, | |
| predictions: Tuple[str], | |
| preprocess_return_metadata: PreprocessReturnMetadata, | |
| **kwargs, | |
| ) -> Any: | |
| return predictions[0] | |
| def predict(self, image_in: Image.Image, prompt="", history=None, **kwargs): | |
| images = [image_in] | |
| if history is None: | |
| history = [] | |
| built_inputs = self.model.build_conversation_input_ids( | |
| self.tokenizer, query=prompt, history=history, images=images | |
| ) # chat mode | |
| inputs = { | |
| "input_ids": built_inputs["input_ids"].unsqueeze(0).to(DEVICE), | |
| "token_type_ids": built_inputs["token_type_ids"].unsqueeze(0).to(DEVICE), | |
| "attention_mask": built_inputs["attention_mask"].unsqueeze(0).to(DEVICE), | |
| "images": [[built_inputs["images"][0].to(DEVICE).to(torch.float16)]], | |
| } | |
| gen_kwargs = {"max_length": 2048, "do_sample": False} | |
| with torch.inference_mode(): | |
| outputs = self.model.generate(**inputs, **gen_kwargs) | |
| outputs = outputs[:, inputs["input_ids"].shape[1] :] | |
| text = self.tokenizer.decode(outputs[0]) | |
| if text.endswith("</s>"): | |
| text = text[:-4] | |
| return (text,) | |
| def infer_from_request(self, request: CogVLMInferenceRequest) -> CogVLMResponse: | |
| t1 = perf_counter() | |
| text = self.infer(**request.dict()) | |
| response = CogVLMResponse(response=text) | |
| response.time = perf_counter() - t1 | |
| return response | |
| if __name__ == "__main__": | |
| m = CogVLM() | |
| m.infer() | |