Stable Diffusion 3.5 réimplémenté depuis zéro en pur PyTorch

(github.com/yousef-rafat)

2 points par GN⁺ 2025-06-15 | 1 commentaires | Partager sur WhatsApp

Le projet miniDiffusion est un projet open source qui réimplémente Stable Diffusion 3.5 depuis zéro en utilisant uniquement PyTorch
La structure du projet se distingue par son orientation vers des objectifs pédagogiques ainsi que vers l’expérimentation et le hacking
L’ensemble de la base de code compte environ 2 800 lignes, avec un code minimal couvrant le VAE, le DiT, ainsi que les scripts d’entraînement et de jeu de données
Parmi les principaux composants figurent le VAE, les encodeurs de texte CLIP et T5, un transformer de diffusion multimodal et la Joint Attention
Le projet inclut encore des fonctionnalités expérimentales et nécessite davantage de tests

Présentation du projet miniDiffusion

miniDiffusion est un projet open source qui réimplémente les fonctionnalités essentielles de Stable Diffusion 3.5 uniquement avec PyTorch
Par rapport à Stable Diffusion 3.5 existant, ce projet présente les avantages suivants

Avec une base de code d’environ 2 800 lignes, il est particulièrement adapté à l’analyse directe de l’architecture et à l’apprentissage
Il peut être utilisé de manière très utile pour diverses expériences de machine learning et pour le hacking de modèles
Il a très peu de dépendances et n’utilise qu’un minimum de bibliothèques

Structure principale et fichiers de configuration

dit.py : implémentation principale du modèle Stable Diffusion
dit_components.py : embeddings, normalisation, patch embedding et fonctions auxiliaires de DiT
attention.py : implémentation de l’algorithme de Joint Attention
noise.py : inclut un ordonnanceur Euler ODE pour Rectified Flow
t5_encoder.py, clip.py : implémentation des encodeurs de texte T5 et CLIP
tokenizer.py : implémentation des tokenizers Byte-Pair et Unigram
metrics.py : implémentation de la métrique d’évaluation FID (Fréchet inception distance)
common.py : fournit des fonctions utilitaires nécessaires à l’entraînement
common_ds.py : implémentation d’un dataset itérable qui convertit les images en données d’entraînement pour DiT
dossier model : stockage des checkpoints du modèle et des logs après l’entraînement
dossier encoders : stockage des checkpoints de modules séparés comme le VAE et CLIP

⚠️ Fonctionnalités expérimentales et nécessité de tests miniDiffusion inclut encore des fonctionnalités expérimentales et nécessite davantage de tests

Détail des fonctionnalités principales

Core Image Generation Modules

Implémentation du VAE, de CLIP et des encodeurs de texte T5
Implémentation des tokenizers Byte-Pair et Unigram

SD3 Components

Multi-Modal Diffusion Transformer Model
Implémentation du Flow-Matching Euler Scheduler
Logit-Normal Sampling
Intégration de l’algorithme Joint Attention

Scripts d’entraînement et d’inférence du modèle

Fournit des scripts d’entraînement et d’inférence pour SD3 (Stable Diffusion 3.5)

Licence

Publié sous licence MIT et conçu à des fins pédagogiques et expérimentales

Signification et avantages de ce projet open source

Il permet d’apprendre et de hacker directement une architecture de génération d’images de pointe au niveau de Stable Diffusion 3.5, en pur PyTorch
Le code est concis et indépendant, ce qui l’optimise pour l’analyse d’architecture / le tuning de modèles / la recherche sur de nouveaux algorithmes
Il permet de mettre directement en pratique les techniques récentes de multimodal, de transformer et d’attention
Il offre une base permettant d’expérimenter en toute sécurité, indépendamment des projets commerciaux

1 commentaires

GN⁺ 2025-06-15

Avis Hacker News

L’implémentation de référence de Flux a une structure vraiment minimaliste, donc cela vaut le coup d’y jeter un œil si le sujet vous intéresse
- GitHub de Flux
- Le projet minRF a l’avantage de permettre de démarrer facilement l’entraînement de petits modèles de diffusion en utilisant le rectified flow
- GitHub de minRF
- L’implémentation de référence de Stable Diffusion 3.5 est elle aussi assez concise, donc bien adaptée comme point de référence
- GitHub de SD 3.5
- Les implémentations de référence sont souvent mal maintenues et pleines de bugs
  - On peut prendre comme exemple le tokenizer CLIP d’OpenAI : ce n’est même pas la version utilisée pour l’entraînement officiel, juste une référence, et pourtant les bugs n’y sont pas corrigés et se retrouvent recopiés tels quels dans plusieurs projets
  - Flux non plus n’est pas celui réellement utilisé pour l’entraînement, et il contient des bugs qui peuvent poser quelques problèmes, notamment avec cudagraphs
  - Comme il dépend de la référence CLIP, et que CLIP lui-même a des bugs, ceux-ci se propagent logiquement aussi
Je me demande si cela veut dire que le projet miniDiffusion utilise le modèle Stable Diffusion 3.5
- Code concerné
- Le jeu de données d’entraînement est très petit et ne contient que des images liées à la mode
- Jeu de données mode
- Ce jeu de données sert à s’exercer au fine-tuning d’un modèle de diffusion
  - Le code réimplémente bien SD3, mais à cause des limites de mon matériel, les poids sont récupérés depuis HuggingFace
Je me demande si l’usage de PyTorch pur apporte un avantage de performance sur des GPU non NVIDIA, ou si PyTorch est tellement optimisé pour CUDA que les autres fabricants de GPU ne peuvent pas vraiment rivaliser
- PyTorch fonctionne plutôt bien aussi sur Apple Silicon
  - En revanche, les GPU d’Apple ont du mal à atteindre la puissance de calcul des GPU NVIDIA haut de gamme, donc la comparaison directe est difficile
  - À noter qu’il y a une particularité quand on utilise PyTorch sur Apple Silicon
    - Chaque tenseur est considéré comme « possédé » par un périphérique précis (CPU ou GPU), donc tout déplacement de données entraîne une copie complète
    - Le Mac a une architecture à mémoire unifiée, mais PyTorch continue malgré tout à fonctionner d’une manière qui effectue des copies de données
- Sur des appareils non NVIDIA comme AMD, il est aussi possible d’exécuter des charges de travail ML via Vulkan
  - Si des extensions récentes comme cooperative matrix, ainsi que la prise en charge de nouvelles fonctionnalités au niveau des pilotes, continuent d’être ajoutées, il est possible que l’écart de performance avec CUDA, de l’ordre de quelques pourcents, disparaisse
- La prise en charge de ROCm par PyTorch progresse très lentement, et même quand cela fonctionne, c’est lent
- PyTorch fonctionne bien sur ROCm, mais je ne sais pas vraiment si cela marche au point d’être totalement « au même niveau »
Dans le code PyTorch, au lieu 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)
```
il serait intéressant d’essayer quelque chose comme
```
    self.qkv = nn.Linear(embed_size, 3*embed_size, bias = False)
    # dans la fonction forward
    qkv = self.qkv(x)
```
- Cela dit, dans ce cas, au lieu que les paramètres de q, k et v soient reliés indépendamment comme à l’origine, les paramètres entre q, k et v se retrouvent liés
  - Désolé si je me trompe complètement parce que je suis trop fatigué en ce moment
Cela a l’air d’être une bonne ressource pour les apprenants
- Je me demande s’il existe un tutoriel ou une documentation qu’un débutant pourrait suivre
- Il existe un cours fast.ai qui montre comment implémenter Stable Diffusion soi-même
  - Cours fast.ai
Je me demande si cela veut dire qu’on peut utiliser Stable Diffusion sans restriction de licence
- Non : l’algorithme d’inférence / d’entraînement (les mathématiques elles-mêmes) n’est pas soumis au droit d’auteur, mais l’OP n’a fait que réimplémenter le code
  - Ce qui pose problème côté droit d’auteur, c’est le modèle lui-même (les poids), et l’OP ne l’a pas entraîné lui-même faute de données et de puissance de calcul
Honnêtement, même si c’est un peu embarrassant à admettre, je me demande ce que nous avons réellement gagné avant et après l’apparition de ce dépôt
- Personnellement, j’ai plutôt évité de fabriquer des modèles moi-même et je me suis surtout contenté d’observer les résultats de loin
- Je supposais vaguement qu’il existait déjà des scripts d’inférence / d’entraînement basés sur PyTorch et accessibles publiquement
- Je pensais au moins que les scripts d’inférence seraient fournis avec la distribution du modèle, et qu’il existerait aussi des scripts de fine-tuning / d’entraînement
- Je ne sais pas très bien si ce projet est une réécriture « clean room » ou « dirty room » de l’existant, ou si au contraire les anciens codes PyTorch eux-mêmes étaient déjà tellement mêlés à CUDA/C que la version en pur PyTorch a une vraie valeur
- Bref, je ne suis pas sûr de bien comprendre, donc si quelqu’un peut l’expliquer, ce serait apprécié
  - La valeur principale de ce projet est une « implémentation à dépendances minimales »
    - Je n’ai jamais fait tourner SD 3.5 moi-même, mais c’est construit autour de la bibliothèque huggingface, et personnellement je trouve que huggingface a des dépendances tellement complexes qu’il est souvent difficile de lancer quoi que ce soit sans reproduire presque exactement l’environnement du développeur
    - En particulier, quelques mois ou quelques années après la sortie originale, faire tourner certains modèles devient très difficile
    - Par exemple, si on regarde le fichier requirements.txt de l’implémentation de référence SD3.5 de Stability AI, il n’y a même pas de versions épinglées, et il inclut des bibliothèques énormes comme transformers, ce qui est vraiment pénible en contexte professionnel
  - Stability AI distribue les modèles Stable Diffusion sous la Stability AI Community License, qui n’est pas « totalement libre » contrairement à la licence MIT
    - Il n’est pas permis de modifier les poids de certaines manières
    - Ce paquet permet de faire tourner le modèle (inférence) ou de faire du fine-tuning à partir de poids IA déjà existants
    - C’est excellent pour l’apprentissage et l’étude, mais les questions de licence restent présentes
Quand je pense à SD 3.5 (ou à n’importe quelle autre version), j’ai tendance à considérer que l’essentiel réside dans les poids générés pendant l’entraînement
- Le code lui-même me paraît relativement moins important du point de vue de la qualité du résultat ou des performances
- Cela dit, ce n’est peut-être pas une appréciation juste, et ce n’est pas une tentative de dénigrer ce genre d’effort
Je me demande dans quelle mesure le code source académique original publié par le groupe CompViz de la Ludwig Maximilian University est exploitable en pratique
Je me demande si l’implémentation du diffusion transformer (DiT) ici reproduit correctement la cross-token attention comme dans la version complète de SD 3.5, ou si cela a été simplifié pour rendre le code plus lisible