fn_registry/python/functions/ml/generation_config.py

"""GenerationConfig — contrato de parametros para generacion de imagenes con difusion."""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

if TYPE_CHECKING:
    pass

_SAMPLER_VALUES = (
    "euler",
    "euler_a",
    "dpm++2m",
    "dpm++2m_v2",
    "heun",
    "dpm2",
    "lcm",
)

SamplerName = Literal[
    "euler",
    "euler_a",
    "dpm++2m",
    "dpm++2m_v2",
    "heun",
    "dpm2",
    "lcm",
]

try:
    from pydantic import BaseModel, ConfigDict

    from lora_ref import LoraRef
    from model_ref import ModelRef

    class GenerationConfig(BaseModel):
        """Contrato de parametros para generacion de imagenes con modelos de difusion.

        Tipo producto central del dominio ml. Usado como contrato compartido entre
        Python (diffusers, sd.cpp wrapper) y Go (orquestador). Serializa a JSON
        canonico via model_dump_json() para intercambio entre servicios.

        Attributes:
            prompt:          Descripcion textual positiva de la imagen a generar.
            negative_prompt: Descripcion de lo que se quiere evitar. None omite
                             el condicionamiento negativo (requiere soporte del modelo).
            seed:            Semilla para reproducibilidad. -1 usa semilla aleatoria.
            steps:           Numero de pasos de denoising. Rango tipico: 20-50.
                             LCM: 4-8 pasos. Valores altos aumentan calidad y tiempo.
            cfg_scale:       Classifier-Free Guidance scale. Controla cuanto el modelo
                             sigue el prompt. Rango tipico: 5.0-12.0.
                             7.5 es el valor clasico. LCM: 1.0-2.0.
            sampler:         Algoritmo de denoising. Ver SamplerName para valores validos.
            width:           Ancho de la imagen en pixeles. Debe ser multiplo de 8.
                             SD1.5: 512. SDXL: 1024. Flux: 1024+.
            height:          Alto de la imagen en pixeles. Mismas restricciones que width.
            model:           Referencia al modelo base. Ver ModelRef.
            loras:           Lista de adaptadores LoRA a aplicar. Lista vacia = sin LoRA.
            clip_skip:       Numero de capas CLIP a saltar desde el final del encoder.
                             None usa el valor por defecto del modelo. Tipico: 1-2 para anime.
        """

        model_config = ConfigDict(frozen=True)

        prompt: str
        negative_prompt: str | None = None
        seed: int
        steps: int
        cfg_scale: float
        sampler: SamplerName
        width: int
        height: int
        model: ModelRef
        loras: list[LoraRef] = []
        clip_skip: int | None = None

except ImportError:
    from dataclasses import dataclass, field

    @dataclass(frozen=True)
    class GenerationConfig:  # type: ignore[no-redef]
        """Contrato de parametros para generacion de imagenes (fallback dataclass).

        Usar la version pydantic cuando este disponible para validacion y
        serializacion JSON canonica compartida con Go.

        Attributes:
            prompt:          Descripcion textual positiva de la imagen.
            negative_prompt: Descripcion de lo que evitar. None = sin condicionamiento negativo.
            seed:            Semilla. -1 = aleatoria.
            steps:           Pasos de denoising (20-50 tipico, 4-8 para LCM).
            cfg_scale:       CFG scale (5.0-12.0 tipico, 1.0-2.0 para LCM).
            sampler:         Algoritmo de denoising (ver SamplerName).
            width:           Ancho en pixeles, multiplo de 8.
            height:          Alto en pixeles, multiplo de 8.
            model:           Referencia al modelo base (ModelRef).
            loras:           Lista de LoRAs a aplicar (LoraRef[]).
            clip_skip:       Capas CLIP a saltar desde el final. None = default del modelo.
        """

        prompt: str
        seed: int
        steps: int
        cfg_scale: float
        sampler: str
        width: int
        height: int
        model: object  # ModelRef
        negative_prompt: str | None = None
        loras: tuple = field(default_factory=tuple)
        clip_skip: int | None = None