"""
🔥 TEST YOUR AUDIO MODEL FROM HUGGING FACE
Simma7/audio_model
"""

import torch
import librosa
from transformers import AutoProcessor, AutoModel
from huggingface_hub import snapshot_download

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ================= DOWNLOAD MODEL =================
print("⬇️ Downloading model from Hugging Face...")

model_dir = snapshot_download("Simma7/audio_model")

print("✅ Download complete")

# ================= LOAD MODEL =================
print("🧠 Loading model...")

processor = AutoProcessor.from_pretrained(model_dir)
model = AutoModel.from_pretrained(model_dir).to(DEVICE)

print("✅ Model loaded")

# ================= PREDICT =================
def predict(audio_path):
    audio, sr = librosa.load(audio_path, sr=16000)

    inputs = processor(audio, sampling_rate=16000, return_tensors="pt", padding=True)

    with torch.no_grad():
        outputs = model(**inputs.to(DEVICE))

    # embedding-based score
    embedding = outputs.last_hidden_state.mean(dim=1)
    prob = torch.sigmoid(embedding.mean()).item()

    return prob

# ================= MAIN =================
if __name__ == "__main__":
    audio_path = "test.wav"   # 🔥 put your audio file

    print("\n🔍 Analyzing audio...")

    prob = predict(audio_path)

    if prob > 0.5:
        print("\n🔴 FAKE AUDIO")
    else:
        print("\n🟢 REAL AUDIO")

    print(f"📊 Confidence: {prob:.4f}")