fn_registry/python/functions/ml/diffusers_unload.py

"""diffusers_unload — libera memoria de un pipeline diffusers y limpia cache global."""

from __future__ import annotations

import gc
import sys
import os
from typing import Any


def diffusers_unload(pipe: Any | None = None) -> None:
    """Libera la memoria ocupada por un pipeline diffusers.

    Si se pasa pipe, lo elimina con del y llama gc.collect() + empty_cache().
    Si pipe es None, limpia ademas el cache global de diffusers_load_pipeline
    (descarga TODOS los pipelines cacheados). En ambos casos invoca
    torch.cuda.empty_cache() si CUDA esta disponible.

    Args:
        pipe:  Pipeline a liberar. Si None, limpia el cache global completo
               de diffusers_load_pipeline ademas de llamar gc + empty_cache.

    Returns:
        None. Efecto secundario: memoria GPU/CPU liberada.
    """
    if pipe is None:
        # Limpiar cache global de diffusers_load_pipeline si esta importado
        try:
            # Importar el modulo para acceder a su cache interno
            load_module_path = os.path.join(os.path.dirname(__file__))
            if load_module_path not in sys.path:
                sys.path.insert(0, load_module_path)
            from diffusers_load_pipeline import _clear_pipeline_cache
            _clear_pipeline_cache()
        except ImportError:
            pass
    else:
        del pipe

    gc.collect()

    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass