Update app.py

app.py

@@ -4,11 +4,63 @@ import torch
 from torchvision import transforms
 import os
 import torchvision.transforms.functional as TF
-from model.anime_gan import Generator
+
 
 st.set_page_config(page_title="Ghibli Style Converter", layout="centered")
 
 @st.cache_resource
+import torch.nn as nn
+
+class ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(ConvLayer, self).__init__()
+        reflection_padding = kernel_size // 2
+        self.layer = nn.Sequential(
+            nn.ReflectionPad2d(reflection_padding),
+            nn.Conv2d(in_channels, out_channels, kernel_size, stride),
+            nn.InstanceNorm2d(out_channels, affine=True),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.block = nn.Sequential(
+            ConvLayer(channels, channels, 3, 1),
+            ConvLayer(channels, channels, 3, 1)
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+class Generator(nn.Module):
+    def __init__(self):
+        super(Generator, self).__init__()
+        self.encoder = nn.Sequential(
+            ConvLayer(3, 32, 7, 1),
+            ConvLayer(32, 64, 3, 2),
+            ConvLayer(64, 128, 3, 2),
+        )
+        self.res_blocks = nn.Sequential(*[ResidualBlock(128) for _ in range(5)])
+        self.decoder = nn.Sequential(
+            nn.Upsample(scale_factor=2),
+            ConvLayer(128, 64, 3, 1),
+            nn.Upsample(scale_factor=2),
+            ConvLayer(64, 32, 3, 1),
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(32, 3, 7, 1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.res_blocks(x)
+        x = self.decoder(x)
+        return x
+
 def load_model():
     model = Generator()
     model.load_state_dict(torch.load("model/miyazaki_hayao.pth", map_location="cpu"))