upload better cycles generator

README.md

@@ -1,43 +1,64 @@
-To load and initialize the `Generator` model from the repository, follow these steps:
+# Loading the Generator Model
 
-1. **Install Required Packages**: Ensure you have the necessary Python packages installed:
+To load and initialize the `Generator` (based on CycleGAN with better cycles) model from the repository, follow these steps:
 
-    ```python
-    pip install torch omegaconf huggingface_hub
-    ```
+## 1. Install Required Packages
 
-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:
+Ensure you have the necessary Python packages installed:
 
-    ```python
-    from huggingface_hub import hf_hub_download
+```bash
+pip install torch==2.5.1 torchvision==0.20.1 safetensors huggingface_hub
+```
 
-    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")
-    ```
+## 2. Download Model Files
 
-3. **Load the Model**: Incorporate the downloaded `model.py` to define the `Generator` class, then load the model's state dictionary and configuration:
+Retrieve the `pytorch_model.safetensors` and `model.py` files from the Hugging Face repository using the `huggingface_hub` library:
 
-    ```python
-    import torch
-    import json
-    from omegaconf import OmegaConf
-    import sys
-    from pathlib import Path
-    from model import Generator
+```python
+from huggingface_hub import hf_hub_download
 
-    # Load configuration
-    with open(config_path, "r") as f:
-        config_dict = json.load(f)
-    cfg = OmegaConf.create(config_dict)
+repo_id = "Kiwinicki/sat2map-generator"
+model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.safetensors")
+generator_code_path = hf_hub_download(repo_id=repo_id, filename="model.py")
+```
 
-    # 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)
-    ```
+## 3. Load the Model
 
-    Here, `generator` is the initialized model ready for inference.
+Import the `Generator` class and load the model weights from the safetensors file:
+
+```python
+import torch
+from safetensors.torch import load_file
+from model import Generator, GeneratorConfig
+
+# Initialize configuration with default values
+cfg = GeneratorConfig(
+    channels=3,
+    num_features=64,
+    num_residuals=12,
+    depth=4
+)
+
+# Load the generator model
+state_dict = load_file(model_path)
+generator = Generator(cfg)
+generator.load_state_dict(state_dict)
+generator.eval()
+
+# Test the model
+x = torch.randn([1, cfg.channels, 256, 256])
+out = generator(x)
+print(f"Output shape: {out.shape}")
+```
+
+## 4. Model Configuration
+
+The model uses the following default configuration:
+- **channels**: 3 (RGB images)
+- **num_features**: 64 (base number of features)
+- **num_residuals**: 12 (number of residual blocks)
+- **depth**: 4 (network depth)
+
+The `generator` is now ready for inference on satellite-to-map translation tasks. 
+
+Model trained by Andrii Norets from "Czarna Magia".

config.json

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

generator.pth

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

model.py

@@ -1,9 +1,17 @@
 from torch import tanh, Tensor
 import torch.nn as nn
-from omegaconf import DictConfig
+from dataclasses import dataclass
 from abc import ABC, abstractmethod
 
 
+@dataclass
+class GeneratorConfig:
+    channels: int = 3
+    num_features: int = 64
+    num_residuals: int = 12
+    depth: int = 4
+
+
 class BaseGenerator(ABC, nn.Module):
     def __init__(self, channels: int = 3):
         super().__init__()
@@ -15,7 +23,7 @@ class BaseGenerator(ABC, nn.Module):
 
 
 class Generator(BaseGenerator):
-    def __init__(self, cfg: DictConfig):
+    def __init__(self, cfg: GeneratorConfig):
         super().__init__(cfg.channels)
         self.cfg = cfg
         self.model = self._construct_model()
@@ -124,3 +132,17 @@ class ResidualBlock(nn.Module):
 
     def forward(self, x: Tensor) -> Tensor:
         return x + self.block(x)
+
+if __name__ == '__main__':
+    import torch
+    from safetensors.torch import load_file
+
+    cfg = GeneratorConfig()
+    state_dict = load_file('generator.safetensors')
+    generator = Generator(cfg)    
+    generator.load_state_dict(state_dict)
+    generator.eval()
+
+    x = torch.randn([1, cfg.channels, 256, 256])
+    out = generator(x)
+    print(out.shape)

model.safetensors

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

requirements.txt

@@ -1,25 +0,0 @@
-antlr4-python3-runtime==4.9.3
-certifi==2024.12.14
-charset-normalizer==3.4.1
-colorama==0.4.6
-filelock==3.17.0
-fsspec==2024.12.0
-huggingface-hub==0.28.0
-idna==3.10
-Jinja2==3.1.5
-MarkupSafe==3.0.2
-mpmath==1.3.0
-networkx==3.4.2
-numpy==2.2.2
-omegaconf==2.3.0
-packaging==24.2
-pillow==11.1.0
-PyYAML==6.0.2
-requests==2.32.3
-setuptools==75.8.0
-sympy==1.13.1
-torch==2.5.1
-torchvision==0.20.1
-tqdm==4.67.1
-typing_extensions==4.12.2
-urllib3==2.3.0