Kod - Min Bild AI

class VectorQuantizerEMA(nn.Module):
    def __init__(self, num_embeddings=1024, embedding_dim=512,
                 commitment_cost=0.25, decay=0.99):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.num_embeddings = num_embeddings
        self.commitment_cost = commitment_cost
        self.decay = decay

        # Codebook embeddings
        embed = torch.randn(num_embeddings, embedding_dim)
        self.register_buffer('embedding', embed)
        self.register_buffer('cluster_size', torch.zeros(num_embeddings))
        self.register_buffer('embed_avg', embed.clone())

    def forward(self, z):
        # z shape: [B, C, H, W] -> [B*H*W, C]
        z_flattened = z.permute(0, 2, 3, 1).reshape(-1, self.embedding_dim)

        # Hitta narmåste codebook entry (L2 distans)
        distances = torch.cdist(z_flattened, self.embedding)
        encoding_indices = distances.argmin(dim=1)

        # Kvantiserad output
        z_q = self.embedding[encoding_indices]
        z_q = z_q.view(z.shape[0], z.shape[2], z.shape[3], -1)
        z_q = z_q.permute(0, 3, 1, 2)  # Tillbaka till [B, C, H, W]

        # Straight-through estimator
        z_q = z + (z_q - z).detach()

        # Commitment loss
        loss = self.commitment_cost * F.mse_loss(z, z_q.detach())

        return z_q, loss, encoding_indices

class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.norm1 = nn.GroupNorm(32, in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.norm2 = nn.GroupNorm(32, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.act = nn.SiLU()

        # Skip connection om dimensioner skiljer
        self.skip = nn.Conv2d(in_channels, out_channels, 1) \
                    if in_channels != out_channels else nn.Identity()

    def forward(self, x):
        h = self.act(self.norm1(x))
        h = self.conv1(h)
        h = self.act(self.norm2(h))
        h = self.conv2(h)
        return h + self.skip(x)  # Residual connection

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model=384, num_heads=6, max_seq_len=4096):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 64 per head

        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(0.1)

        # Causal mask (triangular)
        mask = torch.tril(torch.ones(max_seq_len, max_seq_len))
        self.register_buffer('mask', mask.view(1, 1, max_seq_len, max_seq_len))

    def forward(self, x):
        B, T, C = x.shape

        # Beräkna Q, K, V
        qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.d_k)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)  # [3, B, H, T, d_k]

        # Scaled dot-product attention
        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # Applicera causal mask
        attn = attn.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        # Output
        out = (attn @ v).transpose(1, 2).reshape(B, T, C)
        return self.proj(out)

class TransformerBlock(nn.Module):
    def __init__(self, d_model=384, num_heads=6):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, num_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),  # Expansion 4x
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
            nn.Dropout(0.1)
        )

    def forward(self, x):
        # Pre-norm arkitektur
        x = x + self.attn(self.ln1(x))  # Attention + residual
        x = x + self.mlp(self.ln2(x))   # MLP + residual
        return x

def ssim_loss(img1, img2, window_size=11):
    """Structural Similarity Index Loss"""
    C1 = 0.01 ** 2  # Stabilitetskonstant för luminans
    C2 = 0.03 ** 2  # Stabilitetskonstant för kontrast

    # Gaussian window
    window = create_gaussian_window(window_size, img1.shape[1])
    window = window.to(img1.device)

    # Beräkna medelvärden
    mu1 = F.conv2d(img1, window, padding=window_size//2, groups=img1.shape[1])
    mu2 = F.conv2d(img2, window, padding=window_size//2, groups=img2.shape[1])

    mu1_sq, mu2_sq = mu1 ** 2, mu2 ** 2
    mu1_mu2 = mu1 * mu2

    # Beräkna varianser
    sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=img1.shape[1]) - mu1_sq
    sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=img2.shape[1]) - mu2_sq
    sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=img1.shape[1]) - mu1_mu2

    # SSIM formel
    ssim = ((2*mu1_mu2 + C1) * (2*sigma12 + C2)) / \
           ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))

    return 1 - ssim.mean()  # Loss = 1 - SSIM

def edge_loss(img1, img2):
    """Sobel-baserad kantloss"""
    # Sobel kernels
    sobel_x = torch.tensor([[-1, 0, 1],
                            [-2, 0, 2],
                            [-1, 0, 1]], dtype=torch.float32)
    sobel_y = torch.tensor([[-1, -2, -1],
                            [ 0,  0,  0],
                            [ 1,  2,  1]], dtype=torch.float32)

    # Konvertera till grayscale
    gray1 = 0.299*img1[:,0] + 0.587*img1[:,1] + 0.114*img1[:,2]
    gray2 = 0.299*img2[:,0] + 0.587*img2[:,1] + 0.114*img2[:,2]

    # Applicera Sobel
    edges1 = sobel_magnitude(gray1, sobel_x, sobel_y)
    edges2 = sobel_magnitude(gray2, sobel_x, sobel_y)

    return F.mse_loss(edges1, edges2)

def train_vqvae(model, dataloader, optimizer, epochs=1000):
    model.train()
    för epoch in range(epochs):
        total_loss = 0
        för batch in dataloader:
            images = batch.to(device)

            # Forward pass
            recon, vq_loss, _ = model(images)

            # Beräkna alla losses
            mse = F.mse_loss(recon, images)
            ssim = ssim_loss(recon, images)
            edge = edge_loss(recon, images)
            multi = multiscale_loss(recon, images)

            # Total loss (viktad summa)
            loss = mse + 0.1*ssim + 0.05*edge + 0.1*multi + vq_loss

            # Backprop
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        # Checkpoint var 5:e epok
        if epoch % 5 == 0:
            save_checkpoint(model, optimizer, epoch)

def generate_image(transformer, vqvae_decoder, temperature=1.0, top_k=50):
    """Generera ny bild autoregressivt"""
    transformer.eval()
    tokens = torch.zeros(1, 1, dtype=torch.long, device=device)

    # Generera 4096 tokens (64x64)
    för i in range(4096):
        # Forward genom transformer
        logits = transformer(tokens)[:, -1, :]  # Sista positionen

        # Temperature scaling
        logits = logits / temperature

        # Top-K filtering
        if top_k > 0:
            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1:]
            logits[indices_to_remove] = float('-inf')

        # Sampla nästa token
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        # Lagg till i sekvensen
        tokens = torch.cat([tokens, next_token], dim=1)

    # Dekoda tokens till bild via VQ-VAE decoder
    tokens = tokens[:, 1:].reshape(1, 64, 64)  # Ta bort start-token
    z_q = vqvae_decoder.codebook.embedding[tokens]
    z_q = z_q.permute(0, 3, 1, 2)  # [B, C, H, W]
    image = vqvae_decoder.decode(z_q)

    return image