| import torch |
| from safetensors.torch import load_file |
| from architecture import ValdamVision2_5 |
| import json |
|
|
| class ValdamVisionLocalPredictor: |
| def __init__(self, safetensors_weights_path="model.safetensors", config_path="config.json", vocab_path="vocab.json", device="cpu"): |
| self.device = torch.device(device if torch.cuda.is_available() else "cpu") |
| |
| |
| with open(config_path, "r", encoding="utf-8") as f: |
| self.config = json.load(f) |
| with open(vocab_path, "r", encoding="utf-8") as f: |
| self.vocab_stoi = json.load(f) |
| self.vocab_itos = {int(v): k for k, v in self.vocab_stoi.items() if k is not None} |
|
|
| |
| arch = self.config["architecture"] |
| self.model = ValdamVision2_5( |
| embed_size=arch["embed_size"], |
| vocab_size=len(self.vocab_stoi), |
| num_heads=arch.get("num_heads", 8), |
| num_layers=arch.get("num_layers", 6) |
| ).to(self.device) |
|
|
| |
| weights = load_file(safetensors_weights_path, device=device) |
| self.model.load_state_dict(weights) |
| self.model.eval() |
| print("🚀 Valdam Vision 2.5 cargado de forma ultra-segura con Safetensors.") |
|
|
| def predict(self, image_tensor, max_len=30): |
| |
| with torch.no_grad(): |
| memory = self.model.encoder(image_tensor.to(self.device)) |
| input_tokens = torch.tensor([[self.vocab_stoi["<SOS>"]]], device=self.device) |
|
|
| result = [] |
| for _ in range(max_len): |
| outputs = self.model.decoder(memory, input_tokens) |
| next_token = outputs[0, -1, :].argmax().item() |
| token_str = self.vocab_itos.get(next_token, "<UNK>") |
| |
| if token_str == "<EOS>": |
| break |
| if token_str not in ["<SOS>", "<PAD>"]: |
| result.append(token_str) |
| |
| input_tokens = torch.cat([input_tokens, torch.tensor([[next_token]], device=self.device)], dim=1) |
| |
| return " ".join(result) |
|
|
