fn_registry/python/functions/ml/image_gen_result.py

"""ImageGenResult — resultado de una operacion de generacion de imagen."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

# PIL.Image.Image se importa solo para type-checking estatico (mypy/pyright).
# En runtime NO se importa aqui — el consumidor ya tiene PIL instalado si
# trabaja con imagenes reales. Esto evita ImportError cuando el modulo se
# importa en contextos sin Pillow (ej. el orquestador Go via grpc/json).
if TYPE_CHECKING:
    from PIL.Image import Image as PILImage

try:
    from pydantic import BaseModel, ConfigDict, field_validator

    class ImageGenResult(BaseModel):
        """Resultado de una operacion de generacion de imagen con modelo de difusion.

        El campo `image` contiene el objeto PIL.Image.Image generado. No es
        serializable a JSON — se accede directamente para guardar a disco o
        pasar a pipelines de post-proceso. Para serializar el resultado,
        usar solo el campo `meta` (que incluye la config usada) y guardar
        la imagen por separado.

        Attributes:
            image:        Imagen generada. Tipo PIL.Image.Image en runtime.
                          No incluido en model_dump() ni model_dump_json().
                          None si la generacion fallo (ver meta["error"]).
            meta:         Diccionario con metadata de la generacion. Debe incluir:
                          - "config": GenerationConfig.model_dump() con los params usados.
                          - "model": nombre del modelo.
                          - "seed_used": semilla real usada (util cuando seed=-1).
                          - "sampler": nombre del sampler.
                          Puede incluir campos adicionales del backend.
            duration_ms:  Tiempo total de generacion en milisegundos.
            vram_peak_mb: Pico de VRAM consumida durante la generacion en MiB.
                          None si no se pudo medir (CPU inference o backend sin soporte).
        """

        model_config = ConfigDict(arbitrary_types_allowed=True)

        image: Any  # PIL.Image.Image en runtime; Any para evitar dep dura
        meta: dict[str, Any]
        duration_ms: int
        vram_peak_mb: int | None = None

        @field_validator("image", mode="before")
        @classmethod
        def _validate_image(cls, v: Any) -> Any:
            # Aceptar None (generacion fallida) o cualquier objeto imagen.
            # No forzamos importar PIL aqui — la validacion real la hace el backend.
            return v

        def model_dump(self, **kwargs: Any) -> dict[str, Any]:
            """Serializa a dict excluyendo el campo image (no serializable a JSON).

            Returns:
                dict con meta, duration_ms y vram_peak_mb. El campo image se omite.
            """
            return {
                "meta": self.meta,
                "duration_ms": self.duration_ms,
                "vram_peak_mb": self.vram_peak_mb,
            }

        def model_dump_json(self, **kwargs: Any) -> str:
            """Serializa a JSON excluyendo el campo image.

            Returns:
                String JSON con meta, duration_ms y vram_peak_mb.
            """
            import json

            return json.dumps(self.model_dump())

except ImportError:
    from dataclasses import dataclass

    @dataclass
    class ImageGenResult:  # type: ignore[no-redef]
        """Resultado de generacion de imagen (fallback dataclass).

        El campo `image` es PIL.Image.Image en runtime. No serializable a JSON.
        Usar `meta` + guardar imagen por separado para persistencia.

        Attributes:
            image:        PIL.Image.Image generada. None si fallo.
            meta:         Metadata: config usada, modelo, seed_used, sampler, etc.
            duration_ms:  Duracion total de generacion en milisegundos.
            vram_peak_mb: Pico de VRAM en MiB. None si no se pudo medir.
        """

        image: Any  # PIL.Image.Image
        meta: dict
        duration_ms: int
        vram_peak_mb: int | None = None