initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+.venv
+.idea

README.md

@@ -0,0 +1,43 @@
+To load and initialize the `Generator` model from the repository, follow these steps:
+
+1. **Install Required Packages**: Ensure you have the necessary Python packages installed:
+
+    ```python
+    pip install torch omegaconf huggingface_hub
+    ```
+
+2. **Download Model Files**: Retrieve the `generator.pth`, `config.json`, and `model.py` files from the Hugging Face repository. You can use the `huggingface_hub` library for this:
+
+    ```python
+    from huggingface_hub import hf_hub_download
+
+    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")
+    ```
+
+3. **Load the Model**: Incorporate the downloaded `model.py` to define the `Generator` class, then load the model's state dictionary and configuration:
+
+    ```python
+    import torch
+    import json
+    from omegaconf import OmegaConf
+    import sys
+    from pathlib import Path
+    from model import Generator
+
+    # Load configuration
+    with open(config_path, "r") as f:
+        config_dict = json.load(f)
+    cfg = OmegaConf.create(config_dict)
+
+    # Initialize and load the generator model
+    generator = Generator(cfg)
+    generator.load_state_dict(torch.load(generator_path))
+    generator.eval()
+    x = torch.randn([1, cfg['channels'], 256, 256])
+    out = generator(x)
+    ```
+
+    Here, `generator` is the initialized model ready for inference.

config.json

@@ -0,0 +1,6 @@
+{
+    "channels": 3,
+    "num_features": 64,
+    "num_residuals": 12,
+    "depth": 4
+}

generator.pth

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

model.py

@@ -0,0 +1,126 @@
+from torch import tanh, Tensor
+import torch.nn as nn
+from omegaconf import DictConfig
+from abc import ABC, abstractmethod
+
+
+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)

model.safetensors

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