generator.py

# generator.py
import torch
import torch.nn as nn

class Generator128(nn.Module):
    """
    Unconditional GAN Generator for 128x128 dog-like dream images.
    Input: noise vector (100-dim) → reshaped to (100, 1, 1)
    Output: image (3, 128, 128) in range [-1, 1]
    """
    def __init__(self, nz=100, ngf=64, nc=3):
        super(Generator128, self).__init__()
        self.main = nn.Sequential(
            # Input: Z (100, 1, 1)
            nn.ConvTranspose2d(nz, ngf * 16, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 16),
            nn.ReLU(True),
            # 4x4

            nn.ConvTranspose2d(ngf * 16, ngf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # 8x8

            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # 16x16

            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # 32x32

            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # 64x64

            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # 128x128
        )

    def forward(self, x):
        return self.main(x)

def load_laikadream(model_path="model.pth", device="cpu"):
    model = Generator128()
    state_dict = torch.load(model_path, map_location=device, weights_only=True)
    model.load_state_dict(state_dict)
    model.eval()
    return model