nota-ai/bk-sdm-tiny
nota-ai
Texto a imagen
BK-SDM-Tiny es una versión comprimida arquitectónicamente de Stable Diffusion v1.4 para generación texto-a-imagen eficiente. Reduce la U-Net eliminando bloques residuales y de atención, y usa destilación de conocimiento para imitar el comportamiento del modelo original con solo 212.776 pares imagen-texto de LAION-Aesthetics V2 6.5+. La variante Tiny tiene una U-Net de 0,33B parámetros y un modelo completo de 0,50B parámetros, orientado a inferencia más ligera manteniendo capacidades generales de síntesis visual.
Como usar
Instalación y uso con Diffusers:
pip install -U diffusers transformers accelerate
import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("nota-ai/bk-sdm-tiny", dtype=torch.bfloat16, device_map="cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]
Ejemplo con StableDiffusionPipeline, scheduler PNDM por defecto y 50 pasos de denoising:
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-tiny", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
prompt = "a tropical bird sitting on a branch of a tree"
image = pipe(prompt).images[0]
image.save("example.png")
También puede usarse cargando Stable Diffusion v1.4 y sustituyendo solo la U-Net comprimida:
import torch
from diffusers import StableDiffusionPipeline, UNet2DConditionModel
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe.unet = UNet2DConditionModel.from_pretrained("nota-ai/bk-sdm-tiny", subfolder="unet", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
prompt = "a tropical bird sitting on a branch of a tree"
image = pipe(prompt).images[0]
image.save("example.png")
Funcionalidades
- Generación texto-a-imagen basada en Stable Diffusion/Diffusers.
- U-Net comprimida mediante eliminación de bloques en etapas externas, medias e internas.
- Destilación a nivel de características y salida, además de pérdida de denoising.
- Compatible con StableDiffusionPipeline y con sustitución directa de la U-Net en Stable Diffusion v1.4.
- Formato Safetensors y licencia CreativeML OpenRAIL-M.
- Evaluado en MS-COCO con 512x512 imágenes, scheduler PNDM y métricas FID, IS y CLIP Score.
Casos de uso
- Investigación sobre compresión arquitectónica de modelos de difusión texto-a-imagen.
- Experimentos con generación visual eficiente usando una versión ligera de Stable Diffusion.
- Pruebas educativas o creativas de síntesis de imágenes con recursos más limitados.
- Estudio de limitaciones, sesgos y despliegue seguro de modelos generativos.
- Prototipos que necesiten una U-Net compatible con Stable Diffusion v1.4 pero más pequeña.