fn_registry/python/functions/ml/model_ref.py

"""ModelRef — referencia a un modelo de generacion de imagenes."""

from __future__ import annotations

from typing import Literal

try:
    from pydantic import BaseModel, ConfigDict

    class ModelRef(BaseModel):
        """Referencia a un modelo de generacion de imagenes.

        Identifica el modelo por nombre (HuggingFace hub o ruta local),
        tipo de arquitectura y cuantizacion. Serializable a JSON canonico
        con model_dump() / model_dump_json() para el contrato compartido con Go.

        Attributes:
            name:           Nombre del modelo en HuggingFace Hub o identificador local.
                            Ejemplo: "stabilityai/stable-diffusion-xl-base-1.0".
            model_type:     Arquitectura del modelo. Uno de los literales definidos.
            quantization:   Precision numerica del checkpoint. Por defecto "fp16".
            path:           Ruta local al checkpoint si ya fue descargado. None si
                            se debe descargar del hub.
        """

        model_config = ConfigDict(frozen=True)

        name: str
        model_type: Literal[
            "sd15",
            "sd20",
            "sdxl",
            "sd3",
            "flux_dev",
            "flux_schnell",
            "flux_kontext",
            "qwen_image",
            "chroma",
            "z_image",
        ]
        quantization: Literal[
            "fp32", "fp16", "bf16", "q8_0", "q5_1", "q5_0", "q4_1", "q4_0"
        ] = "fp16"
        path: str | None = None

except ImportError:
    from dataclasses import dataclass

    @dataclass(frozen=True)
    class ModelRef:  # type: ignore[no-redef]
        """Referencia a un modelo de generacion de imagenes (fallback dataclass).

        Usar la version pydantic cuando este disponible para validacion y
        serializacion JSON canonica. Esta version no valida los literales en
        tiempo de ejecucion.

        Attributes:
            name:         Nombre del modelo en HuggingFace Hub o ruta local.
            model_type:   Arquitectura del modelo (sd15|sd20|sdxl|sd3|flux_dev|...).
            quantization: Precision numerica (fp32|fp16|bf16|q8_0|...). Por defecto "fp16".
            path:         Ruta local al checkpoint. None si no esta descargado.
        """

        name: str
        model_type: str
        quantization: str = "fp16"
        path: str | None = None