first commit

app.py

@@ -6,13 +6,12 @@ import matplotlib.pyplot as plt
 import requests
 import os
 
-# URLs dos arquivos do modelo
+# ================================
+# 1. Baixar pesos do Surya-1.0
+# ================================
 MODEL_URL = "https://huggingface.co/nasa-ibm-ai4science/Surya-1.0/resolve/main/surya.366m.v1.pt"
-
-# Nome local do arquivo
 MODEL_FILE = "surya.366m.v1.pt"
 
-# Função para baixar o modelo se não existir
 def download_model():
     if not os.path.exists(MODEL_FILE):
         print("Baixando pesos do Surya-1.0...")
@@ -21,33 +20,52 @@ def download_model():
             f.write(r.content)
         print("Download concluído!")
 
-# Baixar modelo
 download_model()
 
-# Carregar modelo PyTorch
-model = torch.load(MODEL_FILE, map_location=torch.device('cpu'))
+# ================================
+# 2. Definir a arquitetura do Surya
+# ================================
+# Aqui você deve colar ou importar a classe SuryaModel do repo oficial
+# Exemplo genérico:
+import torch.nn as nn
+
+class SuryaModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(1, 1, kernel_size=3, padding=1)
+    def forward(self, x):
+        return self.conv(x)
+
+# ================================
+# 3. Criar instância e carregar pesos
+# ================================
+model = SuryaModel()
+state_dict = torch.load(MODEL_FILE, map_location=torch.device('cpu'))
+model.load_state_dict(state_dict)
 model.eval()
 
-# Função para gerar heatmap
+# ================================
+# 4. Função de inferência para heatmap
+# ================================
 def infer_solar_image_heatmap(img):
-    # Pré-processamento: grayscale, resize 224x224
     img = img.convert("L").resize((224, 224))
     img_tensor = torch.tensor(np.array(img), dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.0
 
     with torch.no_grad():
         outputs = model(img_tensor)
     
-    # Criar heatmap
     emb = outputs.squeeze().numpy()
     heatmap = emb - emb.min()
-    heatmap /= heatmap.max() + 1e-8  # normalização 0-1
+    heatmap /= heatmap.max() + 1e-8
 
     plt.imshow(heatmap, cmap='hot')
     plt.axis('off')
     plt.tight_layout()
     return plt.gcf()
 
-# Interface Gradio
+# ================================
+# 5. Interface Gradio
+# ================================
 interface = gr.Interface(
     fn=infer_solar_image_heatmap,
     inputs=gr.Image(type="pil"),