import hydra import pyrootutils import os import torch from omegaconf import OmegaConf from flask import Flask, request import json from typing import Optional import transformers from dataclasses import dataclass, field import io import base64 from PIL import Image import gc pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) BOI_TOKEN = '' EOI_TOKEN = '' IMG_TOKEN = '' IMG_FLAG = '' NUM_IMG_TOKNES = 32 NUM_IMG_CODES = 8192 app = Flask(__name__) def decode_image(encoded_image: str) -> Image: decoded_bytes = base64.b64decode(encoded_image.encode('utf-8')) buffer = io.BytesIO(decoded_bytes) image = Image.open(buffer) return image def encode_image(image: Image.Image, format: str = 'PNG') -> str: with io.BytesIO() as buffer: image.save(buffer, format=format) encoded_image = base64.b64encode(buffer.getvalue()).decode('utf-8') return encoded_image @dataclass class Arguments: image_transform: Optional[str] = field(default=None, metadata={"help": "config path of image transform"}) tokenizer: Optional[str] = field(default=None, metadata={"help": "config path of tokenizer used to initialize tokenizer"}) model: Optional[str] = field(default=None, metadata={"help": "config path of llm"}) port: Optional[str] = field(default=80, metadata={"help": "network port"}) llm_device: Optional[str] = field(default='cuda:0', metadata={"help": "llm device"}) tokenizer_device: Optional[str] = field(default='cuda:0', metadata={"help": "tokenizer device"}) offload_encoder: Optional[bool] = field(default=False, metadata={"help": "offload image tokenizer"}) offload_decoder: Optional[bool] = field(default=True, metadata={"help": "offload image tokenizer"}) parser = transformers.HfArgumentParser(Arguments) args, = parser.parse_args_into_dataclasses() class LLMService: def __init__(self, args) -> None: image_transform_cfg = OmegaConf.load(args.image_transform) tokenizer_cfg = OmegaConf.load(args.tokenizer) model_cfg = OmegaConf.load(args.model) self.image_id_shift = 32000 self.image_transform = hydra.utils.instantiate(image_transform_cfg) self.tokenizer = hydra.utils.instantiate(tokenizer_cfg, device=args.tokenizer_device, load_diffusion=True) if args.offload_encoder: self.tokenizer.image_tokenizer.model.visual_encoder.to('cpu') if args.offload_decoder: self.tokenizer.image_tokenizer.diffusion_model.to('cpu') # model = hydra.utils.instantiate(model_cfg, torch_dtype=torch.float16) # self.model = model.eval().to(args.llm_device) model = hydra.utils.instantiate(model_cfg, device_map=args.llm_device).eval() self.model = model print(model.get_memory_footprint()) self.llm_device = args.llm_device self.tokenizer_device = args.tokenizer_device self.offload_encoder = args.offload_encoder self.offload_decoder = args.offload_decoder self.boi_token_id = self.tokenizer(BOI_TOKEN, add_special_tokens=False).input_ids[0] self.eoi_token_id = self.tokenizer(EOI_TOKEN, add_special_tokens=False).input_ids[0] print('Init Done...') service = LLMService(args) @app.route('/generate', methods=['GET', 'POST']) def generate(): request_info = request.get_json() text_list = request_info['text'].split(IMG_FLAG) image_list = request_info['images'] temperature = request_info.get('temperature', 0.7) num_beams = request_info.get('num_beams', 1) max_new_tokens = request_info.get('max_new_tokens', 256) top_p = request_info.get('top_p', 0.5) force_boi = request_info.get('force_boi', False) assert len(text_list) == len(image_list) + 1 if len(image_list) > 0: images_tensor_list = [] images_tensor_indices = [] images_ids_list = [] images_ids_indices = [] for idx, image_item in enumerate(image_list): if isinstance(image_item, str): image = decode_image(image_item) image_tensor = service.image_transform(image) images_tensor_list.append(image_tensor) images_tensor_indices.append(idx) else: images_ids_list.append(image_item) images_ids_indices.append(idx) if len(images_tensor_list) > 0: images_tensor = torch.stack(images_tensor_list, dim=0).to(service.tokenizer_device) if service.offload_encoder: service.tokenizer.image_tokenizer.model.visual_encoder.to(service.tokenizer_device) images_ids_1 = service.tokenizer.encode_image(image_torch=images_tensor).cpu() if args.offload_encoder: service.tokenizer.image_tokenizer.model.visual_encoder.to('cpu') torch.cuda.empty_cache() gc.collect() num_image_ids = images_ids_1.shape[-1] else: num_image_ids = len(images_ids_list[-1]) images_ids_2 = torch.tensor(images_ids_list, dtype=torch.long) images_ids = torch.zeros((len(image_list), num_image_ids), dtype=torch.long) if len(images_tensor_indices) > 0: images_ids[images_tensor_indices, :] = images_ids_1 if len(images_ids_indices) > 0: images_ids[images_ids_indices, :] = images_ids_2 input_text = '' for i in range(images_ids.shape[0]): single_image_ids = images_ids[i].view(-1).tolist() image_tokens = BOI_TOKEN + ''.join([IMG_TOKEN.format(int(item)) for item in single_image_ids]) + EOI_TOKEN input_text += text_list[i] + image_tokens input_text = service.tokenizer.bos_token + input_text + text_list[-1] images_ids_list = images_ids.tolist() else: input_text = service.tokenizer.bos_token + ''.join(text_list) images_ids_list = [] if force_boi: input_text += BOI_TOKEN print(input_text) input_ids = service.tokenizer(input_text, add_special_tokens=False, return_tensors='pt').input_ids input_ids = input_ids.to(service.llm_device) generation_config = { 'temperature': temperature, 'num_beams': num_beams, 'max_new_tokens': max_new_tokens, 'top_p': top_p, 'do_sample': True } generate_ids = service.model.generate(input_ids=input_ids, **generation_config) if force_boi: generate_ids = generate_ids[0][input_ids.shape[1] - 1:] else: generate_ids = generate_ids[0][input_ids.shape[1]:] print('generated_ids: ', generate_ids) boi_indices = torch.where(generate_ids == service.boi_token_id)[0].tolist() eoi_indices = torch.where(generate_ids == service.eoi_token_id)[0].tolist() # assert len(boi_indices) == len(eoi_indices) generated_image_base64_list = [] text_mask = torch.ones_like(generate_ids, dtype=torch.bool) error_msg = [] if len(boi_indices) != len(eoi_indices): error_msg.append( f'Num of BOI (begain of image) tokens: {len(boi_indices)} is not equal to EOI(end of image tokens): {len(eoi_indices)}, some image Some images will fail to decode.' ) num_images = min(len(boi_indices), len(eoi_indices)) for idx in range(num_images): boi_index, eoi_index = boi_indices[idx], eoi_indices[idx] # for boi_index, eoi_index in zip(boi_indices, eoi_indices): image_ids = generate_ids[boi_index + 1:eoi_index].unsqueeze(0).to(service.tokenizer_device) image_ids = image_ids - service.image_id_shift if image_ids.shape[-1] != NUM_IMG_TOKNES: error_msg.append(f'Len(image_ids) {image_ids.shape[-1]} is not equal to {NUM_IMG_TOKNES}') image_base64 = '' elif (image_ids < 0).any() or (image_ids >= NUM_IMG_CODES).any(): error_msg.append(f'Some image_id out of range: [0, {NUM_IMG_CODES})') image_base64 = '' else: if service.offload_decoder: service.tokenizer.image_tokenizer.diffusion_model.to(service.tokenizer_device) image = service.tokenizer.decode_image(image_ids)[0] if service.offload_decoder: service.tokenizer.image_tokenizer.diffusion_model.to('cpu') torch.cuda.empty_cache() gc.collect() image_base64 = encode_image(image) generated_image_base64_list.append(image_base64) text_mask[boi_index + 1:eoi_index] = False images_ids_list.append(image_ids.view(-1).tolist()) generate_ids = generate_ids[text_mask] # print('generate_ids: ', generate_ids) # generate_text = service.tokenizer.decode(generate_ids, skip_special_tokens=True) generate_text = service.tokenizer.decode(generate_ids, skip_special_tokens=False) # print('generate_text before: ', generate_text) generate_text = generate_text.replace(BOI_TOKEN + ' ' + EOI_TOKEN + ' ', IMG_FLAG) generate_text = generate_text.replace(service.tokenizer.eos_token, '') print('generate_text: ', generate_text) return {'text': generate_text, 'images': generated_image_base64_list, 'images_ids': images_ids_list, 'error_msg': error_msg} if __name__ == '__main__': app.run(host='0.0.0.0', port=args.port)