| import base64 |
| import io |
| import json |
| import os |
| import time |
|
|
| import numpy as np |
| import torch |
| import triton_python_backend_utils as pb_utils |
| from PIL import Image |
| from transformers import AutoModelForCausalLM, AutoProcessor |
|
|
|
|
| class TritonPythonModel: |
| def initialize(self, args): |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
|
|
| base_path = args['model_repository'] |
| base_path = os.path.join(base_path, "1") |
| model_path = os.path.join(base_path, "icon_handler") |
|
|
| instance_device_id = str(args.get("model_instance_device_id", "0")).strip() |
| default_device = f"cuda:{instance_device_id}" if len(instance_device_id) > 0 else "cuda:0" |
| requested_device = os.environ.get( |
| "FLORENCE2_DEVICE", |
| os.environ.get("ICON_DEVICE", default_device), |
| ).strip().lower() |
| if requested_device == "auto": |
| requested_device = default_device |
| if not requested_device.startswith("cuda"): |
| raise RuntimeError( |
| f"florence2 requires GPU-only execution; invalid device '{requested_device}'. " |
| "Use FLORENCE2_DEVICE=cuda:<id>." |
| ) |
| if not torch.cuda.is_available(): |
| raise RuntimeError("florence2 requires CUDA, but torch.cuda.is_available() is False.") |
|
|
| self.device = requested_device |
| self.dtype = torch.float16 |
| self.model = AutoModelForCausalLM.from_pretrained( |
| model_path, |
| trust_remote_code=True, |
| local_files_only=True, |
| torch_dtype=self.dtype |
| ).to(self.device) |
| self.processor = AutoProcessor.from_pretrained( |
| model_path, |
| trust_remote_code=True, |
| local_files_only=True |
| ) |
| self.model.eval() |
| pb_utils.Logger.log_info(f"florence2 model initialized on {self.device} ({self.dtype})") |
|
|
| self.model_config = json.loads(args["model_config"]) |
| output_config = pb_utils.get_output_config_by_name(self.model_config, "CAPTIONS") |
| self.output_type = pb_utils.triton_string_to_numpy(output_config["data_type"]) |
|
|
| self.batch_size = 8 |
|
|
| def execute(self, requests): |
| logger = pb_utils.Logger |
|
|
| responses = [] |
| st = time.time() |
| for request in requests: |
| |
| prompt = "<CAPTION_ENG>" |
|
|
| image_bytes_list = pb_utils.get_input_tensor_by_name(request, "IMAGE_BYTES_LIST").as_numpy() |
| image_bytes_list = np.array(image_bytes_list).reshape(-1) |
| pil_images = [] |
| for image_bytes in image_bytes_list: |
| image_bytes = base64.b64decode(image_bytes.decode('utf-8')) |
| pil_image = Image.open(io.BytesIO(image_bytes)) |
| pil_images.append(pil_image) |
|
|
| batch_list = [] |
| for i in range(0, len(pil_images), self.batch_size): |
| batch_list.append(pil_images[i:i + self.batch_size]) |
|
|
| generated_text = [] |
| for batch in batch_list: |
| with torch.no_grad(): |
| inputs = self.processor( |
| text=[prompt] * len(batch), |
| images=batch, |
| return_tensors="pt", |
| padding=True, |
| do_resize=True |
| ) |
| generated_ids = self.model.generate( |
| input_ids=inputs["input_ids"].to(self.device), |
| pixel_values=inputs["pixel_values"].to(self.device, dtype=self.dtype), |
| max_new_tokens=8, |
| num_beams=1, |
| do_sample=False |
| ) |
| generated_text.extend(self.processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)) |
|
|
| out_tensor_0 = pb_utils.Tensor("CAPTIONS", np.array(generated_text).astype(self.output_type)) |
|
|
| responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0])) |
|
|
| logger.log_info(f"florence2 execute duration : {int((time.time() - st)*1000)} ms") |
| return responses |
|
|
