upload pth weights

README.md

@@ -7,42 +7,34 @@ To load and initialize the `Generator` (based on CycleGAN with better cycles) mo
 Ensure you have the necessary Python packages installed:
 
 ```bash
-pip install torch==2.5.1 torchvision==0.20.1 safetensors huggingface_hub
+pip install torch==2.5.1 torchvision==0.20.1 huggingface_hub
 ```
 
 ## 2. Download Model Files
 
-Retrieve the `pytorch_model.safetensors` and `model.py` files from the Hugging Face repository using the `huggingface_hub` library:
+Retrieve the `generator.pth` and `model.py` files from the Hugging Face repository using the `huggingface_hub` library:
 
 ```python
 from huggingface_hub import hf_hub_download
 
 repo_id = "Kiwinicki/sat2map-generator"
-model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.safetensors")
+model_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
 generator_code_path = hf_hub_download(repo_id=repo_id, filename="model.py")
 ```
 
 ## 3. Load the Model
 
-Import the `Generator` class and load the model weights from the safetensors file:
+Import the `Generator` class and load the model weights from the `.pth` 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)
+cfg = GeneratorConfig()
 generator = Generator(cfg)
-generator.load_state_dict(state_dict)
+generator.load_state_dict(torch.load('generator.pth'))
 generator.eval()
 
 # Test the model

generator.safetensors → generator.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:662ab517fdeb62d15349409d45f0df28870fcd9f1c3a71f5a4efb8d8d830f144
-size 59680580
+oid sha256:323795e775e9a92e6f3d44cb87c3609911cee6201f586519c2f0f6fdb9361841
+size 59701794

model.py

@@ -135,14 +135,10 @@ class ResidualBlock(nn.Module):
 
 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.load_state_dict(torch.load('generator.pth'))
     generator.eval()
 
-    x = torch.randn([1, cfg.channels, 256, 256])
-    out = generator(x)
+    out = generator(torch.randn([1, cfg.channels, 256, 256]))
     print(out.shape)