Create vision_wrapper.py

751560b verified over 1 year ago

16 kB

import math
import torch
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from PIL import Image
import numpy as np
import io

from transformers import AutoProcessor
from transformers.models.paligemma.modeling_paligemma import (
    PaliGemmaConfig,
    PaliGemmaForConditionalGeneration,
    PaliGemmaPreTrainedModel,
)

from transformers.models.qwen2_vl import (
    Qwen2VLForConditionalGeneration, 
    Qwen2VLConfig,
)

class ColPali(PaliGemmaPreTrainedModel):
    """
    ColPali model implementation from the "ColPali: Efficient Document Retrieval with Vision Language Models" paper.
    """
    def __init__(self, config: PaliGemmaConfig):
        super().__init__(config=config)
        model = PaliGemmaForConditionalGeneration(config=config)
        if model.language_model._tied_weights_keys is not None:
            self._tied_weights_keys = [f"model.language_model.{k}" for k in model.language_model._tied_weights_keys]
        self.model = model
        self.dim = 128
        self.custom_text_proj = nn.Linear(self.model.config.text_config.hidden_size, self.dim)
        self.post_init()

    def forward(self, *args, **kwargs) -> torch.Tensor:
        # Delete output_hidden_states from kwargs
        kwargs.pop("output_hidden_states", None)
        outputs = self.model(*args, output_hidden_states=True, **kwargs)  # (batch_size, sequence_length, hidden_size)
        last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
        proj = self.custom_text_proj(last_hidden_states)  # (batch_size, sequence_length, dim)
        # L2 normalization
        proj = proj / proj.norm(dim=-1, keepdim=True)  # (batch_size, sequence_length, dim)
        proj = proj * kwargs["attention_mask"].unsqueeze(-1)  # (batch_size, sequence_length, dim)
        return proj


class ColPaliRetriever():
    def __init__(self, bs=4, use_gpu=True):
        self.bs = bs
        self.bs_query = 64
        self.model_name = "checkpoint/colpali-v1.1"
        self.base_ckpt = "checkpoint/colpaligemma-3b-mix-448-base"
        # Load model on cuda:0 by default!
        device = "cuda:0" if (torch.cuda.is_available() and use_gpu) else "cpu"
        self.model = ColPali.from_pretrained(
            self.base_ckpt, torch_dtype=torch.bfloat16, device_map=None  # <-- NONE: Don't use device_map
        )
        self.model.load_adapter(self.model_name)
        self.model = self.model.to(device)
        self.model.eval()
        # Multi-GPU with DataParallel
        if torch.cuda.device_count() > 1 and use_gpu:
            print(f"[ColPaliRetriever] Using DataParallel on {torch.cuda.device_count()} GPUs")
            self.model = torch.nn.DataParallel(self.model)
            self.device = torch.device("cuda:0")
        else:
            self.device = torch.device(device)
        print(f"[ColPaliRetriever - init] ColPali loaded from '{self.base_ckpt}' (Adapter '{self.model_name}')...")
        # Load processor
        self.processor = AutoProcessor.from_pretrained(self.model_name)
        self.mock_image = Image.new("RGB", (16, 16), color="black")


    def embed_queries(self, queries, pad=False):
        if isinstance(queries, str):
            queries = [queries]
        embeddings = []
        dataloader = DataLoader(queries, batch_size=self.bs_query, shuffle=False, 
                                collate_fn=lambda x: self.process_queries(x))
        with torch.no_grad():
            for batch in tqdm(dataloader, desc="[ColPaliRetriever] Embedding queries"):
                batch = {k: v.to(self.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                attention_mask = batch["attention_mask"]
                if isinstance(outputs, (tuple, list)): outputs = outputs[0]
                for emb, mask in zip(outputs, attention_mask):
                    if pad:
                        embeddings.append(emb.cpu().float().numpy())
                    else:
                        emb_nonpad = emb[mask.bool()]
                        embeddings.append(emb_nonpad.cpu().float().numpy())
        return embeddings

    def embed_quotes(self, images):
        if isinstance(images, Image.Image):
            images = [images]
        embeddings = []
        dataloader = DataLoader(images, batch_size=self.bs, shuffle=False,
                                collate_fn=lambda x: self.process_images(x))
        with torch.no_grad():
            for batch in tqdm(dataloader, desc="[ColPaliRetriever] Embedding images"):
                batch = {k: v.to(self.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                for emb in torch.unbind(outputs):
                    embeddings.append(emb.cpu().float().numpy())
        return embeddings
     

    def process_queries(self, queries, max_length=512):
        texts_query = [f"Question: {q}" + "<pad>" * 10 for q in queries]
        sl = getattr(self.processor, "image_seq_length", 32) # 1024
        batch_query = self.processor(
            images=[self.mock_image] * len(texts_query),
            text=texts_query,
            return_tensors="pt",
            padding="longest",
            max_length=max_length + sl   # fallback seq len
        )
        if "pixel_values" in batch_query: del batch_query["pixel_values"]
        
        batch_query["input_ids"] = batch_query["input_ids"][..., sl :]
        batch_query["attention_mask"] = batch_query["attention_mask"][..., sl :]
        return batch_query

    def process_images(self, images): 
        pil_images = []
        for img in images:
            if isinstance(img, Image.Image):  # Already a PIL Image
                pil_img = img
            elif isinstance(img, (bytes, bytearray)):  # Binary image (e.g., from buffered.getvalue())
                pil_img = Image.open(io.BytesIO(img))
            else:
                raise ValueError("Each image must be a PIL.Image.Image or bytes.")
            pil_images.append(pil_img.convert("RGB"))
        
        texts = ["Describe the image."] * len(pil_images)
        batch_docs = self.processor(
            text=texts,
            images=pil_images,
            return_tensors="pt",
            padding="longest"
        )
        return batch_docs

    def score(self, query_embs, image_embs):
        """
        Computes (batch) similarity scores MaxSim style.
        Inputs:
           query_embs: [Nq, seq, dim]
           image_embs: [Ni, seq, dim]
        Returns:
           scores: [Nq, Ni] max similarity per query-image (like ColBERT)
        """
        qs = [torch.from_numpy(e) for e in query_embs]
        ds = [torch.from_numpy(e) for e in image_embs]
        # MaxSim/colbert scoring: max dot product over sequence dimension
        # shape: [Q, D]
        scores = np.zeros((len(qs), len(ds)), dtype=np.float32)
        for i, q in enumerate(qs):
            q = q.float()  # [Lq, d]
            for j, d in enumerate(ds):
                d = d.float() # [Ld, d]
                # score = max_{q_token, d_token} q_token @ d_token.T
                sim = torch.matmul(q, d.T)    # [Lq, Ld]
                maxsim = torch.max(sim, dim=1)[0].sum().item()  # colbert-style batch: sum-of-max over query tokens
                scores[i, j] = maxsim
        return scores



class ColQwen2(Qwen2VLForConditionalGeneration):
    """
    ColQwen2 model implementation.
    """
    def __init__(self, config: Qwen2VLConfig):
        super().__init__(config)
        self.dim = 128
        self.custom_text_proj = torch.nn.Linear(self.model.config.hidden_size, self.dim)
        self.padding_side = "left"
        self.post_init()

    def forward(self, *args, **kwargs) -> torch.Tensor:
        kwargs.pop("output_hidden_states", None)
        # scatter hack for DDP, see original code if needed
        if "pixel_values" in kwargs and "image_grid_thw" in kwargs:
            offsets = kwargs["image_grid_thw"][:, 1] * kwargs["image_grid_thw"][:, 2]
            kwargs["pixel_values"] = torch.cat([pv[:o] for pv, o in zip(kwargs["pixel_values"], offsets)], dim=0)

        position_ids, rope_deltas = self.get_rope_index(
            input_ids=kwargs["input_ids"],
            image_grid_thw=kwargs.get("image_grid_thw", None),
            video_grid_thw=None,
            attention_mask=kwargs.get("attention_mask", None),
        )
        outputs = super().forward(*args,
                                  **kwargs,
                                  position_ids=position_ids,
                                  rope_deltas=rope_deltas,
                                  use_cache=False,
                                  output_hidden_states=True)
        last_hidden_states = outputs.hidden_states[-1]
        proj = self.custom_text_proj(last_hidden_states)
        proj = proj / proj.norm(dim=-1, keepdim=True)
        proj = proj * kwargs["attention_mask"].unsqueeze(-1)
        return proj


class ColQwen2Retriever:
    def __init__(self, bs=4, use_gpu=True):
        self.bs = bs
        self.bs_query = 64
        self.model_name = "checkpoint/colqwen2-v1.0"
        self.base_ckpt = "checkpoint/colqwen2-base"
        self.device = "cuda" if torch.cuda.is_available() and use_gpu else "cpu"

        self.model = ColQwen2.from_pretrained(
            self.base_ckpt,
            torch_dtype=torch.bfloat16,
            device_map=self.device
        )
        self.model.load_adapter(self.model_name)
        self.model.eval()

        # DataParallel:
        self.is_parallel = False
        if torch.cuda.device_count() > 1:
            print(f"Using {torch.cuda.device_count()} GPUs with DataParallel")
            self.model = torch.nn.DataParallel(self.model)
            self.is_parallel = True

        self.processor = AutoProcessor.from_pretrained(self.model_name)
        self.min_pixels = 4 * 28 * 28
        self.max_pixels = 768 * 28 * 28
        self.factor = 28
        self.max_ratio = 200

    # ---------- Image Processing Utilities ----------
    @staticmethod
    def round_by_factor(number, factor):
        return round(number / factor) * factor

    @staticmethod
    def ceil_by_factor(number, factor):
        return math.ceil(number / factor) * factor

    @staticmethod
    def floor_by_factor(number, factor):
        return math.floor(number / factor) * factor

    def smart_resize(self, height: int, width: int) -> tuple:
        if max(height, width) / min(height, width) > self.max_ratio:
            raise ValueError(
                f"absolute aspect ratio must be smaller than {self.max_ratio}, "
                f"got {max(height, width) / min(height, width)}"
            )
        h_bar = max(self.factor, self.round_by_factor(height, self.factor))
        w_bar = max(self.factor, self.round_by_factor(width, self.factor))
        if h_bar * w_bar > self.max_pixels:
            beta = math.sqrt((height * width) / self.max_pixels)
            h_bar = self.floor_by_factor(height / beta, self.factor)
            w_bar = self.floor_by_factor(width / beta, self.factor)
        elif h_bar * w_bar < self.min_pixels:
            beta = math.sqrt(self.min_pixels / (height * width))
            h_bar = self.ceil_by_factor(height * beta, self.factor)
            w_bar = self.ceil_by_factor(width * beta, self.factor)
        return h_bar, w_bar

    def process_images(self, images):
        pil_images = []
        for img in images:
            if isinstance(img, Image.Image):
                pil_img = img
            elif isinstance(img, (bytes, bytearray)):
                pil_img = Image.open(io.BytesIO(img))
            else:
                raise ValueError("Each image must be a PIL.Image.Image or bytes.")
            pil_images.append(pil_img.convert("RGB"))
        
        # Resize and convert
        resized_images = []
        for image in pil_images:
            orig_size = image.size
            resized_height, resized_width = self.smart_resize(orig_size[1], orig_size[0])
            out_img = image.resize((resized_width,resized_height)).convert('RGB')
            resized_images.append(out_img)

        texts_doc = [ 
            "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n"
        ] * len(resized_images)

        batch_doc = self.processor(
            text=texts_doc,
            images=resized_images,
            padding="longest",
            return_tensors="pt"
        )
        # The following hack can be skipped during inference unless you run into shape mismatch
        offsets = batch_doc["image_grid_thw"][:, 1] * batch_doc["image_grid_thw"][:, 2]
        pixel_values = torch.split(batch_doc["pixel_values"], offsets.tolist())
        max_length = max([len(pv) for pv in pixel_values])
        pixel_values = [torch.cat([pv,
                                   torch.zeros((max_length - len(pv), pv.shape[1]),
                                               dtype=pv.dtype, device=pv.device)]) for pv in pixel_values]
        batch_doc["pixel_values"] = torch.stack(pixel_values)
        return batch_doc

    def process_queries(self, queries, max_length=50, suffix=None):
        if suffix is None:
            suffix = "<pad>" * 10
        texts_query = []
        for q in queries:
            q_ = f"Query: {q}{suffix}"
            texts_query.append(q_)
        batch_query = self.processor(
            text=texts_query,
            return_tensors="pt",
            padding="longest",
        )
        return batch_query

    def embed_queries(self, queries, pad=False):
        if isinstance(queries, str):
            queries = [queries]
        embeddings = []
        dataloader = DataLoader(
            queries, batch_size=self.bs_query, shuffle=False,
            collate_fn=lambda x: self.process_queries(x)
        )
        with torch.no_grad():
            # Use main device for DataParallel
            dev = self.model.device_ids[0] if self.is_parallel else self.model.device
            for batch in tqdm(dataloader, desc="[ColQwen2Retriever] Embedding queries"):
                batch = {k: v.to(dev) for k, v in batch.items()}
                outputs = self.model(**batch)
                attention_mask = batch["attention_mask"]
                if isinstance(outputs, (tuple, list)):
                    outputs = outputs[0]
                for emb, mask in zip(outputs, attention_mask):
                    if pad:
                        embeddings.append(emb.cpu().float().numpy())
                    else:
                        emb_nonpad = emb[mask.bool()]
                        embeddings.append(emb_nonpad.cpu().float().numpy())
        return embeddings

    def embed_quotes(self, images):
        if isinstance(images, Image.Image):
            images = [images]
        embeddings = []
        dataloader = DataLoader(
            images, batch_size=self.bs, shuffle=False,
            collate_fn=lambda x: self.process_images(x)
        )
        with torch.no_grad():
            dev = self.model.device_ids[0] if self.is_parallel else self.model.device
            for batch in tqdm(dataloader, desc="[ColQwen2Retriever] Embedding images"):
                batch = {k: v.to(dev) for k, v in batch.items()}
                outputs = self.model(**batch)
                for emb in torch.unbind(outputs):
                    embeddings.append(emb.cpu().float().numpy())
        return embeddings
    

    def score(self, query_embs, image_embs):
        qs = [torch.from_numpy(e) for e in query_embs]
        ds = [torch.from_numpy(e) for e in image_embs]
        scores = np.zeros((len(qs), len(ds)), dtype=np.float32)
        for i, q in enumerate(qs):
            q = q.float()
            for j, d in enumerate(ds):
                d = d.float()
                sim = torch.matmul(q, d.T)
                maxsim = torch.max(sim, dim=1)[0].sum().item()
                scores[i, j] = maxsim
        return scores

	import math
	import torch
	from torch import nn
	from torch.utils.data import DataLoader
	from tqdm import tqdm
	from PIL import Image
	import numpy as np
	import io

	from transformers import AutoProcessor
	from transformers.models.paligemma.modeling_paligemma import (
	PaliGemmaConfig,
	PaliGemmaForConditionalGeneration,
	PaliGemmaPreTrainedModel,
	)

	from transformers.models.qwen2_vl import (
	Qwen2VLForConditionalGeneration,
	Qwen2VLConfig,
	)

	class ColPali(PaliGemmaPreTrainedModel):
	"""
	ColPali model implementation from the "ColPali: Efficient Document Retrieval with Vision Language Models" paper.
	"""
	def __init__(self, config: PaliGemmaConfig):
	super().__init__(config=config)
	model = PaliGemmaForConditionalGeneration(config=config)
	if model.language_model._tied_weights_keys is not None:
	self._tied_weights_keys = [f"model.language_model.{k}" for k in model.language_model._tied_weights_keys]
	self.model = model
	self.dim = 128
	self.custom_text_proj = nn.Linear(self.model.config.text_config.hidden_size, self.dim)
	self.post_init()

	def forward(self, args, *kwargs) -> torch.Tensor:
	# Delete output_hidden_states from kwargs
	kwargs.pop("output_hidden_states", None)
	outputs = self.model(args, output_hidden_states=True, *kwargs) # (batch_size, sequence_length, hidden_size)
	last_hidden_states = outputs.hidden_states[-1] # (batch_size, sequence_length, hidden_size)
	proj = self.custom_text_proj(last_hidden_states) # (batch_size, sequence_length, dim)
	# L2 normalization
	proj = proj / proj.norm(dim=-1, keepdim=True) # (batch_size, sequence_length, dim)
	proj = proj * kwargs["attention_mask"].unsqueeze(-1) # (batch_size, sequence_length, dim)
	return proj


	class ColPaliRetriever():
	def __init__(self, bs=4, use_gpu=True):
	self.bs = bs
	self.bs_query = 64
	self.model_name = "checkpoint/colpali-v1.1"
	self.base_ckpt = "checkpoint/colpaligemma-3b-mix-448-base"
	# Load model on cuda:0 by default!
	device = "cuda:0" if (torch.cuda.is_available() and use_gpu) else "cpu"
	self.model = ColPali.from_pretrained(
	self.base_ckpt, torch_dtype=torch.bfloat16, device_map=None # <-- NONE: Don't use device_map
	)
	self.model.load_adapter(self.model_name)
	self.model = self.model.to(device)
	self.model.eval()
	# Multi-GPU with DataParallel
	if torch.cuda.device_count() > 1 and use_gpu:
	print(f"[ColPaliRetriever] Using DataParallel on {torch.cuda.device_count()} GPUs")
	self.model = torch.nn.DataParallel(self.model)
	self.device = torch.device("cuda:0")
	else:
	self.device = torch.device(device)
	print(f"[ColPaliRetriever - init] ColPali loaded from '{self.base_ckpt}' (Adapter '{self.model_name}')...")
	# Load processor
	self.processor = AutoProcessor.from_pretrained(self.model_name)
	self.mock_image = Image.new("RGB", (16, 16), color="black")


	def embed_queries(self, queries, pad=False):
	if isinstance(queries, str):
	queries = [queries]
	embeddings = []
	dataloader = DataLoader(queries, batch_size=self.bs_query, shuffle=False,
	collate_fn=lambda x: self.process_queries(x))
	with torch.no_grad():
	for batch in tqdm(dataloader, desc="[ColPaliRetriever] Embedding queries"):
	batch = {k: v.to(self.device) for k, v in batch.items()}
	outputs = self.model(**batch)
	attention_mask = batch["attention_mask"]
	if isinstance(outputs, (tuple, list)): outputs = outputs[0]
	for emb, mask in zip(outputs, attention_mask):
	if pad:
	embeddings.append(emb.cpu().float().numpy())
	else:
	emb_nonpad = emb[mask.bool()]
	embeddings.append(emb_nonpad.cpu().float().numpy())
	return embeddings

	def embed_quotes(self, images):
	if isinstance(images, Image.Image):
	images = [images]
	embeddings = []
	dataloader = DataLoader(images, batch_size=self.bs, shuffle=False,
	collate_fn=lambda x: self.process_images(x))
	with torch.no_grad():
	for batch in tqdm(dataloader, desc="[ColPaliRetriever] Embedding images"):
	batch = {k: v.to(self.device) for k, v in batch.items()}
	outputs = self.model(**batch)
	for emb in torch.unbind(outputs):
	embeddings.append(emb.cpu().float().numpy())
	return embeddings


	def process_queries(self, queries, max_length=512):
	texts_query = [f"Question: {q}" + "<pad>" * 10 for q in queries]
	sl = getattr(self.processor, "image_seq_length", 32) # 1024
	batch_query = self.processor(
	images=[self.mock_image] * len(texts_query),
	text=texts_query,
	return_tensors="pt",
	padding="longest",
	max_length=max_length + sl # fallback seq len
	)
	if "pixel_values" in batch_query: del batch_query["pixel_values"]

	batch_query["input_ids"] = batch_query["input_ids"][..., sl :]
	batch_query["attention_mask"] = batch_query["attention_mask"][..., sl :]
	return batch_query

	def process_images(self, images):
	pil_images = []
	for img in images:
	if isinstance(img, Image.Image): # Already a PIL Image
	pil_img = img
	elif isinstance(img, (bytes, bytearray)): # Binary image (e.g., from buffered.getvalue())
	pil_img = Image.open(io.BytesIO(img))
	else:
	raise ValueError("Each image must be a PIL.Image.Image or bytes.")
	pil_images.append(pil_img.convert("RGB"))

	texts = ["Describe the image."] * len(pil_images)
	batch_docs = self.processor(
	text=texts,
	images=pil_images,
	return_tensors="pt",
	padding="longest"
	)
	return batch_docs

	def score(self, query_embs, image_embs):
	"""
	Computes (batch) similarity scores MaxSim style.
	Inputs:
	query_embs: [Nq, seq, dim]
	image_embs: [Ni, seq, dim]
	Returns:
	scores: [Nq, Ni] max similarity per query-image (like ColBERT)
	"""
	qs = [torch.from_numpy(e) for e in query_embs]
	ds = [torch.from_numpy(e) for e in image_embs]
	# MaxSim/colbert scoring: max dot product over sequence dimension
	# shape: [Q, D]
	scores = np.zeros((len(qs), len(ds)), dtype=np.float32)
	for i, q in enumerate(qs):
	q = q.float() # [Lq, d]
	for j, d in enumerate(ds):
	d = d.float() # [Ld, d]
	# score = max_{q_token, d_token} q_token @ d_token.T
	sim = torch.matmul(q, d.T) # [Lq, Ld]
	maxsim = torch.max(sim, dim=1)[0].sum().item() # colbert-style batch: sum-of-max over query tokens
	scores[i, j] = maxsim
	return scores



	class ColQwen2(Qwen2VLForConditionalGeneration):
	"""
	ColQwen2 model implementation.
	"""
	def __init__(self, config: Qwen2VLConfig):
	super().__init__(config)
	self.dim = 128
	self.custom_text_proj = torch.nn.Linear(self.model.config.hidden_size, self.dim)
	self.padding_side = "left"
	self.post_init()

	def forward(self, args, *kwargs) -> torch.Tensor:
	kwargs.pop("output_hidden_states", None)
	# scatter hack for DDP, see original code if needed
	if "pixel_values" in kwargs and "image_grid_thw" in kwargs:
	offsets = kwargs["image_grid_thw"][:, 1] * kwargs["image_grid_thw"][:, 2]
	kwargs["pixel_values"] = torch.cat([pv[:o] for pv, o in zip(kwargs["pixel_values"], offsets)], dim=0)

	position_ids, rope_deltas = self.get_rope_index(
	input_ids=kwargs["input_ids"],
	image_grid_thw=kwargs.get("image_grid_thw", None),
	video_grid_thw=None,
	attention_mask=kwargs.get("attention_mask", None),
	)
	outputs = super().forward(*args,
	**kwargs,
	position_ids=position_ids,
	rope_deltas=rope_deltas,
	use_cache=False,
	output_hidden_states=True)
	last_hidden_states = outputs.hidden_states[-1]
	proj = self.custom_text_proj(last_hidden_states)
	proj = proj / proj.norm(dim=-1, keepdim=True)
	proj = proj * kwargs["attention_mask"].unsqueeze(-1)
	return proj


	class ColQwen2Retriever:
	def __init__(self, bs=4, use_gpu=True):
	self.bs = bs
	self.bs_query = 64
	self.model_name = "checkpoint/colqwen2-v1.0"
	self.base_ckpt = "checkpoint/colqwen2-base"
	self.device = "cuda" if torch.cuda.is_available() and use_gpu else "cpu"

	self.model = ColQwen2.from_pretrained(
	self.base_ckpt,
	torch_dtype=torch.bfloat16,
	device_map=self.device
	)
	self.model.load_adapter(self.model_name)
	self.model.eval()

	# DataParallel:
	self.is_parallel = False
	if torch.cuda.device_count() > 1:
	print(f"Using {torch.cuda.device_count()} GPUs with DataParallel")
	self.model = torch.nn.DataParallel(self.model)
	self.is_parallel = True

	self.processor = AutoProcessor.from_pretrained(self.model_name)
	self.min_pixels = 4 * 28 * 28
	self.max_pixels = 768 * 28 * 28
	self.factor = 28
	self.max_ratio = 200

	# ---------- Image Processing Utilities ----------
	@staticmethod
	def round_by_factor(number, factor):
	return round(number / factor) * factor

	@staticmethod
	def ceil_by_factor(number, factor):
	return math.ceil(number / factor) * factor

	@staticmethod
	def floor_by_factor(number, factor):
	return math.floor(number / factor) * factor

	def smart_resize(self, height: int, width: int) -> tuple:
	if max(height, width) / min(height, width) > self.max_ratio:
	raise ValueError(
	f"absolute aspect ratio must be smaller than {self.max_ratio}, "
	f"got {max(height, width) / min(height, width)}"
	)
	h_bar = max(self.factor, self.round_by_factor(height, self.factor))
	w_bar = max(self.factor, self.round_by_factor(width, self.factor))
	if h_bar * w_bar > self.max_pixels:
	beta = math.sqrt((height * width) / self.max_pixels)
	h_bar = self.floor_by_factor(height / beta, self.factor)
	w_bar = self.floor_by_factor(width / beta, self.factor)
	elif h_bar * w_bar < self.min_pixels:
	beta = math.sqrt(self.min_pixels / (height * width))
	h_bar = self.ceil_by_factor(height * beta, self.factor)
	w_bar = self.ceil_by_factor(width * beta, self.factor)
	return h_bar, w_bar

	def process_images(self, images):
	pil_images = []
	for img in images:
	if isinstance(img, Image.Image):
	pil_img = img
	elif isinstance(img, (bytes, bytearray)):
	pil_img = Image.open(io.BytesIO(img))
	else:
	raise ValueError("Each image must be a PIL.Image.Image or bytes.")
	pil_images.append(pil_img.convert("RGB"))

	# Resize and convert
	resized_images = []
	for image in pil_images:
	orig_size = image.size
	resized_height, resized_width = self.smart_resize(orig_size[1], orig_size[0])
	out_img = image.resize((resized_width,resized_height)).convert('RGB')
	resized_images.append(out_img)

	texts_doc = [
	"<\|im_start\|>user\n<\|vision_start\|><\|image_pad\|><\|vision_end\|>Describe the image.<\|im_end\|>\n"
	] * len(resized_images)

	batch_doc = self.processor(
	text=texts_doc,
	images=resized_images,
	padding="longest",
	return_tensors="pt"
	)
	# The following hack can be skipped during inference unless you run into shape mismatch
	offsets = batch_doc["image_grid_thw"][:, 1] * batch_doc["image_grid_thw"][:, 2]
	pixel_values = torch.split(batch_doc["pixel_values"], offsets.tolist())
	max_length = max([len(pv) for pv in pixel_values])
	pixel_values = [torch.cat([pv,
	torch.zeros((max_length - len(pv), pv.shape[1]),
	dtype=pv.dtype, device=pv.device)]) for pv in pixel_values]
	batch_doc["pixel_values"] = torch.stack(pixel_values)
	return batch_doc

	def process_queries(self, queries, max_length=50, suffix=None):
	if suffix is None:
	suffix = "<pad>" * 10
	texts_query = []
	for q in queries:
	q_ = f"Query: {q}{suffix}"
	texts_query.append(q_)
	batch_query = self.processor(
	text=texts_query,
	return_tensors="pt",
	padding="longest",
	)
	return batch_query

	def embed_queries(self, queries, pad=False):
	if isinstance(queries, str):
	queries = [queries]
	embeddings = []
	dataloader = DataLoader(
	queries, batch_size=self.bs_query, shuffle=False,
	collate_fn=lambda x: self.process_queries(x)
	)
	with torch.no_grad():
	# Use main device for DataParallel
	dev = self.model.device_ids[0] if self.is_parallel else self.model.device
	for batch in tqdm(dataloader, desc="[ColQwen2Retriever] Embedding queries"):
	batch = {k: v.to(dev) for k, v in batch.items()}
	outputs = self.model(**batch)
	attention_mask = batch["attention_mask"]
	if isinstance(outputs, (tuple, list)):
	outputs = outputs[0]
	for emb, mask in zip(outputs, attention_mask):
	if pad:
	embeddings.append(emb.cpu().float().numpy())
	else:
	emb_nonpad = emb[mask.bool()]
	embeddings.append(emb_nonpad.cpu().float().numpy())
	return embeddings

	def embed_quotes(self, images):
	if isinstance(images, Image.Image):
	images = [images]
	embeddings = []
	dataloader = DataLoader(
	images, batch_size=self.bs, shuffle=False,
	collate_fn=lambda x: self.process_images(x)
	)
	with torch.no_grad():
	dev = self.model.device_ids[0] if self.is_parallel else self.model.device
	for batch in tqdm(dataloader, desc="[ColQwen2Retriever] Embedding images"):
	batch = {k: v.to(dev) for k, v in batch.items()}
	outputs = self.model(**batch)
	for emb in torch.unbind(outputs):
	embeddings.append(emb.cpu().float().numpy())
	return embeddings


	def score(self, query_embs, image_embs):
	qs = [torch.from_numpy(e) for e in query_embs]
	ds = [torch.from_numpy(e) for e in image_embs]
	scores = np.zeros((len(qs), len(ds)), dtype=np.float32)
	for i, q in enumerate(qs):
	q = q.float()
	for j, d in enumerate(ds):
	d = d.float()
	sim = torch.matmul(q, d.T)
	maxsim = torch.max(sim, dim=1)[0].sum().item()
	scores[i, j] = maxsim
	return scores