first commit

app.py

@@ -1,76 +1,108 @@
-import gradio as gr
+import os
+import requests
+import importlib.util
+import sys
 import torch
 from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
-import requests
-import os
+import gradio as gr
 
 # ================================
-# 1. Baixar pesos do Surya-1.0
+# 1. Função para baixar arquivos
 # ================================
-MODEL_URL = "https://huggingface.co/nasa-ibm-ai4science/Surya-1.0/resolve/main/surya.366m.v1.pt"
-MODEL_FILE = "surya.366m.v1.pt"
-
-def download_model():
-    if not os.path.exists(MODEL_FILE):
-        print("Baixando pesos do Surya-1.0...")
-        r = requests.get(MODEL_URL)
-        with open(MODEL_FILE, "wb") as f:
+def download_file(url, filename):
+    if not os.path.exists(filename):
+        print(f"Baixando {filename}...")
+        r = requests.get(url)
+        with open(filename, "wb") as f:
             f.write(r.content)
-        print("Download concluído!")
+        print(f"{filename} baixado!")
+
+# ================================
+# 2. Baixar arquivos do Surya
+# ================================
+files = {
+    "helio_spectformer.py": "https://raw.githubusercontent.com/NASA-IMPACT/Surya/main/surya/helio_spectformer.py",
+    "spectformer.py": "https://raw.githubusercontent.com/NASA-IMPACT/Surya/main/surya/spectformer.py",
+    "embedding.py": "https://raw.githubusercontent.com/NASA-IMPACT/Surya/main/surya/embedding.py",
+    "flow.py": "https://raw.githubusercontent.com/NASA-IMPACT/Surya/main/surya/flow.py"
+}
+
+for fname, url in files.items():
+    download_file(url, fname)
 
-download_model()
+# ================================
+# 3. Baixar pesos do Surya
+# ================================
+MODEL_FILE = "surya.366m.v1.pt"
+MODEL_URL = "https://huggingface.co/nasa-ibm-ai4science/Surya-1.0/resolve/main/surya.366m.v1.pt"
+download_file(MODEL_URL, MODEL_FILE)
 
 # ================================
-# 2. Definir a arquitetura do Surya
+# 4. Importar dinamicamente as classes
 # ================================
-# Aqui você deve colar ou importar a classe SuryaModel do repo oficial
-# Exemplo genérico:
-import torch.nn as nn
+spec = importlib.util.spec_from_file_location("helio_spectformer", "helio_spectformer.py")
+helio_module = importlib.util.module_from_spec(spec)
+sys.modules["helio_spectformer"] = helio_module
+spec.loader.exec_module(helio_module)
 
-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)
+HelioSpectFormer = helio_module.HelioSpectFormer
 
 # ================================
-# 3. Criar instância e carregar pesos
+# 5. Instanciar o modelo
 # ================================
-model = SuryaModel()
-state_dict = torch.load(MODEL_FILE, map_location=torch.device('cpu'))
+model = HelioSpectFormer(
+    img_size=224,
+    patch_size=16,
+    in_chans=1,
+    embed_dim=366,
+    time_embedding={"type": "linear", "time_dim": 1},
+    depth=8,
+    n_spectral_blocks=4,
+    num_heads=8,
+    mlp_ratio=4.0,
+    drop_rate=0.0,
+    window_size=7,
+    dp_rank=1,
+    learned_flow=False,
+    finetune=True
+)
+
+state_dict = torch.load(MODEL_FILE, map_location=torch.device("cpu"))
 model.load_state_dict(state_dict)
 model.eval()
 
 # ================================
-# 4. Função de inferência para heatmap
+# 6. Função de inferência
 # ================================
 def infer_solar_image_heatmap(img):
-    img = img.convert("L").resize((224, 224))
-    img_tensor = torch.tensor(np.array(img), dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 255.0
+    # Pré-processamento da imagem
+    img_gray = img.convert("L").resize((224,224))
+    ts_tensor = torch.tensor(np.array(img_gray), dtype=torch.float32).unsqueeze(0).unsqueeze(0).unsqueeze(2) / 255.0
+    batch = {"ts": ts_tensor, "time_delta_input": torch.zeros((1,1))}
 
     with torch.no_grad():
-        outputs = model(img_tensor)
-    
+        outputs = model(batch)
+
+    # Criar heatmap da saída
     emb = outputs.squeeze().numpy()
     heatmap = emb - emb.min()
     heatmap /= heatmap.max() + 1e-8
 
-    plt.imshow(heatmap, cmap='hot')
-    plt.axis('off')
+    plt.imshow(heatmap, cmap="hot")
+    plt.axis("off")
     plt.tight_layout()
     return plt.gcf()
 
 # ================================
-# 5. Interface Gradio
+# 7. Interface Gradio
 # ================================
 interface = gr.Interface(
     fn=infer_solar_image_heatmap,
     inputs=gr.Image(type="pil"),
     outputs=gr.Plot(label="Heatmap do embedding Surya"),
-    title="Playground Surya-1.0 com Heatmap",
+    title="Playground Surya-1.0",
     description="Upload de imagem solar → visualize heatmap gerado pelo Surya-1.0"
 )
 

requirements.txt

@@ -1,6 +1,7 @@
 torch
+einops
 pillow
 numpy
 matplotlib
 gradio
-requests
+requests