working version on pth weights

app.py

@@ -5,32 +5,24 @@ import sys
 import os
 from PIL import Image
 import torchvision.transforms as transforms
-from safetensors.torch import load_file
+
 
 photos_folder = "Photos"
 
-# Download model files
+# Download model and config
 repo_id = "Kiwinicki/sat2map-generator"
-model_path = hf_hub_download(repo_id=repo_id, filename="generator.safetensors")
-generator_code_path = hf_hub_download(repo_id=repo_id, filename="model.py")
+generator_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
+model_path = hf_hub_download(repo_id=repo_id, filename="model.py")
 
 # Add path to model
-sys.path.append(os.path.dirname(generator_code_path))
+sys.path.append(os.path.dirname(model_path))
 from model import Generator, GeneratorConfig
 
-# Initialize configuration
-cfg = GeneratorConfig(
-    channels=3,
-    num_features=64,
-    num_residuals=12,
-    depth=4
-)
-
 # Initialize model
+cfg = GeneratorConfig()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 generator = Generator(cfg).to(device)
-state_dict = load_file(model_path)
-generator.load_state_dict(state_dict)
+generator.load_state_dict(torch.load(generator_path))
 generator.eval()
 
 # Transformations
@@ -40,6 +32,7 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 ])
 
+
 def process_image(image):
     if image is None:
         return None
@@ -55,7 +48,6 @@ def process_image(image):
     output_image = output_tensor.squeeze(0).cpu()
     output_image = output_image * 0.5 + 0.5  # Denormalization
     output_image = transforms.ToPILImage()(output_image)
-    
     return output_image
 
 def load_images_from_folder(folder):

requirements.txt

@@ -3,5 +3,4 @@ torch>=2.0.0
 torchvision>=0.15.0
 gradio
 pillow
-pydantic==2.10.6
-safetensors
+pydantic==2.10.6