antonypotapenko/cyclegan-lab5

1from __future__ import annotations2 3from pathlib import Path4from typing import Optional5 6import torch7import torch.nn as nn8 9 10class ConvBlock(nn.Module):11    def __init__(12        self,13        in_channels: int,14        out_channels: int,15        kernel_size: int = 3,16        stride: int = 1,17        padding: int = 1,18        norm_type: Optional[str] = "instance",19        activation: Optional[str] = "relu",20        use_bias: Optional[bool] = None,21    ) -> None:22        super().__init__()23 24        if use_bias is None:25            use_bias = norm_type in ["instance", None]26 27        layers: list[nn.Module] = []28 29        if padding > 0:30            layers.append(nn.ReflectionPad2d(padding))31            conv_padding = 032        else:33            conv_padding = 034 35        layers.append(36            nn.Conv2d(37                in_channels,38                out_channels,39                kernel_size=kernel_size,40                stride=stride,41                padding=conv_padding,42                bias=use_bias,43            )44        )45 46        if norm_type is not None:47            if norm_type == "instance":48                layers.append(nn.InstanceNorm2d(out_channels, affine=True))49            elif norm_type == "batch":50                layers.append(nn.BatchNorm2d(out_channels))51            else:52                raise ValueError(f"Unknown norm_type: {norm_type}")53 54        if activation is not None:55            if activation == "relu":56                layers.append(nn.ReLU(inplace=True))57            elif activation == "lrelu":58                layers.append(nn.LeakyReLU(0.2, inplace=True))59            elif activation == "tanh":60                layers.append(nn.Tanh())61            else:62                raise ValueError(f"Unknown activation: {activation}")63 64        self.block = nn.Sequential(*layers)65 66    def forward(self, x: torch.Tensor) -> torch.Tensor:67        return self.block(x)68 69 70class ResidualBlock(nn.Module):71    def __init__(self, channels: int, norm_type: str = "instance") -> None:72        super().__init__()73        self.block = nn.Sequential(74            ConvBlock(75                channels,76                channels,77                kernel_size=3,78                stride=1,79                padding=1,80                norm_type=norm_type,81                activation="relu",82            ),83            ConvBlock(84                channels,85                channels,86                kernel_size=3,87                stride=1,88                padding=1,89                norm_type=norm_type,90                activation=None,91            ),92        )93 94    def forward(self, x: torch.Tensor) -> torch.Tensor:95        return x + self.block(x)96 97 98class UpsampleBlock(nn.Module):99    def __init__(self, in_channels: int, out_channels: int, norm_type: str = "instance") -> None:100        super().__init__()101        self.block = nn.Sequential(102            nn.Upsample(scale_factor=2, mode="nearest"),103            ConvBlock(104                in_channels,105                out_channels,106                kernel_size=3,107                stride=1,108                padding=1,109                norm_type=norm_type,110                activation="relu",111            ),112        )113 114    def forward(self, x: torch.Tensor) -> torch.Tensor:115        return self.block(x)116 117 118class Generator(nn.Module):119    def __init__(120        self,121        in_channels: int = 3,122        out_channels: int = 3,123        base_channels: int = 64,124        n_res_blocks: int = 6,125        norm_type: str = "instance",126    ) -> None:127        super().__init__()128 129        layers: list[nn.Module] = [130            ConvBlock(131                in_channels,132                base_channels,133                kernel_size=7,134                stride=1,135                padding=3,136                norm_type=norm_type,137                activation="relu",138            ),139            ConvBlock(140                base_channels,141                base_channels * 2,142                kernel_size=3,143                stride=2,144                padding=1,145                norm_type=norm_type,146                activation="relu",147            ),148            ConvBlock(149                base_channels * 2,150                base_channels * 4,151                kernel_size=3,152                stride=2,153                padding=1,154                norm_type=norm_type,155                activation="relu",156            ),157        ]158 159        for _ in range(n_res_blocks):160            layers.append(ResidualBlock(base_channels * 4, norm_type=norm_type))161 162        layers.extend(163            [164                UpsampleBlock(base_channels * 4, base_channels * 2, norm_type=norm_type),165                UpsampleBlock(base_channels * 2, base_channels, norm_type=norm_type),166                ConvBlock(167                    base_channels,168                    out_channels,169                    kernel_size=7,170                    stride=1,171                    padding=3,172                    norm_type=None,173                    activation=None,174                ),175                nn.Tanh(),176            ]177        )178 179        self.model = nn.Sequential(*layers)180 181    def forward(self, x: torch.Tensor) -> torch.Tensor:182        return self.model(x)183 184 185def build_generator() -> Generator:186    """Matches the notebook architecture used to save gen_a2b_custom.pth / gen_b2a_custom.pth."""187    model = Generator(188        in_channels=3,189        out_channels=3,190        base_channels=64,191        n_res_blocks=6,192        norm_type="instance",193    )194    return model195 196 197def load_generator(weights_path: str | Path, device: str | torch.device = "cpu") -> nn.Module:198    path = Path(weights_path)199    map_location = torch.device(device)200 201    # Optional fast path: use TorchScript if the user exported it locally.202    if path.suffix == ".pt" and path.name.endswith(("_jit.pt", ".jit.pt")):203        model = torch.jit.load(str(path), map_location=map_location)204        model.eval()205        return model206 207    model = build_generator().to(map_location)208    state = torch.load(str(path), map_location=map_location)209    model.load_state_dict(state, strict=True)210    model.eval()211    return model212