from huggingface_hub import hf_hub_download
from transformers import ViTForImageClassification, AutoImageProcessor



from fastapi import FastAPI, File, UploadFile
from transformers import ViTForImageClassification, AutoImageProcessor
import io
from PIL import Image
import torch

model_file = hf_hub_download(
    repo_id="iwin10s/leak-detection-model",
    filename="model.safetensors"
)

processor_config = hf_hub_download(
    repo_id="iwin10s/leak-detection-model",
    filename="preprocessor_config.json"
)

# Load model & processor
processor = AutoImageProcessor.from_pretrained("iwin10s/leak-detection-model")
model = ViTForImageClassification.from_pretrained(
    pretrained_model_name_or_path="iwin10s/leak-detection-model",
    local_files_only=False
)

app = FastAPI()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    data = await file.read()

    img = Image.open(io.BytesIO(data)).convert("RGB")
    inputs = processor(images=img, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1).cpu().numpy()[0]

    return {"leak_probability": float(probs[1])}