File size: 2,704 Bytes
b161cf1 75a5fc5 b161cf1 75a5fc5 b161cf1 75a5fc5 b161cf1 75a5fc5 b161cf1 75a5fc5 b161cf1 75a5fc5 b161cf1 75a5fc5 b161cf1 75a5fc5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """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()
|