fn_registry/python/functions/ml/diffusers_set_scheduler.py

"""diffusers_set_scheduler — cambia el scheduler de un pipeline diffusers."""

from __future__ import annotations

from typing import Any


# Mapping canónico sampler -> (scheduler_class_name, kwargs_extra)
_SCHEDULER_MAP: dict[str, tuple[str, dict]] = {
    "euler":       ("EulerDiscreteScheduler", {}),
    "euler_a":     ("EulerAncestralDiscreteScheduler", {}),
    "dpm++2m":     ("DPMSolverMultistepScheduler", {"algorithm_type": "dpmsolver++"}),
    "dpm++2m_v2":  ("DPMSolverMultistepScheduler", {"algorithm_type": "dpmsolver++", "solver_order": 2}),
    "heun":        ("HeunDiscreteScheduler", {}),
    "dpm2":        ("KDPM2DiscreteScheduler", {}),
    "lcm":         ("LCMScheduler", {}),
}


def diffusers_set_scheduler(pipe: Any, sampler: str) -> Any:
    """Reemplaza el scheduler de un pipeline diffusers por el correspondiente al sampler.

    Usa <SchedulerClass>.from_config(pipe.scheduler.config) para heredar la
    configuracion base del modelo (betas, clip_sample, etc.) y aplica encima
    los kwargs especificos del sampler. Modifica pipe.scheduler in-place y
    retorna el mismo pipe para composicion.

    Args:
        pipe:    Pipeline diffusers cargado (StableDiffusionPipeline,
                 StableDiffusionXLPipeline, etc.). Debe tener atributo
                 pipe.scheduler con .config.
        sampler: Nombre del sampler. Valores validos: euler, euler_a,
                 dpm++2m, dpm++2m_v2, heun, dpm2, lcm.

    Returns:
        El mismo pipe con pipe.scheduler reemplazado por la clase
        correspondiente al sampler solicitado.

    Raises:
        ImportError: Si diffusers no esta instalado.
        ValueError:  Si el sampler no esta en el mapping soportado.
    """
    try:
        import diffusers
    except ImportError as exc:
        raise ImportError(
            "diffusers_set_scheduler requiere diffusers. "
            "Instalar con: pip install diffusers"
        ) from exc

    if sampler not in _SCHEDULER_MAP:
        supported = ", ".join(sorted(_SCHEDULER_MAP.keys()))
        raise ValueError(
            f"Sampler '{sampler}' no soportado. Valores validos: {supported}"
        )

    class_name, extra_kwargs = _SCHEDULER_MAP[sampler]
    scheduler_cls = getattr(diffusers, class_name, None)

    if scheduler_cls is None:
        raise ImportError(
            f"La clase '{class_name}' no esta disponible en la version de diffusers "
            f"instalada. Actualizar diffusers para usar el sampler '{sampler}'."
        )

    pipe.scheduler = scheduler_cls.from_config(
        pipe.scheduler.config,
        **extra_kwargs,
    )
    return pipe