| from typing import Union |
| import torch |
| from PIL import Image |
| from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor |
| from torchvision.transforms import InterpolationMode |
|
|
| BICUBIC = InterpolationMode.BICUBIC |
| ImageInput = Union[Image.Image, list[Image.Image], tuple[Image.Image, ...]] |
|
|
| def _convert_image_to_rgb(image): |
| return image.convert("RGB") |
|
|
| def _image_transform(image_size): |
| return Compose( |
| [ |
| Resize(image_size, interpolation=BICUBIC), |
| CenterCrop(image_size), |
| _convert_image_to_rgb, |
| ToTensor(), |
| Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), |
| ] |
| ) |
|
|
| def _as_list(value): |
| if isinstance(value, (list, tuple)): |
| return list(value) |
| return [value] |
|
|
|
|
| class ImageRewardMLP(torch.nn.Module): |
| def __init__(self, input_size): |
| super().__init__() |
| self.layers = torch.nn.Sequential( |
| torch.nn.Linear(input_size, 1024), |
| torch.nn.Dropout(0.2), |
| torch.nn.Linear(1024, 128), |
| torch.nn.Dropout(0.2), |
| torch.nn.Linear(128, 64), |
| torch.nn.Dropout(0.1), |
| torch.nn.Linear(64, 16), |
| torch.nn.Linear(16, 1), |
| ) |
|
|
| for name, param in self.layers.named_parameters(): |
| if "weight" in name: |
| torch.nn.init.normal_(param, mean=0.0, std=1.0 / (input_size + 1)) |
| if "bias" in name: |
| torch.nn.init.constant_(param, val=0) |
|
|
| def forward(self, x): |
| return self.layers(x) |
|
|
|
|
| class ImageRewardModel(torch.nn.Module): |
| def __init__(self, blip=None, tokenizer=None, image_size=224, max_length=35, mean=0.16717362830052426, std=1.0333394966054072): |
| super().__init__() |
| if blip is None: |
| blip = self.default_blip_model() |
| |
| self.blip = blip |
| self.tokenizer = tokenizer |
| self.preprocess = _image_transform(image_size) |
| self.max_length = max_length |
| self.mlp = ImageRewardMLP(blip.config.text_config.hidden_size) |
| |
| self.register_buffer("score_mean", torch.tensor(float(mean)), persistent=False) |
| self.register_buffer("score_std", torch.tensor(float(std)), persistent=False) |
|
|
| @staticmethod |
| def default_blip_model(): |
| from transformers import BlipConfig, BlipForImageTextRetrieval |
|
|
| vision_hidden_size = 1024 |
| text_config = ImageRewardModel._load_text_config(None) |
| config = BlipConfig( |
| vision_config={ |
| "hidden_size": vision_hidden_size, |
| "intermediate_size": vision_hidden_size * 4, |
| "num_hidden_layers": 24, |
| "num_attention_heads": 16, |
| "image_size": 224, |
| "patch_size": 16, |
| "hidden_act": "gelu", |
| "layer_norm_eps": 1e-6, |
| }, |
| text_config={ |
| **text_config, |
| "vocab_size": 30524, |
| "encoder_hidden_size": vision_hidden_size, |
| "add_cross_attention": True, |
| "is_decoder": True, |
| }, |
| projection_dim=256, |
| ) |
| return BlipForImageTextRetrieval(config) |
|
|
| @staticmethod |
| def _load_text_config(med_config_path): |
| return { |
| "hidden_size": 768, |
| "intermediate_size": 3072, |
| "num_hidden_layers": 12, |
| "num_attention_heads": 12, |
| "max_position_embeddings": 512, |
| "vocab_size": 30524, |
| "hidden_act": "gelu", |
| "layer_norm_eps": 1e-12, |
| "attention_probs_dropout_prob": 0.1, |
| "hidden_dropout_prob": 0.1, |
| "pad_token_id": 0, |
| "type_vocab_size": 2, |
| } |
|
|
| @staticmethod |
| def convert_key_value(key, value): |
| if key.startswith("blip.visual_encoder."): |
| suffix = key[len("blip.visual_encoder.") :] |
| |
| if suffix == "cls_token": |
| return "blip.vision_model.embeddings.class_embedding", value |
| if suffix == "pos_embed": |
| return "blip.vision_model.embeddings.position_embedding", value |
| if suffix.startswith("patch_embed.proj."): |
| return "blip.vision_model.embeddings.patch_embedding." + suffix[len("patch_embed.proj.") :], value |
| |
| if suffix.startswith("blocks."): |
| parts = suffix.split(".") |
| layer = parts[1] |
| rest = ".".join(parts[2:]) |
| prefix = f"blip.vision_model.encoder.layers.{layer}." |
| mapping = { |
| "norm1.": "layer_norm1.", |
| "attn.qkv.": "self_attn.qkv.", |
| "attn.proj.": "self_attn.projection.", |
| "norm2.": "layer_norm2.", |
| "mlp.fc1.": "mlp.fc1.", |
| "mlp.fc2.": "mlp.fc2.", |
| } |
| for source, target in mapping.items(): |
| if rest.startswith(source): |
| return prefix + target + rest[len(source) :], value |
| |
| if suffix.startswith("norm."): |
| return "blip.vision_model.post_layernorm." + suffix[len("norm.") :], value |
| |
| return None, value |
| return key, value |
|
|
| @property |
| def device(self): |
| return next(self.parameters(), torch.tensor([])).device |
|
|
| @property |
| def dtype(self): |
| return next(self.parameters(), torch.tensor(0.0)).dtype |
|
|
| def _tokenize(self, prompts): |
| return self.tokenizer( |
| prompts, |
| padding="max_length", |
| truncation=True, |
| max_length=self.max_length, |
| return_tensors="pt", |
| ).to(self.device) |
|
|
| def _preprocess_images(self, images): |
| tensors = [self.preprocess(image.convert("RGB")) for image in images] |
| return torch.stack(tensors, dim=0).to(device=self.device, dtype=self.dtype) |
|
|
| def _normalize_inputs(self, prompts, images): |
| images = _as_list(images) |
| prompts = _as_list(prompts) |
| |
| if len(prompts) == 1 and len(images) > 1: |
| prompts = prompts * len(images) |
| if len(images) == 1 and len(prompts) > 1: |
| images = images * len(prompts) |
| |
| if len(prompts) != len(images): |
| raise ValueError(f"Expected the same number of prompts and images, got {len(prompts)} and {len(images)}.") |
| |
| return prompts, images |
|
|
| @torch.no_grad() |
| def forward(self, prompts: Union[str, list[str]], images): |
| prompts, images = self._normalize_inputs(prompts, images) |
| text_input = self._tokenize(prompts) |
| image_tensor = self._preprocess_images(images) |
| |
| image_output = self.blip.vision_model(pixel_values=image_tensor, return_dict=True) |
| image_embeds = image_output.last_hidden_state |
| image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=self.device) |
| |
| text_output = self.blip.text_encoder( |
| input_ids=text_input.input_ids, |
| attention_mask=text_input.attention_mask, |
| encoder_hidden_states=image_embeds, |
| encoder_attention_mask=image_atts, |
| return_dict=True, |
| ) |
| |
| text_features = text_output.last_hidden_state[:, 0, :] |
| rewards = self.mlp(text_features).squeeze(-1) |
| rewards = (rewards - self.score_mean) / self.score_std |
| |
| return rewards |
