Update app.py

app.py

@@ -1,34 +1,36 @@
 import gradio as gr
 import torch
 from huggingface_hub import hf_hub_download
-import json
-from omegaconf import OmegaConf
 import sys
 import os
 from PIL import Image
 import torchvision.transforms as transforms
+from safetensors.torch import load_file
 
 photos_folder = "Photos"
 
-# Download model and config
+# Download model files
 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")
+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")
 
 # Add path to model
-sys.path.append(os.path.dirname(model_path))
-from model import Generator
+sys.path.append(os.path.dirname(generator_code_path))
+from model import Generator, GeneratorConfig
 
-# Load configuration
-with open(config_path, "r") as f:
-    config_dict = json.load(f)
-cfg = OmegaConf.create(config_dict)
+# Initialize configuration
+cfg = GeneratorConfig(
+    channels=3,
+    num_features=64,
+    num_residuals=12,
+    depth=4
+)
 
 # Initialize model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 generator = Generator(cfg).to(device)
-generator.load_state_dict(torch.load(generator_path, map_location=device))
+state_dict = load_file(model_path)
+generator.load_state_dict(state_dict)
 generator.eval()
 
 # Transformations
@@ -86,12 +88,11 @@ def app():
                 gallery = gr.Gallery(
                     label="Image Gallery",
                     value=gallery_images,
-                    columns=3,  # Set number of columns directly in the constructor
+                    columns=3,
                     rows=2,
                     height="auto"
                 )
 
-
             with gr.Column():
                 output_image = gr.Image(label="Result Image", type="pil")