VAE + U-Net: a arquitetura por trás do Stable Diffusion

# Inferência simplificada do Stable Diffusion
def generate(prompt, steps=30, cfg=7.5):
    # 1. Texto → embedding
    c = text_encoder(prompt)            # shape: [seq, dim]
    c_uncond = text_encoder("")         # para classifier-free guidance

    # 2. Inicia no latente (não no pixel!)
    z = torch.randn(1, 4, 64, 64)       # latente 64×64×4

    # 3. Loop de denoising no latente
    for t in scheduler.timesteps:       # 30 passos com DPM-Solver
        eps_cond   = unet(z, t, c)
        eps_uncond = unet(z, t, c_uncond)
        eps        = eps_uncond + cfg * (eps_cond - eps_uncond)
        z          = scheduler.step(eps, t, z)

    # 4. Decoder VAE devolve para pixel
    image = vae_decoder(z)              # shape: [1, 3, 512, 512]
    return image

class ResBlock(nn.Module):
    def __init__(self, ch_in, ch_out, time_dim):
        self.gn1   = GroupNorm(32, ch_in)
        self.conv1 = Conv2d(ch_in, ch_out, 3, padding=1)
        self.time  = Linear(time_dim, ch_out)      # FiLM-like
        self.gn2   = GroupNorm(32, ch_out)
        self.conv2 = Conv2d(ch_out, ch_out, 3, padding=1)
        self.skip  = Conv2d(ch_in, ch_out, 1) if ch_in != ch_out else nn.Identity()

    def forward(self, x, t_emb):
        h = self.conv1(F.silu(self.gn1(x)))
        h = h + self.time(F.silu(t_emb))[:, :, None, None]  # injeta t
        h = self.conv2(F.silu(self.gn2(h)))
        return h + self.skip(x)

class SpatialTransformer(nn.Module):
    def __init__(self, ch, n_heads, ctx_dim):
        self.norm = GroupNorm(32, ch)
        self.proj_in  = Conv2d(ch, ch, 1)
        self.attn_self  = MultiHeadAttention(ch, n_heads)            # tokens espaciais entre si
        self.attn_cross = MultiHeadAttention(ch, n_heads, ctx_dim)   # tokens espaciais ↔ texto
        self.ff   = FeedForward(ch)
        self.proj_out = Conv2d(ch, ch, 1)

    def forward(self, x, context):                # context = CLIP/T5 embedding
        B, C, H, W = x.shape
        h = self.proj_in(self.norm(x))
        h = h.reshape(B, C, H*W).transpose(1, 2)  # [B, HW, C] como tokens
        h = h + self.attn_self(h)                 # self-attn entre patches
        h = h + self.attn_cross(h, context)       # cross-attn com prompt
        h = h + self.ff(h)
        h = h.transpose(1, 2).reshape(B, C, H, W)
        return self.proj_out(h) + x