Stable Diffusion 3.5 reimplementado do zero em PyTorch puro

• miniDiffusion é um projeto que reimplementa o modelo Stable Diffusion 3.5 em PyTorch puro com dependências mínimas, projetado para fins de educação, experimentação e hacking
• A implementação completa tem cerca de 2.800 linhas, do VAE ao DiT, incluindo scripts de treinamento e de dataset, com foco em minimizar o código necessário para reproduzir o Stable Diffusion 3.5 do zero
• O código principal do modelo está em dit.py, dit_components.py e attention.py, com Joint Attention, embeddings, normalização, patch embedding e funções auxiliares do DiT separados
• Os componentes incluem VAE, CLIP, codificador de texto T5, tokenizadores Byte-Pair e Unigram, Multi-Modal Diffusion Transformer, Flow-Matching Euler Scheduler e Logit-Normal Sampling
• O repositório ainda tem recursos experimentais, precisa de mais testes e é oferecido sob a MIT License para fins educacionais e experimentais

---

Objetivo e escopo do miniDiffusion

• miniDiffusion é um projeto que reimplementa o modelo Stable Diffusion 3.5 em PyTorch puro com dependências mínimas
• Foi criado para educação, experimentação e hacking, com foco em reduzir a quantidade de código necessária para reproduzir o Stable Diffusion 3.5 do zero
• A implementação tem cerca de 2.800 linhas, incluindo VAE, DiT, scripts de treinamento e scripts de dataset

Estrutura principal de arquivos

• O código central do modelo Stable Diffusion está nos seguintes arquivos
  • dit.py: código principal do modelo DiT
  • dit_components.py: embeddings, normalização, patch embedding e funções auxiliares do DiT
  • attention.py: implementação de Joint Attention
• noise.py contém o Euler Scheduler para resolver a ODE do Rectified Flow
• Os codificadores de texto e tokenizadores estão organizados em arquivos separados
  • t5_encoder.py: codificador de texto T5
  • clip.py: implementação do CLIP
  • tokenizer.py: tokenizadores do T5 e do CLIP
• metrics.py implementa o Fréchet Inception Distance (FID)
• O código auxiliar de treinamento e de transformação de dados está nos seguintes arquivos
  • common.py: funções auxiliares para treinamento
  • common_ds.py: implementação de iterable dataset que converte dados de imagem em dados de treinamento para o DiT

Pastas e checkpoints

• A pasta model armazena checkpoints do modelo e logs após o treinamento
• A pasta encoders armazena checkpoints de outros módulos, como VAE e CLIP

Componentes incluídos

• Módulos principais de geração de imagem

  • VAE

  • CLIP

  • T5 Text Encoders

    • tokenizadores Byte-Pair e Unigram
    • componentes relacionados ao Stable Diffusion 3
    • modelo Multi-Modal Diffusion Transformer
    • Flow-Matching Euler Scheduler
    • Logit-Normal Sampling
    • Joint Attention
    • inclui scripts de treinamento e inferência para Stable Diffusion 3

Instalação e preparação antes do uso

• Clonar o repositório

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

• Instalar as dependências

pip install -r requirements.txt

• Antes de instalar os checkpoints do modelo, é necessário adicionar o token do Hugging Face em get_checkpoints.py

python3 encoders/get_checkpoints.py

Estado e licença

• O repositório ainda contém recursos experimentais e precisa de mais testes
• O projeto é fornecido sob a MIT License para fins educacionais e experimentais