wahib-elbessali/cgan-stylegan2-ada

Conditional StyleGAN2-ADA — Medical Imaging

Class-conditional image synthesis across three medical imaging domains, trained with parallel distributed weight averaging across 5 GPU workers.

Models

File	Dataset	Classes
brain-mri/checkpoint.pt	Brain Tumor MRI	glioma_tumor, meningioma_tumor, no_tumor, pituitary_tumor
chestx-ray14/checkpoint.pt	NIH ChestX-ray14 (4-class subset)	No Finding, Effusion, Atelectasis, Pneumonia
oasis/checkpoint.pt	OASIS Alzheimer MRI	Non Demented, Very mild Dementia, Mild Dementia, Moderate Dementia

Usage

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

# ── Architecture ──────────────────────────────────────────────────────────────

class EqualLinear(nn.Module):
    def __init__(self, in_dim, out_dim, bias=True, bias_init=0.0, lr_mul=1.0):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_dim, in_dim) / lr_mul)
        self.bias   = nn.Parameter(torch.full([out_dim], float(bias_init))) if bias else None
        self.scale  = (1 / math.sqrt(in_dim)) * lr_mul
        self.lr_mul = lr_mul
    def forward(self, x):
        b = self.bias * self.lr_mul if self.bias is not None else None
        return F.linear(x, self.weight * self.scale, b)

class ModulatedConv2d(nn.Module):
    def __init__(self, in_ch, out_ch, ks, w_dim, demodulate=True, upsample=False):
        super().__init__()
        self.in_ch = in_ch; self.out_ch = out_ch; self.ks = ks
        self.demodulate = demodulate; self.upsample = upsample; self.padding = ks // 2
        self.affine = EqualLinear(w_dim, in_ch, bias_init=1.0)
        self.weight = nn.Parameter(torch.randn(out_ch, in_ch, ks, ks))
        self.scale  = 1 / math.sqrt(in_ch * ks ** 2)
    def forward(self, x, w):
        B, C, H, W = x.shape
        s = self.affine(w)
        weight = self.weight.unsqueeze(0) * self.scale * s.reshape(B, 1, -1, 1, 1)
        if self.demodulate:
            weight = weight * torch.rsqrt(weight.pow(2).sum([2,3,4], keepdim=True) + 1e-8)
        if self.upsample:
            x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
            H, W = H*2, W*2
        x = x.reshape(1, B*C, H, W)
        weight = weight.reshape(B*self.out_ch, self.in_ch, self.ks, self.ks)
        out = F.conv2d(x, weight, padding=self.padding, groups=B)
        return out.reshape(B, self.out_ch, out.shape[-2], out.shape[-1])

class NoiseInjection(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(1))
    def forward(self, x):
        B, C, H, W = x.shape
        return x + self.weight * torch.randn(B, 1, H, W, device=x.device, dtype=x.dtype)

def _nf(stage): return min(512, int(512 / 2 ** max(0, stage - 1)))

class MappingNetwork(nn.Module):
    def __init__(self, z_dim, embed_dim, w_dim, num_classes, num_layers=8):
        super().__init__()
        self.embed = nn.Embedding(num_classes, embed_dim)
        layers = [EqualLinear(z_dim + embed_dim, w_dim, lr_mul=0.01), nn.LeakyReLU(0.2)]
        for _ in range(num_layers - 1):
            layers += [EqualLinear(w_dim, w_dim, lr_mul=0.01), nn.LeakyReLU(0.2)]
        self.net = nn.Sequential(*layers)
    def forward(self, z, c):
        return self.net(torch.cat([F.normalize(z, dim=1),
                                   F.normalize(self.embed(c), dim=1)], dim=1))

class SynthesisBlock(nn.Module):
    def __init__(self, in_ch, out_ch, w_dim, img_ch=1, is_first=False):
        super().__init__()
        self.is_first = is_first
        if is_first:
            self.const = nn.Parameter(torch.randn(1, in_ch, 4, 4))
            self.conv1 = ModulatedConv2d(in_ch, out_ch, 3, w_dim)
        else:
            self.conv1 = ModulatedConv2d(in_ch, out_ch, 3, w_dim, upsample=True)
        self.conv2  = ModulatedConv2d(out_ch, out_ch, 3, w_dim)
        self.noise1 = NoiseInjection(); self.noise2 = NoiseInjection()
        self.act    = nn.LeakyReLU(0.2)
        self.to_rgb = ModulatedConv2d(out_ch, img_ch, 1, w_dim, demodulate=False)
        self.up_rgb = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
    def forward(self, x, w, prev_rgb=None):
        if self.is_first:
            x = self.const.repeat(w.shape[0], 1, 1, 1)
        x   = self.act(self.noise1(self.conv1(x, w)))
        x   = self.act(self.noise2(self.conv2(x, w)))
        rgb = self.to_rgb(x, w)
        if prev_rgb is not None:
            rgb = rgb + self.up_rgb(prev_rgb)
        return x, rgb

class Generator(nn.Module):
    def __init__(self, z_dim, w_dim, embed_dim, num_classes, img_channels, resolution):
        super().__init__()
        self.mapping = MappingNetwork(z_dim, embed_dim, w_dim, num_classes)
        self.blocks  = nn.ModuleList()
        _prev = _nf(0)
        for s in range(int(math.log2(resolution)) - 1):
            _out = _nf(s)
            self.blocks.append(SynthesisBlock(_prev, _out, w_dim, img_channels, s == 0))
            _prev = _out
    def forward(self, z, c):
        w = self.mapping(z, c)
        x = rgb = None
        for block in self.blocks:
            x, rgb = block(x, w, prev_rgb=rgb)
        return rgb

# ── Load & generate ───────────────────────────────────────────────────────────

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Choose: "brain-mri/checkpoint.pt", "chestx-ray14/checkpoint.pt", "oasis/checkpoint.pt"
ckpt_path = "brain-mri/checkpoint.pt"
classes   = ["glioma_tumor", "meningioma_tumor", "no_tumor", "pituitary_tumor"]

local = hf_hub_download(repo_id="wahib-elbessali/cgan-stylegan2-ada", filename=ckpt_path)
ckpt  = torch.load(local, map_location=device, weights_only=False)

G = Generator(z_dim=512, w_dim=512, embed_dim=512, num_classes=4,
              img_channels=1, resolution=128).to(device).eval()
sd = ckpt["G_ema"]
if next(iter(sd)).startswith("_orig_mod."):
    sd = {k[len("_orig_mod."):]: v for k, v in sd.items()}
G.load_state_dict(sd)

# Generate 4 images of class 0 (glioma_tumor)
with torch.no_grad():
    z = torch.randn(4, 512, device=device)
    c = torch.zeros(4, dtype=torch.long, device=device)
    imgs = G(z, c).cpu()

for i, img in enumerate(imgs):
    arr = ((img.squeeze().numpy() + 1) * 127.5).clip(0, 255).astype("uint8")
    Image.fromarray(arr, "L").save(f"{classes[0]}_{i}.png")

Results

Model	FID ↓	KID ↓	IS ↑
Brain MRI	30.15	0.0158	2.99
ChestX-ray14	30.62	0.0306	2.48
OASIS	16.52	0.0044	1.81

Architecture

• Base: StyleGAN2-ADA (Karras et al., NeurIPS 2020) — adaptive discriminator augmentation for limited-data regimes
• Conditioning: Class embedding injected into generator (AdaIN) and discriminator (projection head), following Mirza & Osindero 2014
• Resolution: 128 × 128 px, grayscale
• Training: Parallel distributed weight averaging across 5 workers via FedAvg (McMahan et al., 2017). Brain MRI trained from scratch; ChestX-ray14 and OASIS each independently fine-tuned from the Brain MRI checkpoint.