app/main.py

import time
import os
import base64
from io import BytesIO
import concurrent.futures
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch_neuronx
import transformers
from transformers import AutoConfig
from llava.constants import MM_TOKEN_INDEX, DEFAULT_VIDEO_TOKEN, DEFAULT_VIDEO_START_TOKEN, DEFAULT_VIDEO_END_TOKEN
from llava.conversation import conv_templates
from llava.model.utils import LayerNorm
from llava.mm_utils import tokenizer_image_token, process_images_v2
from llava.model.multimodal_encoder.processor import Blip2ImageTrainProcessor

from typing import Dict
from fastapi import FastAPI, Request, HTTPException
from backend_model import MistralModel

# Suppress transformers logging
transformers.logging.set_verbosity_error()

NUM_SEGMENTS = 10  # Number of frame segments to use
WEIGHT_ROOT = '/root/inf2_dir/'  # Root directory for model weights
CONFIG_DIR = os.path.join(WEIGHT_ROOT, "llava-mistral_videollava_ptv12_350k_samep_65k_sopv2_065th_sopv1_fps2_scratch")  # Tokenizer directory
NEURON_VISION_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", "neuron_eva_vit_batch7.pth")  # Vision model weights (Neuron format)
NEURON_BERT_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", "neuron_bert.pth")  # BERT model weights (Neuron format)
PROJECTOR_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", 'projector.pth')  # Projector weights
EMBED_TOKEN_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", 'embed_tokens.pth')  # Embedding weights
QUERY_TOKEN_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", 'query_tokens.pth')
LAYERNORM_SAVE_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", 'ln_state_dict.pth')
POSITION_ENCODING_SAVE_PATH = os.path.join(WEIGHT_ROOT, "inf2_weights", 'frame_position_encoding.pth')


def generate_input_ids(tokenizer):
    """
    Generate input token IDs and conversation template for the model.

    Args:
        tokenizer (AutoTokenizer): Tokenizer instance.

    Returns:
        tuple: (input_ids, conv)
            input_ids (torch.Tensor): Input token IDs.
            conv (Conversation): Conversation template.
    """
    conv = conv_templates['thoth'].copy()  # Copy the conversation template
    qs = "Please describe this video in detail."
    qs = DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_TOKEN + DEFAULT_VIDEO_END_TOKEN + '\n' + qs  # Prepend video tokens
    conv.append_message(conv.roles[0], qs)  # Add the question to the conversation
    conv.append_message(conv.roles[1], None)  # Add a placeholder for the response
    prompt = conv.get_prompt()  # Get the conversation prompt
    input_ids = tokenizer_image_token(prompt, tokenizer, MM_TOKEN_INDEX, return_tensors='pt').unsqueeze(
        0)  # Tokenize and convert to tensor
    return input_ids


def uniform_sample(frames, num_segments):
    """
    Uniformly sample frames from the provided list.

    Args:
        frames (list): List of frame images.
        num_segments (int): Number of segments to sample.

    Returns:
        list: List of sampled frames.
    """
    indices = np.linspace(start=0, stop=len(frames) - 1, num=num_segments).astype(
        int)  # Calculate indices for uniform sampling
    frames = [frames[ind] for ind in indices]  # Sample frames based on indices
    return frames


def image_open_byteio(byte_data):
    output = Image.open(BytesIO(byte_data)).convert('RGB')
    return output

# Create FastAPI app
app = FastAPI()
mistral_model = MistralModel(model_name=CONFIG_DIR)  # Load Mistral model

# Load model configuration and tokenizer
config = AutoConfig.from_pretrained(CONFIG_DIR, trust_remote_code=True)
tokenizer = mistral_model.tokenizer

input_ids = generate_input_ids(tokenizer)  # Generate input IDs and conversation template
input_ids = input_ids[0].to('cpu')  # [token_len]

# Initialize image processor and vision module
image_processor = Blip2ImageTrainProcessor(
    image_size=config.img_size,
    is_training=False)

# vision_module = MASPVision()
# new_vision_state_dict = torch.load(VISION_STATE_DICT, map_location='cpu')  # Load vision state dict
# vision_module.load_state_dict(new_vision_state_dict)

with torch_neuronx.experimental.neuron_cores_context(start_nc=0, nc_count=2):  # Use Neuron cores for inference
    vision_module_neuron = torch.jit.load(NEURON_VISION_PATH)
vision_module_neuron = vision_module_neuron.eval()

# Load embedding weights and set up embedding module
padding_idx = config.pad_token_id
embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx)
embed_weight = torch.load(EMBED_TOKEN_PATH)
embed_tokens.load_state_dict(embed_weight)
embed_tokens = embed_tokens.eval()
embed_tokens.to(torch.float16).to('cpu')

#layer norm
vision_width = 1408
ln_vision = LayerNorm(vision_width)
ln_vision_weight = torch.load(LAYERNORM_SAVE_PATH)
ln_vision.load_state_dict(ln_vision_weight)
ln_vision = ln_vision.eval()
ln_vision = ln_vision.to(torch.float32)

num_query_token = 32
query_tokens = nn.Parameter(
    torch.zeros(1, num_query_token, 768)
)
query_tokens.data.normal_(mean=0.0, std=0.02)
query_tokens_weight = torch.load(QUERY_TOKEN_PATH)['query_tokens']
query_tokens.data = query_tokens_weight

frame_position_encoding = nn.Embedding(10, 768)
frame_position_encoding_weight = torch.load(POSITION_ENCODING_SAVE_PATH)
frame_position_encoding.load_state_dict(frame_position_encoding_weight)

projector = nn.Linear(config.mm_hidden_size, config.hidden_size)
projector_weight = torch.load(PROJECTOR_PATH)
projector.load_state_dict(projector_weight)

neuron_bert = torch.jit.load(NEURON_BERT_PATH)
neuron_bert = neuron_bert.eval()
def vision_module_neuron_output(image):
    """
    Get output from the vision module (Neuron format).

    Args:
        image (torch.Tensor): Input image.

    Returns:
        torch.Tensor: Vision module output.
    """
    output = vision_module_neuron(image)
    return output


@app.post("/generate")
async def generate(request: Request) -> Dict[str, str]:
    """
    Generate text using the Mistral model.

    Args:
        request (Request): The incoming request object.

    Returns:
        Dict[str, str]: A dictionary containing the generated text or an error message.
    """
    try:
        s1 = time.time()
        request_payload = await request.json()
        request_payload_keys = request_payload.keys()

        if "images" in request_payload_keys:  # If input is a list of images
            packed_data = request_payload.get("images")
            # unpacked_data = [base64.b64decode(item) for item in packed_data]  # Decode base64 images
            # input_images = [Image.open(BytesIO(byte_data)).convert('RGB') for byte_data in unpacked_data]  # Load images
            with concurrent.futures.ThreadPoolExecutor() as executor:
                unpacked_data = list(executor.map(base64.b64decode, packed_data))

            with concurrent.futures.ThreadPoolExecutor() as executor:
                input_images = list(executor.map(image_open_byteio, unpacked_data))

            input_images = uniform_sample(input_images, NUM_SEGMENTS)  # Sample frames
            input_images = process_images_v2(input_images, image_processor, config)  # Process images
            print("s1 - input_images time: ", time.time() - s1)

            si = time.time()
            with torch.inference_mode():  # Enable inference mode
                with concurrent.futures.ThreadPoolExecutor() as executor:  # Use thread pool for parallel processing
                    image_features_list = list(executor.map(vision_module_neuron_output, input_images))
                image_features = torch.cat(image_features_list, dim=0)  # Concatenate image features
                print("si - image_features neuron time: ", time.time() - si)

                s2 = time.time()
                # projector = nn.Linear(config.mm_hidden_size, config.hidden_size)  # Initialize projector
                # projector_weight = torch.load(PROJECTOR_DIR)
                # projector.load_state_dict(projector_weight)
                # image_features = vision_module.forward_features(input_images, image_features,
                #                                                 projector)  # Process vision features
                image_features = ln_vision(image_features)
                attn_mask = torch.ones(image_features.size()[:-1], dtype=torch.long).to(image_features.device)
                query_tokens_inputs = query_tokens.expand(image_features.shape[0], -1, -1)

                image_features = neuron_bert(
                    query_tokens_inputs.to(torch.float32),
                    image_features.to(torch.float32),
                    attn_mask.to(torch.int64)
                )["last_hidden_state"].to(torch.float32)

                frame_ids = torch.arange(input_images.shape[0], dtype=torch.long, device=image_features.device).unsqueeze(1)
                frame_ids = frame_ids.repeat(1, input_images.shape[1]).flatten(0, 1)  # [num_frames * num_patches]
                image_features += frame_position_encoding(frame_ids).unsqueeze(-2)  # [num_frames, 1, 768]
                projected_features = projector(image_features)

                image_features = projected_features.flatten(0, 1)
                print(image_features.shape)

                image_features.to(device='cpu', dtype=torch.float16)  # Convert to float16 and move to CPU
                print("s2 - image_features prepare time: ", time.time() - s2)

                s3 = time.time()
                vision_token_indice = torch.where(input_ids == MM_TOKEN_INDEX)[0][0]  # Get index of vision token
                pre_text_token = embed_tokens(input_ids[:vision_token_indice])  # Embed tokens before vision token
                post_text_token = embed_tokens(input_ids[vision_token_indice + 1:])  # Embed tokens after vision token
                print("s3 - text_token time: ", time.time() - s3)

                s4 = time.time()
                inputs_embeds = torch.cat([pre_text_token, image_features, post_text_token]).unsqueeze(
                    0)  # Concatenate input embeddings
                inputs = inputs_embeds.detach().cpu().numpy().tolist()  # Convert to list for input to Mistral model
                print("s4 - inputs time: ", time.time() - s4)

        elif "inputs" in request_payload_keys:  # If input is normal text or embedding
            si = time.time()
            inputs = request_payload.get("inputs")
        else:
            raise HTTPException(status_code=400, detail="Please provide correct input")

        if not inputs:
            raise HTTPException(status_code=400, detail="No input provided")

        s5 = time.time()
        parameters = request_payload.get("parameters", {})  # Get additional parameters
        generated_text = mistral_model.generate(inputs, parameters)  # Generate text using Mistral model
        print("s5 - generated_text time: ", time.time() - s5)
        print("total inference time: ", time.time() - si)
        return {"generated_text": generated_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating text: {str(e)}")