MODELO_BYTE_DREAM.md

Byte Dream - Modelo de IA Generativa

📋 Visão Geral

Byte Dream é um modelo de difusão para geração de imagens a partir de descrições textuais (text-to-image), baseado na arquitetura UNet com mecanismos de atenção.

---

🏗️ Arquitetura do Modelo

Componentes Principais

1. UNet2DConditionModel (571M parâmetros)

• Função: Predição de ruído condicionado por texto
• Arquitetura: Encoder-Decoder com skip connections

Estrutura:

Input (4 canais latentes)
    ↓
Down Block 0: 320 canais → 2 camadas ResNet + Atenção
    ↓
Down Block 1: 640 canais → 2 camadas ResNet + Atenção
    ↓
Down Block 2: 1280 canais → 2 camadas ResNet + Atenção
    ↓
Down Block 3: 1280 canais → 2 camadas ResNet + Atenção
    ↓
Middle Block: 1280 canais → ResNet + Atenção + ResNet
    ↓
Up Block 0: 1280 canais → 2 camadas ResNet + Atenção
    ↓
Up Block 1: 1280 canais → 2 camadas ResNet + Atenção
    ↓
Up Block 2: 640 canais → 2 camadas ResNet + Atenção (+skip)
    ↓
Up Block 3: 320 canais → 2 camadas ResNet + Atenção (+skip)
    ↓
Output: 4 canais

2. AutoencoderKL (VAE)

• Encoder: Comprime imagem 512x512x3 → 4x4x4 latentes
• Decoder: Reconstrói latentes → imagem 512x512x3
• Canais latentes: 4 (usando apenas a média do VAE)

3. Text Encoder

• Modelo: CLIP ViT-L/14 (Hugging Face)
• Dimensão: 768 dimensões
• Max tokens: 77 tokens

4. DDIM Scheduler

• Timesteps: 1000 passos de treino
• Schedule: Scaled linear (β: 0.00085 → 0.012)
• Amostragem: Determinística (η=0)

---

🔧 Componentes Detalhados

TimestepEmbedding

Input: timestep escalar (int)
  ↓
Embedding sinusoidal (320 dim)
  ↓
MLP: 320 → 1280 → 1280
  ↓
Output: embedding temporal (1280 dim)

ResnetBlock2D

Input: (B, C, H, W) + temb
  ↓
GroupNorm → SiLU → Conv2d
  ↓
Adicionar temb (projetado)
  ↓
GroupNorm → SiLU → Dropout → Conv2d
  ↓
Skip connection (conv 1x1 se necessário)
  ↓
Output: (B, C', H, W)

AttentionBlock (Cross-Attention)

Input: 
  - hidden_states: (B, C, H, W) ou (B, L, C)
  - encoder_hidden_states: (B, 77, 768)
  
Processamento:
  1. Se 4D: reshape (B,C,H,W) → (B, H*W, C)
  2. Projeções Q, K, V
  3. Multi-head attention (8 heads)
  4. Reshape de volta para 4D (se necessário)
  5. Projeção de saída
  
Output: (B, C, H, W) ou (B, L, C)

---

📊 Especificações Técnicas

Parâmetros

• UNet: 571,138,564 parâmetros
• Blocos: 4 down, 4 up, 1 middle
• Camadas por bloco: 2 ResNets + atenção cruzada
• Heads de atenção: 8
• Dimensão head: 40 (320/8)

Dimensões

Estágio	Canais	Resolução	Skip channels
Entrada	4	64x64	-
Down 0	320	64x64	320
Down 1	640	32x32	640
Down 2	1280	16x16	1280
Down 3	1280	8x8	1280
Middle	1280	8x8	-
Up 0	1280	8x8	1280
Up 1	1280	16x16	1280
Up 2	640	32x32	640
Up 3	320	64x64	320
Saída	4	64x64	-

---

🔄 Fluxo de Treinamento

Pipeline

1. Carregar dataset (imagens + captions)
   ↓
2. Pré-processamento:
   - Resize para 512x512
   - Normalização [-1, 1]
   - Data augmentation (flip, crop)
   ↓
3. Encoding:
   - Imagem → VAE → latentes (4x64x64)
   - Texto → CLIP → embeddings (77x768)
   ↓
4. Diffusion forward process:
   - Adicionar ruído nos latentes
   - Sample timestep aleatório
   ↓
5. Predição UNet:
   - Input: noisy_latents + timestep + text_embeds
   - Output: noise_pred
   ↓
6. Loss: MSE(noise_pred, noise_true)
   ↓
7. Backprop + otimizador Adam

Hiperparâmetros (config.yaml)

training:
  epochs: 100
  batch_size: 4
  learning_rate: 1e-5
  gradient_accumulation_steps: 1
  max_grad_norm: 1.0
  
data_augmentation:
  random_flip: true
  center_crop: true
  image_size: 512

---

🚀 Fluxo de Inferência

Geração de Imagens

1. Prompt → Text Encoder → embeddings
   ↓
2. Sample noise: (1, 4, 64, 64) ~ N(0,1)
   ↓
3. DDIM sampling loop (50 steps):
   for t in timesteps:
     noise_pred = UNet(noisy_latents, t, text_embeds)
     noisy_latents = scheduler.step(noise_pred, t, noisy_latents)
   ↓
4. Decode: latentes → VAE decoder → imagem
   ↓
5. Output: imagem 512x512 RGB

---

🛠️ Correções Implementadas

1. Timestep Embedding

• Problema: Linear layer esperava float mas recebia long
• Solução: Criar TimestepEmbedding com embedding sinusoidal

2. Atenção Cruzada

• Problema: Attention esperava 3D mas recebia 4D
• Solução: Reshape automático (B,C,H,W) ↔ (B,H*W,C)

3. Skip Connections

• Problema: Mismatch de canais entre down/up blocks
• Solução:
  • Projetar skip connections dinamicamente
  • Interpolar spatial dimensions se necessário
  • Equalizar número de camadas (layers_per_block=2)

4. VAE Latents

• Problema: VAE output tinha 8 canais (mean + logvar)
• Solução: Usar apenas primeiros 4 canais (mean)

---

📁 Estrutura de Arquivos

Byte Dream/
├── bytedream/
│   ├── model.py          # UNet, VAE, Text Encoder
│   ├── scheduler.py      # DDIM scheduler
│   ├── pipeline.py       # Pipeline de geração
│   └── utils.py          # Utilitários
├── train.py              # Script de treinamento
├── infer.py              # Inferência/generation
├── config.yaml           # Configurações
├── app.py                # Interface web
└── dataset/              # Dataset de treino
    ├── *.jpg             # Imagens
    └── *.txt             # Captions

---

🎯 Features Únicas

1. Otimizado para CPU: Focado em execução local
2. Arquitetura eficiente: 571M parâmetros (balanceado)
3. Skip connections adaptativas: Projeção automática de canais
4. Atenção flexível: Suporta tensores 3D e 4D
5. Dataset sintético: Geração automática de dados de treino

---

📈 Performance Esperada

• Treino: ~2-3 horas/epoch (CPU, 100 imagens)
• Inferência: ~30-60 segundos/imagem (50 steps, CPU)
• Qualidade: Bom para dataset pequeno, melhora com mais dados

---

🔮 Próximos Passos

1. Otimização:

   • Mixed precision (FP16)
   • Gradient checkpointing
   • ONNX runtime

2. Melhorias de Qualidade:

   • Mais dados de treino
   • Aumentar resolução
   • Fine-tuning de modelos pré-treinados

3. Deploy:

   • Exportar para Hugging Face
   • API REST
   • Interface Gradio/Streamlit

---

📝 Referências

• Stable Diffusion (RunwayML)
• DDIM Paper
• CLIP (OpenAI)
• U-Net Architecture

---

Autor: Byte Dream Team
Versão: 1.0.0
Licença: MIT