Reimplementan Stable Diffusion 3.5 desde cero en PyTorch puro

 (github.com/yousef-rafat)
2 puntos por GN⁺ 2025-06-15 | 1 comentarios | Compartir por WhatsApp
  • El proyecto miniDiffusion es un proyecto open source que reimplementa desde cero el modelo Stable Diffusion 3.5 usando únicamente PyTorch
  • La estructura de este proyecto se caracteriza por estar enfocada en fines educativos y en experimentación y hacking
  • Todo el código base tiene alrededor de 2800 líneas y está compuesto con código mínimo, desde el VAE hasta el DiT, además de scripts de entrenamiento y de datasets
  • Entre sus componentes principales se incluyen VAE, CLIP, codificadores de texto T5, transformer de difusión multimodal, joint attention y más
  • Aún incluye funcionalidades experimentales, por lo que necesita más pruebas

Introducción al proyecto miniDiffusion

miniDiffusion es un proyecto open source que reimplementa las funciones principales de Stable Diffusion 3.5 usando solo PyTorch
Frente a Stable Diffusion 3.5 existente, este proyecto tiene las siguientes ventajas

  • El código base tiene unas 2,800 líneas, por lo que es pequeño y muy adecuado para analizar su estructura directamente y aprender de ella
  • Puede aprovecharse de forma muy útil para diversos experimentos de aprendizaje automático y hacking de modelos
  • Tiene muy pocas dependencias y usa solo un conjunto mínimo de librerías

Estructura principal y archivos de configuración

  • dit.py : implementación principal del modelo Stable Diffusion
  • dit_components.py : embeddings, normalización, patch embedding y funciones auxiliares de DiT
  • attention.py : implementación del algoritmo de Joint Attention
  • noise.py : incluye el scheduler Euler ODE para Rectified Flow
  • t5_encoder.py, clip.py : implementación de los codificadores de texto T5 y CLIP
  • tokenizer.py : implementación de tokenizadores Byte-Pair y Unigram
  • metrics.py : implementación de la métrica de evaluación FID (Fréchet inception distance)
  • common.py : proporciona funciones auxiliares necesarias para el entrenamiento
  • common_ds.py : implementación de un dataset iterable que convierte imágenes en datos de entrenamiento para DiT
  • carpeta model : guarda checkpoints del modelo y logs después del entrenamiento
  • carpeta encoders : guarda checkpoints de módulos separados como VAE y CLIP

⚠️ Funciones experimentales y necesidad de pruebas miniDiffusion aún incluye funciones experimentales y requiere más pruebas

Desglose detallado por funcionalidad principal

Core Image Generation Modules

  • Implementación de VAE, CLIP y codificadores de texto T5
  • Implementación de tokenizadores Byte-Pair y Unigram

Componentes de SD3

  • Modelo Transformer de Difusión Multimodal
  • Implementación de Flow-Matching Euler Scheduler
  • Logit-Normal Sampling
  • Incorporación del algoritmo de Joint Attention

Scripts de entrenamiento e inferencia del modelo

  • Incluye scripts de entrenamiento e inferencia para SD3 (Stable Diffusion 3.5)

Licencia

  • Se publica bajo la licencia MIT y fue creado con fines educativos y experimentales

Significado y ventajas de este proyecto open source

  • Permite entrenar y hacer hacking directamente sobre una arquitectura de modelo moderno de generación de imágenes al nivel de Stable Diffusion 3.5 usando PyTorch puro
  • Su código es conciso e independiente, por lo que está optimizado para análisis de arquitectura / ajuste de modelos / investigación de nuevos algoritmos
  • Permite practicar directamente técnicas modernas de multimodalidad, transformers y attention
  • Proporciona una base para experimentar con seguridad al margen de proyectos comerciales

Lecturas relacionadas

1 comentarios

 
GN⁺ 2025-06-15
Comentarios de Hacker News

  • La implementación de referencia de Flux tiene una estructura realmente minimalista, así que si a alguien le interesa vale la pena echarle un vistazo

    • GitHub de Flux

    • El proyecto minRF tiene la ventaja de que permite empezar fácilmente a entrenar modelos pequeños de difusión usando rectified flow

    • GitHub de minRF

    • La implementación de referencia de Stable Diffusion 3.5 también está escrita de forma bastante concisa, así que sirve bien como referencia

    • GitHub de SD 3.5

    • Las implementaciones de referencia muchas veces no están bien mantenidas y suelen tener bastantes bugs

      • Como ejemplo está el tokenizador CLIP de OpenAI: no es la versión usada en el entrenamiento oficial, solo una referencia, pero aun así el bug no se corrige y termina copiándose igual en varios proyectos
      • Flux tampoco fue lo que realmente se usó en entrenamiento, y tiene bugs que pueden causar algunos problemas con cudagraphs y similares
      • Depende de la referencia de CLIP, pero como CLIP mismo tiene bugs, la estructura hace que esos bugs también se propaguen

  • Me surgió la duda de si el proyecto miniDiffusion significa que usa el modelo Stable Diffusion 3.5

    • Código relacionado

    • El dataset de entrenamiento es muy pequeño e incluye solo fotos relacionadas con moda

    • Dataset de moda

    • Ese dataset es para practicar el fine-tuning de un modelo de difusión

      • Es una reimplementación en código de SD3, pero por las limitaciones de mi hardware usé pesos tomados de HuggingFace

  • Me pregunto si usar PyTorch puro trae ventajas de rendimiento en GPUs que no sean de NVIDIA, o si PyTorch está tan optimizado para CUDA que otros fabricantes de GPU no pueden competir

    • PyTorch funciona bastante bien también en Apple Silicon

      • Eso sí, como las GPUs de Apple difícilmente alcanzan el nivel de cómputo de las GPUs tope de gama de NVIDIA, es difícil compararlas directamente
      • Como referencia, hay una peculiaridad al usar PyTorch en Apple Silicon
        • Cada tensor se considera "propiedad" de un dispositivo específico (CPU o GPU), así que al mover datos se hace una copia completa
        • La Mac tiene memoria unificada, pero PyTorch igual mantiene una estructura donde copia los datos

    • También es posible correr cargas de trabajo de ML en dispositivos no NVIDIA como AMD a través de Vulkan

      • Si se siguen sumando extensiones introducidas recientemente, como cooperative matrix, y soporte para nuevas funciones a nivel de driver, podría desaparecer esa diferencia de rendimiento de unos pocos puntos porcentuales frente a CUDA

    • El soporte de PyTorch para ROCm avanza muy lentamente, y aun cuando logras hacerlo funcionar, es lento

    • PyTorch sí funciona bien en ROCm, pero no sé si llega a funcionar al nivel de ser realmente "equivalente"

  • En el código de PyTorch, en vez de

        self.q = nn.Linear(embed_size, embed_size, bias = False)
    self.k = nn.Linear(embed_size, embed_size, bias = False)
    self.v = nn.Linear(embed_size, embed_size, bias = False)

    sugieren que estaría bien probar algo como

        self.qkv = nn.Linear(embed_size, 3*embed_size, bias = False)
    # dentro de la función forward
    qkv = self.qkv(x)
    • Si haces eso, en vez de que los parámetros de q, k y v estén conectados de forma independiente, los parámetros entre q, k y v quedarían conectados
      • Si estoy confundido por estar demasiado cansado en este momento, disculpen

  • Parece un buen material para gente que está aprendiendo

    • Me pregunto si hay algún tutorial o guía que también pueda seguir un principiante

    • En fast.ai hay una clase donde implementan Stable Diffusion directamente

  • Me surge la duda de si esto significa que se puede usar Stable Diffusion sin restricciones de licencia

    • No; el algoritmo de inferencia/entrenamiento (las matemáticas en sí) no está sujeto a copyright, pero el OP solo reimplementó el código
      • Lo que sí entra en temas de copyright es el modelo mismo (los pesos), y el OP no lo entrenó por su cuenta porque no tiene ni los datos ni la potencia de cómputo

  • La verdad, aunque me da un poco de pena, me pregunto qué es lo nuevo que ganamos al comparar el antes y el después de que existieran estos repositorios

    • Personalmente he evitado meterme a crear modelos y más bien he visto los resultados desde fuera

    • Ya daba por hecho que antes también existían scripts públicos de inferencia/entrenamiento basados en PyTorch

    • Al menos pensaba que los scripts de inferencia venían junto con la distribución del modelo, y que también habría scripts de fine-tuning/entrenamiento

    • No me queda claro si este proyecto es una reescritura tipo "clean room" o "dirty room", o si incluso el código PyTorch existente era tan complejo por todo lo CUDA/C que una versión en PyTorch puro sí tiene mucho valor

    • En cualquier caso, no lo tengo claro y agradecería si alguien lo puede explicar

      • El valor principal de este proyecto es que es una "implementación con dependencias mínimas"

        • No he corrido SD 3.5 en la práctica, pero está construido sobre la biblioteca de huggingface y, en lo personal, siento que huggingface tiene dependencias tan complejas que si no reproduces casi exactamente el entorno del desarrollador, a veces ni siquiera corre
        • Sobre todo varios meses o años después del lanzamiento original, ejecutar ciertos modelos se vuelve muy difícil
        • Por ejemplo, si ves el archivo requirements.txt de la implementación de referencia de Stability AI para SD3.5, ni siquiera fija versiones e incluye bibliotecas enormes como transformers, así que en trabajo real esto sí resulta bastante problemático

      • Stability AI distribuye los modelos Stable Diffusion bajo la Stability AI Community License, que a diferencia de MIT no es "totalmente libre"

        • No está permitido modificar los pesos de ciertas maneras
        • Este paquete está hecho de modo que puedes ejecutar el modelo (inferencia) o hacer fine-tuning con pesos de IA ya existentes
        • Es excelente para estudiar o aprender, pero los temas de licencia siguen ahí

  • Cuando pienso en SD 3.5, o en cualquier versión, veo como lo más importante la parte de los pesos generados durante el entrenamiento

    • El código en sí me parece relativamente menos importante en términos de calidad del resultado o rendimiento
    • Pero quizá no sea una evaluación precisa, y no intento menospreciar este esfuerzo

  • Me pregunto qué tan utilizable en la práctica es el código fuente académico original que publicó el grupo CompViz de Ludwig Maximilian University

  • Me pregunto si la implementación de diffusion transformer (DiT) aquí implementa correctamente la cross-token attention como en la versión completa de SD 3.5, o si la simplificaron para hacer el código más legible