add stylegan

model/batik_stylegan.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba46e6a4db0ed1f2390b297d5272142a4f904ae8eadfd54f2d790995a8386215
+size 115666401

stylegan.py

@@ -1,368 +0,0 @@
-from torch import nn, optim
-import torch
-from torch.nn import functional as F
-from typing import Any, Callable, Optional
-import math
-
-class WSLinear(nn.Module):
-    '''
-    Weighted scale linear for equalized learning rate.
-
-    Args:
-        in_features (int): The number of input features.
-        out_features (int): The number of output features.
-    '''
-
-    def __init__(self, in_features: int, out_features: int) -> None:
-        super(WSLinear, self).__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-
-        self.linear = nn.Linear(self.in_features, self.out_features)
-        self.scale = (2 / self.in_features) ** 0.5
-        self.bias = self.linear.bias
-        self.linear.bias = None
-
-        self._init_weights()
-
-    def _init_weights(self) -> None:
-        nn.init.normal_(self.linear.weight)
-        nn.init.zeros_(self.bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.linear(x * self.scale) + self.bias
-    
-class WSConv2d(nn.Module):
-    """
-    Weight-scaled Conv2d layer for equalized learning rate.
-
-    Args:
-        in_channels (int): Number of input channels.
-        out_channels (int): Number of output channels.
-        kernel_size (int, optional): Size of the convolving kernel. Default: 3.
-        stride (int, optional): Stride of the convolution. Default: 1.
-        padding (int, optional): Padding added to all sides of the input. Default: 1.
-        gain (float, optional): Gain factor for weight initialization. Default: 2.
-    """
-    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, gain=2):
-        super().__init__()
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
-        self.scale = (gain / (in_channels * kernel_size ** 2)) ** 0.5
-        self.bias = self.conv.bias
-        self.conv.bias = None  # Remove bias to apply it after scaling
-
-        # Initialize weights
-        nn.init.normal_(self.conv.weight)
-        nn.init.zeros_(self.bias)
-
-    def forward(self, x):
-        return self.conv(x * self.scale) + self.bias.view(1, self.bias.shape[0], 1, 1)
-
-class Mapping(nn.Module):
-    '''
-    Mapping network.
-
-    Args:
-        features (int): Number of features in the input and output.
-        num_layers (int): Number of layers in the feed forward network.
-        num_styles (int): Number of styles to generate.
-    '''
-
-    def __init__(
-        self,
-        features: int,
-        num_styles: int,
-        num_layers: int = 8,
-    ) -> None:
-        super(Mapping, self).__init__()
-        self.features = features
-        self.num_layers = num_layers
-        self.num_styles = num_styles
-
-        layers = []
-        for _ in range(self.num_layers):
-            layers.append(WSLinear(self.features, self.features))
-            layers.append(nn.LeakyReLU(0.2))
-
-        self.fc = nn.Sequential(*layers)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        '''
-        Args:
-            x (torch.Tensor): Input tensor of shape (b, l).
-
-        Returns:
-            torch.Tensor: Output tensor with the same shape as input.
-        '''
-
-        x = self.fc(x) # (b, l)
-        return x
-
-class AdaIN(nn.Module):
-    '''
-    Adaptive Instance Normalization (AdaIN)
-    AdaIN(x_i, y) = y_s,i * (x_i - mean(x_i)) / std(x_i) + y_b,i
-
-    Args:
-        eps (float, optional): Small value to avoid division by zero. Default value is 0.00001.
-    '''
-
-    def __init__(self, eps: float= 1e-5) -> None:
-        super(AdaIN, self).__init__()
-        self.eps = eps
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        scale: torch.Tensor,
-        shift: torch.Tensor
-    ) -> torch.Tensor:
-        '''
-        Args:
-            x (torch.Tensor): Input tensor of shape (b, c, h, w).
-            scale (torch.Tensor): Scale tensor of shape (b, c).
-            shift (torch.Tensor): Shift tensor of shape (b, c).
-
-        Returns:
-            torch.Tensor: Output tensor of shape (b, c, h, w).
-        '''
-
-        b, c, *_ = x.shape
-
-        mean = x.mean(dim=(2, 3), keepdim=True) # (b, c, 1, 1)
-        std = x.std(dim=(2, 3), keepdim=True) # (b, c, 1, 1)
-        x_norm = (x - mean) / (std ** 2 + self.eps) ** .5
-
-        scale = scale.view(b, c, 1, 1) # (b, c, 1, 1)
-        shift = scale.view(b, c, 1, 1) # (b, c, 1, 1)
-        outputs = scale * x_norm + shift # (b, c, h, w)
-
-        return outputs
-
-class SynthesisLayer(nn.Module):
-    '''
-    Synthesis network layer which consist of:
-    - Conv2d.
-    - AdaIN.
-    - Affine transformation.
-    - Noise injection.
-
-    Args:
-        in_channels (int): The number of input channels.
-        out_channels (int): The number of output channels.
-        latent_features (int): The number of latent features.
-        use_conv (bool, optional): Whether to use convolution or not. Default value is True.
-    '''
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        latent_features: int,
-        use_conv: bool = True
-    ) -> None:
-        super(SynthesisLayer, self).__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.latent_features = latent_features
-        self.use_conv = use_conv
-
-        self.conv = nn.Sequential(
-            WSConv2d(self.in_channels, self.out_channels, kernel_size=3, padding=1),
-            nn.LeakyReLU(0.2)
-        ) if self.use_conv else nn.Identity()
-        self.norm = AdaIN()
-        self.scale_transform = WSLinear(self.latent_features, self.out_channels)
-        self.shift_transform = WSLinear(self.latent_features, self.out_channels)
-        self.noise_factor = nn.Parameter(torch.zeros(1, self.out_channels, 1, 1))
-
-        self._init_weights()
-
-    def _init_weights(self) -> None:
-        for m in self.modules():
-            if isinstance(m, (nn.Conv2d, nn.Linear)):
-                nn.init.normal_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-        nn.init.ones_(self.scale_transform.bias)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        w: torch.Tensor,
-        noise: Optional[torch.Tensor] = None
-    ) -> torch.Tensor:
-        '''
-        Args:
-            x (torch.Tensor): Input tensor of shape (b, c, h, w).
-            w (torch.Tensor): Latent space vector of shape (b, l).
-            noise (torch.Tensor, optional): Noise tensor of shape (b, 1, h, w). Default value is None.
-
-        Returns:
-            torch.Tensor: Output tensor of shape (b, c, h, w).
-        '''
-
-        b, _, h, w_ = x.shape
-        x = self.conv(x) # (b, o_c, h, w)
-        if noise is None:
-            noise = torch.randn(b, 1, h, w_, device=x.device) # (b, 1, h, w)
-        x += self.noise_factor * noise # (b, o_c, h, w)
-        y_s = self.scale_transform(w) # (b, o_c)
-        y_b = self.shift_transform(w) # (b, o_c)
-        x = self.norm(x, y_s, y_b) # (b, i_c, h, w)
-
-        return x
-    
-
-class SynthesisBlock(nn.Module):
-    '''
-    Synthesis network block which consist of:
-    - Optional upsampling.
-    - 2 Synthesis Layers.
-
-    Args:
-        in_channels (int): The number of input channels.
-        out_channels (int): The number of output channels.
-        latent_features (int): The number of latent features.
-        use_conv (bool, optional): Whether to use convolution or not. Default value is True.
-        upsample (bool, optional): Whether to use upsampling or not. Default value is True.
-    '''
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        latent_features: int,
-        *,
-        use_conv: bool = True,
-        upsample: bool = True
-     ) -> None:
-        super(SynthesisBlock, self).__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.latent_features = latent_features
-        self.use_conv = use_conv
-        self.upsample = upsample
-
-        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear') if self.upsample else nn.Identity()
-        self.layers = nn.ModuleList([
-            SynthesisLayer(self.in_channels, self.in_channels, self.latent_features, use_conv=self.use_conv),
-            SynthesisLayer(self.in_channels, self.out_channels, self.latent_features)
-        ])
-
-    def forward(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
-        '''
-        Args:
-            x (torch.Tensor): Input tensor of shape (b, c, h, w).
-            w (torch.Tensor): Latent vector of shape (b, l).
-
-        Returns:
-            torch.Tensor: Output tensor of shape (b, c, h, w) if not upsample else (b, c, 2h, 2w).
-        '''
-
-        x = self.upsample(x) # (b, c, h, w) if not upsample else (b, c, 2h, 2w)
-
-        for layer in self.layers:
-            x = layer(x, w) # (b, c, h, w) if not upsample else (b, c, 2h, 2w)
-
-        return x
-
-class Synthesis(nn.Module):
-    '''
-    Synthesis network which consist of:
-    - Constant tensor.
-    - Synthesis blocks.
-    - ToRGB convolutions.
-
-    Args:
-        resolution (int): The resolution of the image.
-        const_channels (int): The number of channels in the constant tensor. Default value is 512.
-    '''
-
-    def __init__(self, resolution: int, const_channels: int = 512) -> None:
-        super(Synthesis, self).__init__()
-        self.const_channels = const_channels
-        self.resolution = resolution
-
-        self.resolution_levels = int(math.log2(resolution) - 1)
-
-        self.constant = nn.Parameter(torch.ones(1, self.const_channels, 4, 4)) # (c, 4, 4)
-
-        in_channels = self.const_channels
-        blocks = [ SynthesisBlock(in_channels, in_channels, self.const_channels, use_conv=False, upsample=False) ]
-        to_rgb = [ WSConv2d(in_channels, 3, kernel_size=1, padding=0) ]
-
-        for _ in range(self.resolution_levels - 1):
-            blocks.append(SynthesisBlock(in_channels, in_channels // 2, self.const_channels))
-            to_rgb.append(WSConv2d(in_channels // 2, 3, kernel_size=1, padding=0))
-            in_channels //= 2
-
-        self.blocks = nn.ModuleList(blocks)
-        self.to_rgb = nn.ModuleList(to_rgb)
-
-    def forward(self, w: torch.Tensor, alpha: float, steps: int) -> torch.Tensor:
-        '''
-        Args:
-            w (torch.Tensor): Latent space vector of shape (b, l).
-            alpha (float): Fade in alpha value.
-            steps (int): The number of steps starting from 0.
-
-        Returns:
-            torch.Tensor: Output tensor of shape (b, 3, h, w).
-        '''
-
-        b = w.size(0)
-        x = self.constant.expand(b, -1, -1, -1).clone() # (b, c, h, w)
-
-        if steps == 0:
-            x = self.blocks[0](x, w) # (b, c, h, w)
-            x = self.to_rgb[0](x) # (b, c, h, w)
-            return x
-
-        for i in range(steps):
-            x = self.blocks[i](x, w) # (b, c, h/2, w/2)
-
-        old_rgb = self.to_rgb[steps - 1](x) # (b, 3, h/2, w/2)
-
-        x = self.blocks[steps](x, w) # (b, 3, h, w)
-        new_rgb = self.to_rgb[steps](x) # (b, 3, h, w)
-        old_rgb = F.interpolate(old_rgb, scale_factor=2, mode='bilinear', align_corners=False) # (b, 3, h, w)
-
-        x = (1 - alpha) * old_rgb + alpha * new_rgb # (b, 3, h, w)
-
-        return x
-
-class StyleGAN(nn.Module):
-    '''
-    StyleGAN implementation.
-
-    Args:
-        num_features (int): The number of features in the latent space vector.
-        resolution (int): The resolution of the image.
-        num_blocks (int, optional): The number of blocks in the synthesis network. Default value is 10.
-    '''
-
-    def __init__(self, num_features: int, resolution: int, num_blocks: int = 10):
-        super(StyleGAN, self).__init__()
-        self.num_features = num_features
-        self.resolution = resolution
-        self.num_blocks = num_blocks
-
-        self.mapping = Mapping(self.num_features, self.num_blocks)
-        self.synthesis = Synthesis(self.resolution, self.num_features)
-
-    def forward(self, x: torch.Tensor, alpha: float, steps: int) -> torch.Tensor:
-        '''
-        Args:
-            x (torch.Tensor): Random input tensor of shape (b, l).
-            alpha (float): Fade in alpha value.
-            steps (int): The number of steps starting from 0.
-
-        Returns:
-            torch.Tensor: Output tensor of shape (b, c, h, w).
-        '''
-
-        w = self.mapping(x) # (b, l)
-        outputs = self.synthesis(w, alpha, steps) # (b, c, h, w)
-
-        return outputs

stylegan_generator.py

@@ -1,22 +1,35 @@
-from stylegan import StyleGAN
-import torch
+import os
+import io
+import math
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+import gradio as gr
 
-LATENT_FEATURE = 512
-RESOLUTION = 128
+MODEL_NAME = 'batik_stylegan.onnx'
+CHANNELS = 3
+LATENT_FEATURES = 512
+RESOLUTION = 256
+LAST_INDEX = math.log2(RESOLUTION) - 2
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+MODEL_PATH = os.path.join("model", "")
 
-generator = StyleGAN(LATENT_FEATURES, RESOLUTION).to(device)
-last_checkpoint = torch.load(path, map_location=DEVIdeviceCE)
-generator.load_state_dict(last_checkpoint['generator'], strict=False)
-generator.eval()
+alpha = np.array([1.0], dtype=np.float32)
+steps = np.array([LAST_INDEX], dtype=np.int64)
 
-def generate_stylegan():
-    with torch.no_grad():
-        image = generator(torch.randn(1, LATENT_FEATURES, device=device), alpha=1.0, steps=5)
-        image = image.tanh()
-        image = (image + 1) / 2 
-        image = image.permute(1, 2, 0).cpu().numpy()
-        pil_img = Image.fromarray(image)
+def generate_image():
+    z = np.random.randn(1, LATENT_FEATURES).astype(np.float32)
 
-        return pil_img.resize((512, 512), Image.LANCZOS)
+    output = model.run(None, {
+        'z': z,
+        'alpha': alpha,
+        'steps': steps
+    })[0]
+
+    image = output.squeeze(0)
+    image = (image * 0.5 + 0.5) * 255
+    image = image.astype(np.uint8)
+    image = np.transpose(image, (1, 2, 0))
+    pil_image = Image.fromarray(image, 'RGB')
+    
+    return pil_img.resize((512, 512), Image.LANCZOS)