MiniMax-AI/VTP-Base-f16d64 on Ascend NPU

pip install -U atomgit
pip install torch torchvision transformers timm omegaconf
# torch-npu 与 CANN 驱动需预先安装

from atomgit_hub import snapshot_download
snapshot_download("MiniMax-AI/VTP-Base-f16d64", local_dir="./VTP-Base-f16d64")

git clone https://github.com/MiniMax-AI/VTP.git
export PYTHONPATH="${PYTHONPATH}:$(pwd)/VTP"

import torch
import torch_npu
from PIL import Image
from torchvision import transforms

from vtp.models.vtp_hf import VTPConfig, VTPModel
from vtp.tokenizers import get_tokenizer

MODEL_PATH = "./VTP-Base-f16d64"
DEVICE = "npu:0"

model = VTPModel.from_pretrained(MODEL_PATH).to(DEVICE).eval()

preprocess = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image = preprocess(Image.open("dog.png")).unsqueeze(0).to(DEVICE)

# ---------- 1. 图像重建 ----------
with torch.no_grad():
    latents = model.get_reconstruction_latents(image)
    recon = model.get_latents_decoded_images(latents)
print("Latents shape:", latents.shape)  # [1, 64, 16, 16]

# ---------- 2. CLIP 零样本分类 ----------
tokenizer = get_tokenizer('ViT-B-32', context_length=model.config.text_context_length)
text = tokenizer(["a diagram", "a dog", "a cat", "a person"]).to(DEVICE)
with torch.no_grad():
    image_features = model.get_clip_image_feature(image, normalize=True)
    text_features = model.get_clip_text_feature(text, normalize=True)
    text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
print("Label probs:", [f"{p:.4f}" for p in text_probs[0].tolist()])

# ---------- 3. SSL 特征提取 ----------
with torch.no_grad():
    features = model.get_last_layer_feature(image)
    cls_token = features['cls_token']       # [1, 768]
    patch_tokens = features['patch_tokens'] # [1, 256, 768]
print("CLS shape:", cls_token.shape)