from PIL import Image
import torch
import torch.nn as nn
import os
from mantis.models.mllava import chat_mllava
from mantis.models.mllava import MLlavaProcessor, LlavaForConditionalGeneration
import argparse

class Med_Mantis_Inference(nn.Module):
    def __init__(
        self,
        device="cuda:0",
        model_path="path_to_your_download_model",
    ):
        super(Med_Mantis_Inference, self).__init__()
        self.device = device
        self.model_path = model_path
        self.processor, self.model = self.load_model()

    def load_model(self):
        processor = MLlavaProcessor.from_pretrained(self.model_path)
        model = LlavaForConditionalGeneration.from_pretrained(
            self.merged_model_path,
            device_map=self.device,
            torch_dtype=torch.bfloat16,
            # attn_implementation="flash_attention_2",
        )
        return processor, model

    def forward(self, images_path, usr_text_list):
        if isinstance(usr_text_list, str):
            usr_text_list = [usr_text_list]

        images = []
        for img_path in images_path:
            images.append(Image.open(img_path).convert('RGB'))

        generation_kwargs = {
            "max_new_tokens": 1024,
            "num_beams": 1,
            "do_sample": False
        }

        response_list = []
        for qs_id, text in enumerate(usr_text_list):
            if qs_id == 0:
                response, history = chat_mllava(
                    text, images, self.model, self.processor, **generation_kwargs
                )
            else:
                response, history = chat_mllava(
                    text, images, self.model, self.processor, history=history, **generation_kwargs
                )
            response_list.append(response)

        if len(response_list) == 1:
            return response_list[0]
        else:
            return response_list


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument(
        "--model_path",
        type=str,
        default="path_to_your_download_model"
    )
    args = parser.parse_args()

    Med_Mantis_model = Med_Mantis_Inference(
        device=args.device,
        model_path=args.model_path
    )
    print("Successfully loaded the merged model")

    images_path = [
        "path_to_test_img_1",
        "path_to_test_img_2"
    ]
    usr_text_list = "User question"

    # Inference
    result = Med_Mantis_model(images_path, usr_text_list)
    print(result)

yangxk/Med-Mantis

Medical Large Vision Language Models with Multi-Image Visual Ability

Inference

yangxk/
Med-Mantis