# 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
```python
Input: timestep escalar (int)
  ↓
Embedding sinusoidal (320 dim)
  ↓
MLP: 320 → 1280 → 1280
  ↓
Output: embedding temporal (1280 dim)
```

### ResnetBlock2D
```python
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)
```python
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)
```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
```python
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