import torch
import torchreid

archivo_salida = "osnet_x1_0_msmt17_batch1.onnx"

print("1. Descargando los pesos originales de la Bestia (OSNet x1_0)...")
# ⚡ Al poner pretrained=True, Python usa el enlace directo interno de la librería
# saltándose los bloqueos del navegador web.
model = torchreid.models.build_model(
    name='osnet_x1_0',
    num_classes=1000,
    loss='softmax',
    pretrained=True 
)
model.eval()

print("2. Compilando el modelo ONNX blindado a Batch=1...")
dummy_input = torch.randn(1, 3, 256, 128)

torch.onnx.export(
    model, 
    dummy_input, 
    archivo_salida,
    export_params=True,
    opset_version=18,  # Mantenemos la versión 18 que ya nos funcionó perfecto
    do_constant_folding=True,
    input_names=['images'],
    output_names=['features']
)

print(f"¡Éxito total! Tu modelo de grado comercial máximo está listo: {archivo_salida}")