microsoft

/

kosmos-2.5

like 169

Follow

Microsoft 4,991

Model is great, but for batch inferencing, the problem is, even for different images, the output is the same, it seems to be just doing inferencing on a single image.
This is my test code:

from transformers import AutoProcessor, Kosmos2_5ForConditionalGeneration
from PIL import Image
import torch
from tqdm import tqdm
import os

DTYPE = torch.bfloat16

def initialize_model(model_name: str):
    print(f"Initializing {model_name} model")
    model = Kosmos2_5ForConditionalGeneration.from_pretrained(model_name, device_map="auto", torch_dtype=DTYPE)
    processor = AutoProcessor.from_pretrained(model_name)
    return model, processor

def test_kosmos_batch(
        model, 
        processor, 
        image_dir,
        prompt="<md>",
        batch_size=1,
        device = "cuda",
        dtype = torch.bfloat16,
        max_new_tokens=1024,
    ):

    # Read image paths from the image_dir string
    image_paths = [os.path.join(image_dir, img) for img in os.listdir(image_dir)]

    outputs = []
    num_batches = (len(image_paths) + batch_size - 1) // batch_size

    for i in tqdm(range(num_batches)):
        batch_paths = image_paths[i * batch_size : (i + 1) * batch_size]
        images = [Image.open(path) for path in batch_paths]

        inputs = processor(
            text=[prompt] * len(images),
            images=images,
            return_tensors="pt",
            padding=True,
        )

        inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
        inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
        print("Processing inputs completed")
        try:
            del inputs["width"]
            del inputs["height"]
        except KeyError:
            pass

        generated_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
        )
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
        outputs.extend(generated_text)
    
    for i, output in enumerate(outputs):
        print("=========================================")
        print(f"Output {i}: {output}")

    return outputs

if __name__ == "__main__":
    model_name = "microsoft/kosmos-2.5"
    model, processor = initialize_model(model_name)
    print("Model initialized successfully!")
    image_paths = "./images"
    test_kosmos_batch(model, processor, image_paths, batch_size=3, max_new_tokens=50)

Output is for the same image, which should not be the case:

=========================================
Output 0: <md>**Section 10.** **Effectivity.** This Circular shall take effect fifteen (15) calendar days following its publication either in the Official Gazette or in a newspaper of general circulation.

# Classification: GENERAL
=========================================
Output 1: <md>10\. *Effectivity*. This Circular shall take effect fifteen (15) calendar days following its publication either in the Official Gazette or in a newspaper of general circulation.

# Classification: GENERAL
=========================================
Output 2: <md># Section 10. Effectivity. This Circular shall take effect fifteen (15) calendar days following its publication in the Official Gazette or in a newspaper of general circulation.

# 1. The effectivity of this Circular shall be determined by the following