Spaces:
Sleeping
Sleeping
File size: 2,412 Bytes
868d6df 1aab979 868d6df 1aab979 868d6df 69c116f 868d6df 1aab979 868d6df 1aab979 868d6df b80b731 868d6df | 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 | import os
from pathlib import Path
from PIL import Image
MODEL_ID = "akhilaaa3/Jev-Omni"
BASE_ID = "google/gemma-4-12B-it"
_classifier = None
# Pre-cache weights on disk at container launch (outside @spaces.GPU)
try:
from huggingface_hub import snapshot_download
print(f"[Engine] Pre-caching model weights to local disk: {MODEL_ID}...")
snapshot_download(MODEL_ID)
print(f"[Engine] Pre-caching base weights: {BASE_ID}...")
snapshot_download(BASE_ID)
print("[Engine] Model weights successfully cached to disk.")
except Exception as e:
print(f"[Engine] Disk pre-caching notice: {e}")
def get_classifier():
global _classifier
if _classifier is not None:
return _classifier
import torch
if torch.cuda.is_available():
print(f"[Engine] CUDA detected. Assembling {MODEL_ID} into GPU memory...")
from huggingface_hub import snapshot_download
import importlib.util
merged = snapshot_download(MODEL_ID)
spec = importlib.util.spec_from_file_location("jev_omni_loader", Path(merged) / "jev_omni.py")
loader_mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(loader_mod)
_classifier = loader_mod.load_jev_omni(device="cuda")
print("[Engine] Jev-Omni loaded successfully into CUDA memory.")
else:
print("[Engine] Running without CUDA. Using MockJevOmni.")
class MockJevOmni:
def predict(self, *, state, question, options, media=None, modality="image", video_frames=16):
p = 0.85 if "spill" in question.lower() or "wet" in question.lower() else 0.15
if len(options) == 2 and options[0] == "Yes" and options[1] == "No":
return {
"prediction": "Yes" if p >= 0.5 else "No",
"confidence": p if p >= 0.5 else 1.0 - p,
"probabilities": {"Yes": p, "No": round(1.0 - p, 4)}
}
else:
probs = [round(1.0 / len(options), 4) for _ in options]
probs[0] = round(1.0 - sum(probs[1:]), 4)
return {
"prediction": options[0],
"confidence": probs[0],
"probabilities": dict(zip(options, probs))
}
_classifier = MockJevOmni()
return _classifier
|