| """Extract independent-image MiniMax features from this downloaded bundle.""" |
| import argparse |
| from pathlib import Path |
| import sys |
|
|
| from PIL import Image |
| import torch |
| from torchvision import transforms |
| from transformers import MiniMaxM3VLVisionModel |
|
|
| BUNDLE = Path(__file__).absolute().parents[1] |
| sys.path.insert(0, str(BUNDLE)) |
| from projector import load_projector |
|
|
|
|
| def preprocess(image: Image.Image) -> torch.Tensor: |
| pixels = transforms.Compose([ |
| transforms.Resize(448, interpolation=transforms.InterpolationMode.BICUBIC), |
| transforms.CenterCrop(448), |
| transforms.ToTensor(), |
| transforms.Normalize((0.48145466, 0.4578275, 0.40821073), |
| (0.26862954, 0.26130258, 0.27577711)), |
| ])(image) |
| frames = pixels.unsqueeze(0).expand(2, -1, -1, -1) |
| patches = frames.reshape(1, 2, 3, 16, 2, 14, 16, 2, 14) |
| return patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).reshape(1024, 1176) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("image", type=Path) |
| parser.add_argument("--device", choices=("cuda:0", "cpu"), default="cuda:0") |
| args = parser.parse_args() |
| torch.set_num_threads(2) |
| torch.set_num_interop_threads(2) |
| device, dtype = torch.device(args.device), torch.bfloat16 |
| with Image.open(args.image) as original: |
| original_size = original.size |
| pixels = preprocess(original.convert("RGB")).to(device=device, dtype=dtype) |
| model = MiniMaxM3VLVisionModel.from_pretrained( |
| BUNDLE, dtype=dtype, local_files_only=True, |
| ).to(device).eval() |
| projector = load_projector(BUNDLE).to(device=device, dtype=dtype).eval() |
| with torch.inference_mode(): |
| packed = model(pixels, torch.tensor([[1, 32, 32]], device=device)).last_hidden_state |
| packed = packed.reshape(1, 1024, 1280) |
| tower = packed.view(1, 16, 16, 2, 2, 1280).permute(0, 1, 3, 2, 4, 5).reshape(1, 1024, 1280) |
| merged = tower.view(1, 16, 2, 16, 2, 1280).permute(0, 1, 3, 2, 4, 5).reshape(1, 256, 5120) |
| projected = projector(packed.reshape(1024, 1280)).reshape(1, 256, 6144) |
| print("bundle", BUNDLE) |
| print("device", device, "dtype", dtype, "batch", 1, "image_size", original_size, "crop", (448, 448)) |
| for name, tokens, shape in ( |
| ("tower", tower, (1, 1024, 1280)), |
| ("merged", merged, (1, 256, 5120)), |
| ("projected", projected, (1, 256, 6144)), |
| ): |
| assert tuple(tokens.shape) == shape and tokens.dtype == torch.bfloat16 |
| assert torch.isfinite(tokens).all().item(), name |
| print(name, tuple(tokens.shape), "dtype", tokens.dtype, "device", tokens.device) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
