Flow-Based Generative Models

# ---Importing Necessary Libraries---
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.datasets import FashionMNIST
from torchvision.transforms import ToTensor, Normalize, Compose
from torch.utils.data import DataLoader

# --- Conditional Coupling Layer ---
class ConditionalCouplingLayer(nn.Module):
    def __init__(self, in_channels, cond_channels, hidden_dim=128):
        super().__init__()
        self.scale_net = nn.Sequential(
            nn.Conv2d(in_channels // 2 + cond_channels, hidden_dim, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_dim, in_channels // 2, kernel_size=3, padding=1),
            nn.Tanh()
        )
        self.shift_net = nn.Sequential(
            nn.Conv2d(in_channels // 2 + cond_channels, hidden_dim, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_dim, in_channels // 2, kernel_size=3, padding=1),
        )

    def forward(self, x, cond, invert=False):
        x1, x2 = x.chunk(2, dim=1)  # split channels into 2 parts
        h = torch.cat([x1, cond], dim=1)  # concat condition on channels

        s = self.scale_net(h)
        t = self.shift_net(h)

        if not invert:
            y2 = x2 * torch.exp(s) + t
        else:
            y2 = (x2 - t) * torch.exp(-s)

        return torch.cat([x1, y2], dim=1)

# --- Simple Conditional Flow Model ---
class ConditionalFlow(nn.Module):
    def __init__(self, in_channels, cond_channels, num_layers=4):
        super().__init__()
        self.layers = nn.ModuleList([
            ConditionalCouplingLayer(in_channels, cond_channels) for _ in range(num_layers)
        ])

    def forward(self, x, cond, invert=False):
        if not invert:
            for layer in self.layers:
                x = layer(x, cond, invert=False)
        else:
            for layer in reversed(self.layers):
                x = layer(x, cond, invert=True)
        return x

# --- Helper: normalize input images and condition ---
transform = Compose([ToTensor(), Normalize((0.5,), (0.5,))])

# --- Load FashionMNIST ---
train_dataset = FashionMNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# --- Function to duplicate channel to get 2 channels ---
def duplicate_channel(x):
    return torch.cat([x, x], dim=1)  # duplicate along channel dim

# --- Instantiate model ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
in_channels = 2       # doubled channels now
cond_channels = 1     # condition channel remains 1 (original image)
model = ConditionalFlow(in_channels=in_channels, cond_channels=cond_channels).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# --- Training Loop ---
epochs = 5
for epoch in range(epochs):
    for x, _ in train_loader:
        cond = x.to(device)            # condition: 1 channel
        target = duplicate_channel(cond)  # target: 2 channels

        z = model(target, cond, invert=False)
        loss = torch.mean(z ** 2)  # encourage latent ~ N(0)

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

    print(f"Epoch {epoch + 1}/{epochs} Loss: {loss.item():.4f}")

# --- Generate new images conditioned on input ---
model.eval()
with torch.no_grad():
    sample_input, _ = next(iter(train_loader))
    cond_img = sample_input[:1].to(device)               # 1 channel cond image
    z_sample = torch.randn(1, in_channels, 28, 28).to(device)  # latent with 2 channels

    generated_img = model(z_sample, cond_img, invert=True)  # generate image conditioned on cond_img

import matplotlib.pyplot as plt

def show(img_tensor, title=''):
    img = img_tensor.squeeze().cpu().numpy()
    plt.imshow(img, cmap='gray')
    plt.title(title)
    plt.axis('off')
    plt.show()

show(cond_img, title='Condition Image (1 channel)')
show(generated_img[:, 0:1, :, :], title='Generated Image (first channel)')