Update app.py

app.py

@@ -1,35 +1,35 @@
 import gradio as gr
-import torch.nn as nn
-from torch import tanh, Tensor
-from abc import ABC, abstractmethod
-from huggingface_hub import hf_hub_download
 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
 
+# Pobierz model i config
 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")
 
-import sys, os
+# Dodaj ścieżkę do modelu
 sys.path.append(os.path.dirname(model_path))
 from model import Generator
 
-
-
+# Załaduj konfigurację
 with open(config_path, "r") as f:
     config_dict = json.load(f)
 cfg = OmegaConf.create(config_dict)
 
-generator = Generator(cfg)
-generator.load_state_dict(torch.load(generator_path, map_location=torch.device('cpu')))
+# Inicjalizacja modelu
+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))
 generator.eval()
 
-from PIL import Image
-import torchvision.transforms as transforms
-
-
+# Transformacje
 transform = transforms.Compose([
     transforms.Resize((256, 256)),
     transforms.ToTensor(),
@@ -37,12 +37,25 @@ transform = transforms.Compose([
 ])
 
 def process_image(image):
-    image_tensor = transform(image).unsqueeze(0)
+    # Konwersja do tensora
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Inferencja
     with torch.no_grad():
         output_tensor = generator(image_tensor)
-    output_image = output_tensor.squeeze(0)
+    
+    # Przygotowanie wyjścia
+    output_image = output_tensor.squeeze(0).cpu()
+    output_image = output_image * 0.5 + 0.5  # Denormalizacja
     output_image = transforms.ToPILImage()(output_image)
+    
     return output_image
 
-iface = gr.Interface(fn=process_image, inputs="image", outputs="image", title="Image Generator")
-iface.launch()
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil"),
+    outputs="image",
+    title="Satellite to Map Generator"
+)
+
+iface.launch()