Update app.py

app.py

@@ -1,7 +1,154 @@
 import gradio as gr
+import torch.nn as nn
+from torch import tanh, Tensor
+from abc import ABC, abstractmethod
+from huggingface_hub import hf_hub_download
+import torch
+import json
+from omegaconf import OmegaConf
+from model import Generator
 
-def greet(name):
-    return f"Hello {name}!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+class BaseGenerator(ABC, nn.Module):
+    def __init__(self, channels: int = 3):
+        super().__init__()
+        self.channels = channels
+
+    @abstractmethod
+    def forward(self, x: Tensor) -> Tensor:
+        pass
+
+
+class Generator(BaseGenerator):
+    def __init__(self, cfg: DictConfig):
+        super().__init__(cfg.channels)
+        self.cfg = cfg
+        self.model = self._construct_model()
+
+    def _construct_model(self):
+        initial_layer = nn.Sequential(
+            nn.Conv2d(
+                self.cfg.channels,
+                self.cfg.num_features,
+                kernel_size=7,
+                stride=1,
+                padding=3,
+                padding_mode="reflect",
+            ),
+            nn.ReLU(inplace=True),
+        )
+
+        down_blocks = nn.Sequential(
+            ConvBlock(
+                self.cfg.num_features,
+                self.cfg.num_features * 2,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+            ConvBlock(
+                self.cfg.num_features * 2,
+                self.cfg.num_features * 4,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+            ),
+        )
+
+        residual_blocks = nn.Sequential(
+            *[
+                ResidualBlock(self.cfg.num_features * 4)
+                for _ in range(self.cfg.num_residuals)
+            ]
+        )
+
+        up_blocks = nn.Sequential(
+            ConvBlock(
+                self.cfg.num_features * 4,
+                self.cfg.num_features * 2,
+                down=False,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                output_padding=1,
+            ),
+            ConvBlock(
+                self.cfg.num_features * 2,
+                self.cfg.num_features,
+                down=False,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                output_padding=1,
+            ),
+        )
+
+        last_layer = nn.Conv2d(
+            self.cfg.num_features,
+            self.cfg.channels,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+            padding_mode="reflect",
+        )
+
+        return nn.Sequential(
+            initial_layer, down_blocks, residual_blocks, up_blocks, last_layer
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return tanh(self.model(x))
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, down=True, use_activation=True, **kwargs
+    ):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, padding_mode="reflect", **kwargs)
+            if down
+            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),
+            nn.InstanceNorm2d(out_channels),
+            nn.ReLU(inplace=True) if use_activation else nn.Identity(),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(x)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            ConvBlock(channels, channels, kernel_size=3, padding=1),
+            ConvBlock(
+                channels, channels, use_activation=False, kernel_size=3, padding=1
+            ),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.block(x)
+
+
+repo_id = "Kiwinicki/sat2map-generator"
+generator_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
+config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+model_path = hf_hub_download(repo_id=repo_id, filename="model.py")
+
+
+with open(config_path, "r") as f:
+    config_dict = json.load(f)
+cfg = OmegaConf.create(config_dict)
+
+generator = Generator(cfg)
+generator.load_state_dict(torch.load(generator_path))
+generator.eval()
+
+
+
+def greet(iamge):
+    return image
+
+iface = gr.Interface(fn=greet, inputs="image", outputs="image")
 iface.launch()