Spaces:
Paused
Paused
File size: 5,156 Bytes
37b0e8b | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import time
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
from pathlib import Path
from typing import Optional, Dict, Any
import sys
sys.path.insert(0, str(Path(__file__).parent.parent / "model" / "src"))
from model import MFFTWithExplainability, build_mfft
class ModelServer:
"""
Wraps the MFFT model for production inference.
Handles model loading, preprocessing, prediction, and warmup.
"""
def __init__(self, checkpoint_path: Optional[str] = None, variant: str = "base",
image_size: int = 384):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.variant = variant
self.image_size = image_size
print(f"[ModelServer] Device: {self.device} | Variant: {variant} | Size: {image_size}")
self.model = build_mfft(variant)
self.model = self.model.to(self.device)
self.is_loaded = False
if checkpoint_path and Path(checkpoint_path).exists():
self.load_checkpoint(checkpoint_path)
self.is_loaded = True
print(f"[ModelServer] Model loaded from {checkpoint_path}")
else:
print(f"[ModelServer] No checkpoint found at {checkpoint_path}")
print(f"[ModelServer] Running with untrained weights (random initialization)")
self.warmup()
def load_checkpoint(self, path: str):
try:
state = torch.load(path, map_location=self.device, weights_only=True)
except Exception:
state = torch.load(path, map_location=self.device, weights_only=False)
raw_state = state["model_state_dict"] if isinstance(state, dict) and "model_state_dict" in state else state
# Remap keys between cluster training pipeline and local model.py naming conventions
cleaned_state = {}
for k, v in raw_state.items():
nk = k.replace(".stages.", ".blocks.").replace("fga.gate.", "freq_guided_attn.freq_gate.")
if "num_batches_tracked" in nk:
continue
cleaned_state[nk] = v
self.model.load_state_dict(cleaned_state, strict=False)
self.is_loaded = True
def warmup(self, n_warmup: int = 3):
dummy = torch.randn(1, 3, self.image_size, self.image_size).to(self.device)
self.model.eval()
with torch.no_grad():
for _ in range(n_warmup):
_ = self.model(dummy)
print(f"[ModelServer] Warmup complete ({n_warmup} iterations)")
@torch.no_grad()
def predict(self, image: Image.Image) -> Dict[str, Any]:
start = time.time()
input_tensor = self._preprocess(image)
logits, heatmaps = self.model(input_tensor, return_heatmap=True)
probs = F.softmax(logits, dim=-1)
pred = torch.argmax(probs, dim=-1).item()
bands = self.model.decomposer(input_tensor)
freq_magnitudes = [band.abs().mean().item() for band in bands]
if len(freq_magnitudes) == 4:
band_contribs = {
"low_frequency": freq_magnitudes[0],
"low_mid_frequency": freq_magnitudes[1],
"mid_high_frequency": freq_magnitudes[2],
"high_frequency": freq_magnitudes[3],
}
else:
band_contribs = {
"low_frequency": freq_magnitudes[0],
"mid_frequency": freq_magnitudes[1] if len(freq_magnitudes) > 1 else 0.0,
"high_frequency": freq_magnitudes[2] if len(freq_magnitudes) > 2 else 0.0,
}
elapsed = (time.time() - start) * 1000
result = {
"prediction": pred,
"real_prob": probs[0, 0].item(),
"ai_prob": probs[0, 1].item(),
"confidence": probs.max(dim=-1).values.item(),
"heatmaps": heatmaps,
"frequency_band_contributions": band_contribs,
"processing_time_ms": elapsed,
}
return result
def _preprocess(self, image: Image.Image) -> torch.Tensor:
# Aspect-ratio preserving resize + center crop (matching training/eval pipeline)
w, h = image.size
target = self.image_size + 16
if w < h:
new_w, new_h = target, max(target, int(h * target / w))
else:
new_h, new_w = target, max(target, int(w * target / h))
image = image.convert("RGB").resize((new_w, new_h), Image.Resampling.BILINEAR)
left = max(0, (new_w - self.image_size) // 2)
top = max(0, (new_h - self.image_size) // 2)
image = image.crop((left, top, left + self.image_size, top + self.image_size))
img_array = np.array(image, dtype=np.float32) / 255.0
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
img_array = ((img_array - mean) / std).astype(np.float32)
tensor = torch.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0)
return tensor.to(self.device)
|