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

• miniDiffusion est un projet qui réimplémente le modèle Stable Diffusion 3.5 en pur PyTorch avec un minimum de dépendances, conçu pour l’apprentissage, l’expérimentation et le bidouillage
• L’implémentation complète, du VAE au DiT en passant par les scripts d’entraînement et de jeu de données, représente environ 2 800 lignes et vise à minimiser le code nécessaire pour reproduire Stable Diffusion 3.5 depuis zéro
• Le code principal du modèle se trouve dans dit.py, dit_components.py et attention.py, avec une séparation entre Joint Attention, les embeddings, la normalisation, le patch embedding et les fonctions auxiliaires du DiT
• Les composants incluent VAE, CLIP, les encodeurs de texte T5, les tokenizers Byte-Pair et Unigram, un Multi-Modal Diffusion Transformer, un Flow-Matching Euler Scheduler et le Logit-Normal Sampling
• Le dépôt contient encore des fonctionnalités expérimentales et nécessite davantage de tests ; il est proposé sous licence MIT à des fins éducatives et expérimentales

---

Objectif et périmètre de miniDiffusion

• miniDiffusion est un projet qui réimplémente le modèle Stable Diffusion 3.5 en pur PyTorch avec un minimum de dépendances
• Il a été conçu pour l’apprentissage, l’expérimentation et le bidouillage, avec pour objectif de réduire la quantité de code nécessaire pour reproduire Stable Diffusion 3.5 depuis zéro
• L’implémentation représente environ 2 800 lignes, en incluant le VAE, le DiT, les scripts d’entraînement et les scripts de jeu de données

Principaux fichiers

• Le code central du modèle Stable Diffusion se trouve dans les fichiers suivants
  • dit.py : code principal du modèle DiT
  • dit_components.py : embeddings, normalisation, patch embedding et fonctions auxiliaires du DiT
  • attention.py : implémentation de Joint Attention
• noise.py contient un Euler Scheduler pour résoudre l’ODE du Rectified Flow
• Les encodeurs de texte et les tokenizers sont organisés dans des fichiers séparés
  • t5_encoder.py : encodeur de texte T5
  • clip.py : implémentation de CLIP
  • tokenizer.py : tokenizers T5 et CLIP
• metrics.py implémente le Fréchet Inception Distance (FID)
• Le code d’assistance à l’entraînement et de transformation des données se trouve dans les fichiers suivants
  • common.py : fonctions utilitaires pour l’entraînement
  • common_ds.py : implémentation d’un iterable dataset transformant les images en données d’entraînement pour le DiT

Dossiers et checkpoints

• Le dossier model stocke les checkpoints du modèle et les logs après l’entraînement
• Le dossier encoders stocke les checkpoints d’autres modules comme le VAE et CLIP

Composants inclus

• Modules centraux de génération d’images

  • VAE

  • CLIP

  • Encodeurs de texte T5

    • tokenizers Byte-Pair et Unigram
    • composants liés à Stable Diffusion 3
    • modèle Multi-Modal Diffusion Transformer
    • Flow-Matching Euler Scheduler
    • Logit-Normal Sampling
    • Joint Attention
    • inclut des scripts d’entraînement et d’inférence pour Stable Diffusion 3

Installation et préparation avant utilisation

• Cloner le dépôt

git clone "https://github.com/yousef-rafat/miniDiffusion";

• Installer les dépendances

pip install -r requirements.txt

• Avant d’installer les checkpoints du modèle, il faut ajouter un token Hugging Face dans get_checkpoints.py

python3 encoders/get_checkpoints.py

État et licence

• Le dépôt contient encore des fonctionnalités expérimentales et nécessite davantage de tests
• Le projet est proposé sous licence MIT à des fins éducatives et expérimentales