Generative AI for Image

# GAN: two networks competing
# z (noise) -> [Generator] -> fake_image
# real_image + fake_image -> [Discriminator] -> real/fake

# Minimax game:
# min_G max_D V(D,G) =
#   E[log D(x)] + E[log(1 - D(G(z)))]

# Training:
# 1. Train D: maximize log(D(x)) + log(1-D(G(z)))
# 2. Train G: maximize log(D(G(z)))

# At equilibrium:
# D(x) = 0.5 (can't tell real from fake)
# G generates perfect images

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, latent_dim=100, img_dim=784):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, img_dim),
            nn.Tanh()  # output [-1, 1]
        )
    def forward(self, z):
        return self.net(z)

class Discriminator(nn.Module):
    def __init__(self, img_dim=784):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(img_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    def forward(self, x):
        return self.net(x)

# Create models
G = Generator(latent_dim=100, img_dim=784)
D = Discriminator(img_dim=784)

# Training loop
criterion = nn.BCELoss()
opt_D = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))
opt_G = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))

for epoch in range(100):
    for real_imgs, _ in dataloader:
        batch_size = real_imgs.size(0)
        real_imgs = real_imgs.view(batch_size, -1)

        # Labels
        real_labels = torch.ones(batch_size, 1)
        fake_labels = torch.zeros(batch_size, 1)

        # Train Discriminator
        z = torch.randn(batch_size, 100)
        fake_imgs = G(z).detach()

        real_loss = criterion(D(real_imgs), real_labels)
        fake_loss = criterion(D(fake_imgs), fake_labels)
        d_loss = (real_loss + fake_loss) / 2

        opt_D.zero_grad()
        d_loss.backward()
        opt_D.step()

        # Train Generator
        z = torch.randn(batch_size, 100)
        fake_imgs = G(z)
        g_loss = criterion(D(fake_imgs), real_labels)

        opt_G.zero_grad()
        g_loss.backward()
        opt_G.step()

class DCGenerator(nn.Module):
    def __init__(self, latent_dim=100):
        super().__init__()
        self.net = nn.Sequential(
            # z: (batch, 100, 1, 1)
            nn.ConvTranspose2d(100, 512, 4, 1, 0),
            nn.BatchNorm2d(512),
            nn.ReLU(),
            # -> (batch, 512, 4, 4)

            nn.ConvTranspose2d(512, 256, 4, 2, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            # -> (batch, 256, 8, 8)

            nn.ConvTranspose2d(256, 128, 4, 2, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            # -> (batch, 128, 16, 16)

            nn.ConvTranspose2d(128, 3, 4, 2, 1),
            nn.Tanh()
            # -> (batch, 3, 32, 32)
        )
    def forward(self, z):
        return self.net(z.view(-1, 100, 1, 1))

class DCDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            # (batch, 3, 32, 32)
            nn.Conv2d(3, 64, 4, 2, 1),
            nn.LeakyReLU(0.2),
            # -> (batch, 64, 16, 16)

            nn.Conv2d(64, 128, 4, 2, 1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            # -> (batch, 128, 8, 8)

            nn.Conv2d(128, 256, 4, 2, 1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),
            # -> (batch, 256, 4, 4)

            nn.Conv2d(256, 1, 4, 1, 0),
            nn.Sigmoid()
            # -> (batch, 1, 1, 1)
        )
    def forward(self, x):
        return self.net(x).view(-1, 1)

# Usage
D = DCDiscriminator()
z = torch.randn(8, 100)
fake_imgs = G(z)
pred = D(fake_imgs)  # (8, 1) - real/fake score

# Common issues:

# 1. Mode collapse
# G generates same image every time
# Fix: add noise, use WGAN

# 2. Training instability
# D loss -> 0 (D too strong)
# G loss -> infinity
# Fix: lower lr for D, train D less

# 3. No convergence
# Loss oscillates, never stabilizes
# Fix: use BatchNorm, adjust architecture

# Evaluation metrics:
# - FID (Fréchet Inception Distance)
#   lower = better
# - IS (Inception Score)
#   higher = better

# Tips:
# - Use Adam optimizer
# - Train D more than G
# - Add noise to labels

# Diffusion: two processes

# Forward: add noise gradually
# x_0 -> x_1 -> x_2 -> ... -> x_T (pure noise)
# q(x_t | x_{t-1}) = N(x_t; sqrt(1-beta_t)*x_{t-1}, beta_t*I)

# Reverse: learn to denoise
# x_T -> x_{T-1} -> ... -> x_1 -> x_0 (clean image)
# p(x_{t-1} | x_t) = N(x_{t-1}; mu_theta(x_t, t), sigma_t*I)

# Training:
# 1. Take image x_0
# 2. Sample random timestep t
# 3. Add noise: x_t = sqrt(alpha_t)*x_0 + sqrt(1-alpha_t)*eps
# 4. Predict noise: eps_theta(x_t, t)
# 5. Loss = MSE(eps, eps_theta(x_t, t))

class UNet(nn.Module):
    def __init__(self, in_channels=3, time_dim=256):
        super().__init__()
        # Encoder
        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)

        # Bottleneck
        self.bottleneck = self.conv_block(256, 512)

        # Decoder
        self.dec3 = self.conv_block(512 + 256, 256)
        self.dec2 = self.conv_block(256 + 128, 128)
        self.dec1 = self.conv_block(128 + 64, 64)

        # Time embedding
        self.time_mlp = nn.Sequential(
            nn.Linear(time_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 512)
        )

    def conv_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.GroupNorm(8, out_ch),
            nn.GELU()
        )

    def forward(self, x, t):
        # x: (batch, 3, H, W)
        # t: (batch, time_dim)
        t_emb = self.time_mlp(t)
        # ... encoder, bottleneck, decoder
        return noise_pred

def train_diffusion(model, dataloader, epochs=100):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    for epoch in range(epochs):
        for images, _ in dataloader:
            batch_size = images.size(0)

            # Sample random timestep
            t = torch.randint(0, T, (batch_size,))

            # Add noise
            noise = torch.randn_like(images)
            alpha_t = alphas[t].view(-1, 1, 1, 1)
            alpha_bar_t = alphas_cumprod[t].view(-1, 1, 1, 1)

            noisy_images = torch.sqrt(alpha_bar_t) * images + \
                          torch.sqrt(1 - alpha_bar_t) * noise

            # Predict noise
            noise_pred = model(noisy_images, t)

            # Loss
            loss = F.mse_loss(noise_pred, noise)

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

# alphas, alphas_cumprod = schedule for noise

@torch.no_grad()
def sample(model, num_samples, image_size):
    # Start from pure noise
    x = torch.randn(num_samples, 3, image_size, image_size)

    for t in reversed(range(T)):
        # Predict noise
        noise_pred = model(x, torch.tensor([t]))

        # Denoise
        alpha_t = alphas[t]
        alpha_bar_t = alphas_cumprod[t]
        beta_t = betas[t]

        # Remove noise
        x = (1 / torch.sqrt(alpha_t)) * \
            (x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * noise_pred)

        # Add noise (except at t=0)
        if t > 0:
            noise = torch.randn_like(x)
            x = x + torch.sqrt(beta_t) * noise

    return x

# Generate 16 images
samples = sample(model, 16, 32)
save_image(samples, 'generated.png', nrow=4)

from diffusers import StableDiffusionPipeline

# Load Stable Diffusion
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

# Generate from text prompt
prompt = "a photo of an astronaut riding a horse"
image = pipe(prompt, num_inference_steps=50).images[0]
image.save("astronaut.png")

# Steps:
# 1. Text -> CLIP -> text embeddings
# 2. Random noise in latent space
# 3. U-Net denoise with text condition
# 4. VAE decode latent -> image

# Parameters:
# - num_inference_steps: 20-50 (more = better)
# - guidance_scale: 7.5 (higher = follow prompt more)
# - seed: random seed for reproducibility

# Comparison:
#                 GAN        Diffusion
# Speed:          Fast       Slow
# Training:       Hard       Easy
# Quality:        Good       Excellent
# Diversity:      Low        High
# Mode collapse:  Yes        No
# Evaluation:     FID        FID/IS

# Use cases:
# GAN:
# - Real-time generation
# - Style transfer
# - Super-resolution

# Diffusion:
# - Text-to-image
# - Image editing
# - Inpainting
# - Video generation

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Load MNIST
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
])
dataset = datasets.MNIST('./data', train=True,
                         download=True, transform=transform)
loader = DataLoader(dataset, batch_size=128, shuffle=True)

# DCGAN for MNIST
G = DCGenerator(latent_dim=100)
D = DCDiscriminator()
criterion = nn.BCELoss()
opt_G = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))
opt_D = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))

# Train
for epoch in range(50):
    for real_imgs, _ in loader:
        batch_size = real_imgs.size(0)
        real_imgs = real_imgs.view(batch_size, -1)

        # Train D
        z = torch.randn(batch_size, 100)
        fake_imgs = G(z).detach()
        d_loss = criterion(D(real_imgs), torch.ones(batch_size, 1)) + \
                 criterion(D(fake_imgs), torch.zeros(batch_size, 1))
        opt_D.zero_grad()
        d_loss.backward()
        opt_D.step()

        # Train G
        z = torch.randn(batch_size, 100)
        fake_imgs = G(z)
        g_loss = criterion(D(fake_imgs), torch.ones(batch_size, 1))
        opt_G.zero_grad()
        g_loss.backward()
        opt_G.step()

uv add django torch torchvision

uv run django-admin startproject wk14 .
uv run manage.py startapp dashboard

# dashboard/views.py
import json, torch
from django.http import StreamingHttpResponse

def train(request, model_type):
    def event_stream():
        if model_type == 'gan':
            G = Generator()
            D = Discriminator()
        elif model_type == 'dcgan':
            G = DCGenerator()
            D = DCDiscriminator()

        criterion = nn.BCELoss()
        opt_G = torch.optim.Adam(G.parameters(), lr=0.0002)
        opt_D = torch.optim.Adam(D.parameters(), lr=0.0002)

        for epoch in range(50):
            # Training loop
            # ...

            # Generate sample images
            with torch.no_grad():
                z = torch.randn(16, 100)
                fake_imgs = G(z)

            yield f"data: {json.dumps({'epoch': epoch, 'd_loss': d_loss.item(), 'g_loss': g_loss.item(), 'images': fake_imgs.tolist()})}\n\n"

    return StreamingHttpResponse(event_stream(), content_type='text/event-stream')

from django.http import StreamingHttpResponse
import json, time

def train_stream(request, model_type):
    def event_stream():
        G, D = build_models(model_type)
        criterion = nn.BCELoss()
        opt_G = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))
        opt_D = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))

        for epoch in range(100):
            # Simplified training
            batch_size = 64

            # Train D
            z = torch.randn(batch_size, 100)
            fake_imgs = G(z).detach()
            d_loss = criterion(D(fake_imgs), torch.zeros(batch_size, 1))

            opt_D.zero_grad()
            d_loss.backward()
            opt_D.step()

            # Train G
            z = torch.randn(batch_size, 100)
            fake_imgs = G(z)
            g_loss = criterion(D(fake_imgs), torch.ones(batch_size, 1))

            opt_G.zero_grad()
            g_loss.backward()
            opt_G.step()

            # Generate samples
            with torch.no_grad():
                z = torch.randn(4, 100)
                samples = G(z).view(-1, 1, 28, 28)
                samples = (samples + 1) / 2  # denormalize

            yield f"data: {json.dumps({\n"
                  f"  'epoch': {epoch},\n"
                  f"  'd_loss': {d_loss.item():.4f},\n"
                  f"  'g_loss': {g_loss.item():.4f},\n"
                  f"  'samples': samples.numpy().tolist()\n"
                  f"})}\n\n"
            time.sleep(0.05)

    resp = StreamingHttpResponse(event_stream(),
        content_type='text/event-stream')
    resp['Cache-Control'] = 'no-cache'
    return resp

<!-- dashboard/templates/index.html -->
<div x-data="{ model: 'gan', training: false }">
  <h1>GAN Training Dashboard</h1>

  <select x-model="model">
    <option value="gan">Vanilla GAN</option>
    <option value="dcgan">DCGAN</option>
  </select>

  <button @click="startTrain()">
    Start Training
  </button>

  <div>
    <p>Epoch: <span x-text="epoch"></span></p>
    <p>D Loss: <span x-text="dLoss"></span></p>
    <p>G Loss: <span x-text="gLoss"></span></p>
  </div>

  <!-- Generated images grid -->
  <div class="grid grid-cols-4 gap-2">
    <template x-for="img in samples">
      <img :src="img" class="w-full" />
    </template>
  </div>

  <canvas id="lossChart" width="420" height="200"></canvas>
</div>

<script>
function startTrain() {
  const es = new EventSource('/train/' + model);
  es.onmessage = function(e) {
    const m = JSON.parse(e.data);
    epoch = m.epoch;
    dLoss = m.d_loss.toFixed(4);
    gLoss = m.g_loss.toFixed(4);
    samples = m.samples;
    drawLossChart(m.epoch, m.d_loss, m.g_loss);
  };
}
</script>