87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
"""sdcpp_python_load — carga un StableDiffusion (stable-diffusion-cpp-python) con cache global."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import os
|
|
from typing import Any
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
from model_ref import ModelRef
|
|
|
|
# Cache global: (model_key, wtype, n_threads) -> StableDiffusion object
|
|
_SD_CACHE: dict[tuple[str, str, int], Any] = {}
|
|
|
|
|
|
def _get_model_key(model: ModelRef) -> str:
|
|
"""Retorna la clave de cache para un ModelRef."""
|
|
return model.path if model.path else model.name
|
|
|
|
|
|
def sdcpp_python_load(
|
|
model: ModelRef,
|
|
n_threads: int = -1,
|
|
wtype: str = "default",
|
|
rng_type: str = "cuda",
|
|
) -> Any:
|
|
"""Carga un StableDiffusion via stable-diffusion-cpp-python con cache global.
|
|
|
|
Instancia StableDiffusion con el checkpoint indicado por model.path (o model.name
|
|
si path es None). Los objetos se cachean en memoria por (model_key, wtype, n_threads)
|
|
— una segunda llamada con los mismos parametros retorna la instancia cacheada sin
|
|
recargar el modelo del disco.
|
|
|
|
Args:
|
|
model: Referencia al modelo. model.path se usa si esta presente;
|
|
si no, model.name se pasa como model_path (ruta local o hub).
|
|
n_threads: Numero de hilos de CPU para inferencia. -1 usa todos los disponibles.
|
|
wtype: Tipo de pesos / cuantizacion en memoria.
|
|
Valores: "default" | "f32" | "f16" | "q8_0" | "q5_1" | "q5_0" | "q4_1" | "q4_0".
|
|
"default" respeta el tipo original del checkpoint.
|
|
rng_type: Tipo de generador de numeros aleatorios.
|
|
Valores: "std_default" | "cuda".
|
|
"cuda" produce resultados compatibles con la implementacion CUDA
|
|
incluso en CPU.
|
|
|
|
Returns:
|
|
Instancia StableDiffusion lista para llamar a generate_image().
|
|
|
|
Raises:
|
|
ImportError: Si stable_diffusion_cpp no esta instalado.
|
|
Instalar con: pip install stable-diffusion-cpp-python
|
|
OSError: Si el path del modelo no existe o es invalido.
|
|
"""
|
|
try:
|
|
from stable_diffusion_cpp import StableDiffusion
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"sdcpp_python_load requiere stable-diffusion-cpp-python. "
|
|
"Instalar con: pip install stable-diffusion-cpp-python\n"
|
|
"Para compilar sin CUDA: "
|
|
"CMAKE_ARGS='-DSD_CUDA=OFF' pip install stable-diffusion-cpp-python"
|
|
) from exc
|
|
|
|
model_key = _get_model_key(model)
|
|
cache_key = (model_key, wtype, n_threads)
|
|
|
|
if cache_key in _SD_CACHE:
|
|
return _SD_CACHE[cache_key]
|
|
|
|
load_path = model.path if model.path else model.name
|
|
|
|
sd = StableDiffusion(
|
|
model_path=load_path,
|
|
wtype=wtype,
|
|
n_threads=n_threads,
|
|
rng_type=rng_type,
|
|
)
|
|
|
|
_SD_CACHE[cache_key] = sd
|
|
return sd
|
|
|
|
|
|
def _clear_sd_cache() -> None:
|
|
"""Limpia el cache global de instancias StableDiffusion (uso interno y tests)."""
|
|
_SD_CACHE.clear()
|