import os
from models import FromZero, UNetSR
from safetensors.torch import save_model
from utils import cargar_etiquetas


def main():
    # Crear la carpeta model_test si no existe
    os.makedirs("model_test", exist_ok=True)
    _, num_classes, _ = cargar_etiquetas()

    from_zero_model = FromZero(num_classes=num_classes)
    unet_sr_model = UNetSR()

    save_model(from_zero_model, "model_test/from_zero_model.safetensor")
    save_model(unet_sr_model, "model_test/unet_sr_model.safetensor")

    print(
        "Los modelos de clasificación y super-resolución fueron creados en la carpeta 'model_test'"
    )


if __name__ == "__main__":
    main()