takes last model from repo

app.py

@@ -1,32 +1,85 @@
-import gradio as gr 
-from transformers import AutoModelForImageClassification, AutoImageProcessor 
-from PIL import Image 
-import torch 
-import os 
- 
-MODEL_ID = "nicolasrl/deepfake_vs_real_ViTlarge" 
-HF_TOKEN = os.getenv("HF_TOKEN")  
- 
-model = AutoModelForImageClassification.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN) 
-processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN) 
- 
-def predict(image: Image.Image):     
-    inputs = processor(images=image, return_tensors="pt")     
-    with torch.no_grad():         
-        outputs = model(**inputs)     
-    logits = outputs.logits     
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]     
-    results = {model.config.id2label[i]: float(probs[i]) for i in range(len(probs))}     
-    return dict(sorted(results.items(), key=lambda x: x[1], reverse=True)) 
- 
-iface = gr.Interface(     
-    fn=predict,     
-    inputs=gr.Image(type="pil"),     
-    outputs=gr.Label(num_top_classes=2),     
-    title="Deepfake o Real",     
-    description="Sube una imagen y el modelo hace magia para predecir si es Deepfake o Real",     
-    api_name="predict" 
-) 
- 
-if __name__ == "__main__":     
-    iface.launch(show_error=True)
+import os
+import glob
+import torch
+import onnxruntime as ort
+import numpy as np
+from PIL import Image
+import gradio as gr
+from huggingface_hub import snapshot_download
+
+# ----------------------------------------------------------------------
+# 1. Leer token de entorno
+# ----------------------------------------------------------------------
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise EnvironmentError("No se encontró la variable de entorno HF_TOKEN. Añádela en los Secrets del Space.")
+
+# ----------------------------------------------------------------------
+# 2. Descargar el repositorio autenticado
+# ----------------------------------------------------------------------
+repo_id = "nicolasrl/df-vs-real"
+local_dir = snapshot_download(
+    repo_id,
+    local_dir="repo",
+    repo_type="model",
+    use_auth_token=HF_TOKEN
+)
+
+# ----------------------------------------------------------------------
+# 3. Localizar el modelo ONNX más reciente
+# ----------------------------------------------------------------------
+onnx_paths = glob.glob(os.path.join(local_dir, "**/model.onnx"), recursive=True)
+if not onnx_paths:
+    raise FileNotFoundError("No se encontró ningún archivo model.onnx en el repositorio.")
+
+latest_model_path = max(onnx_paths, key=os.path.getmtime)
+print(f"Usando modelo: {latest_model_path}")
+
+# ----------------------------------------------------------------------
+# 4. Configurar sesión ONNX Runtime
+# ----------------------------------------------------------------------
+providers = ['CPUExecutionProvider']
+session = ort.InferenceSession(latest_model_path, providers=providers)
+
+# ----------------------------------------------------------------------
+# 5. Preprocesamiento de imágenes
+# ----------------------------------------------------------------------
+def preprocess(image: Image.Image):
+    image = image.convert("RGB").resize((224, 224))
+    img_array = np.array(image).astype(np.float32) / 255.0
+    img_array = img_array.transpose(2, 0, 1)  # CHW
+    img_array = np.expand_dims(img_array, 0)  # batch
+    return img_array
+
+# ----------------------------------------------------------------------
+# 6. Inferencia
+# ----------------------------------------------------------------------
+def predict(image: Image.Image):
+    input_data = preprocess(image)
+    input_name = session.get_inputs()[0].name
+    output_name = session.get_outputs()[0].name
+    output = session.run([output_name], {input_name: input_data})[0]
+
+    prob_fake = float(output[0][0])
+    label = "Deepfake" if prob_fake > 0.5 else "Real"
+    return {"Deepfake": prob_fake, "Real": 1 - prob_fake}, f"Predicción: {label}"
+
+# ----------------------------------------------------------------------
+# 7. Interfaz Gradio
+# ----------------------------------------------------------------------
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Sube una imagen"),
+    outputs=[
+        gr.Label(label="Probabilidades"),
+        gr.Textbox(label="Resultado")
+    ],
+    title="Detección de Deepfakes - ONNX",
+    description="Modelo ONNX cargado automáticamente desde la versión más reciente del repo."
+)
+
+# ----------------------------------------------------------------------
+# 8. Lanzar app
+# ----------------------------------------------------------------------
+if __name__ == "__main__":
+    demo.launch()