first commit

app.py

@@ -7,12 +7,37 @@ import numpy as np
 import matplotlib.pyplot as plt
 import requests
 import os
+import sys
+import warnings
+
+# Silenciar aviso depreciação do timm visto no HF Spaces
+warnings.filterwarnings(
+    "ignore",
+    message="Importing from timm.models.layers is deprecated, please import via timm.layers",
+    category=FutureWarning,
+)
+
+# Garantir import local do pacote `surya` mesmo se CWD for diferente
+sys.path.append(os.path.dirname(__file__))
 
 # ================================
 # 1. Baixar pesos do Surya-1.0
 # ================================
 MODEL_URL = "https://huggingface.co/nasa-ibm-ai4science/Surya-1.0/resolve/main/surya.366m.v1.pt"
-MODEL_FILE = "surya.366m.v1.pt"
+
+# Preferir checkpoint local se existir
+MODEL_CANDIDATES = [
+    os.path.join(os.path.dirname(__file__), "surya_model.pt"),
+    os.path.join(os.path.dirname(__file__), "surya.366m.v1.pt"),
+]
+
+def _pick_model_file():
+    for p in MODEL_CANDIDATES:
+        if os.path.exists(p):
+            return p
+    return MODEL_CANDIDATES[-1]
+
+MODEL_FILE = _pick_model_file()
 
 def download_model():
     if not os.path.exists(MODEL_FILE):
@@ -52,9 +77,33 @@ model = HelioSpectFormer(
     finetune=True
 )
 
-# Carregar pesos
-state_dict = torch.load(MODEL_FILE, map_location=torch.device('cpu'))
-model.load_state_dict(state_dict, strict=False)
+# Carregar pesos de forma resiliente (strict=False) e logar diferenças
+def _try_load_weights(m: nn.Module, path: str) -> None:
+    if os.environ.get("NO_WEIGHTS", "").lower() in {"1", "true", "yes"}:
+        print("NO_WEIGHTS=1 -> pulando carregamento de pesos")
+        return
+    try:
+        raw_sd = torch.load(path, map_location=torch.device('cpu'))
+        model_sd = m.state_dict()
+        filtered = {}
+        dropped = []
+        for k, v in raw_sd.items():
+            if k in model_sd and model_sd[k].shape == v.shape:
+                filtered[k] = v
+            else:
+                dropped.append((k, tuple(v.shape) if hasattr(v, 'shape') else None, tuple(model_sd.get(k, torch.tensor(())).shape) if k in model_sd else None))
+        missing, unexpected = m.load_state_dict(filtered, strict=False)
+        print(f"Pesos carregados parcialmente. Ok={len(filtered)} Missing={len(missing)} Unexpected={len(unexpected)} Dropped={len(dropped)}")
+        if dropped:
+            print("Algumas chaves foram descartadas por mismatch (ex.:)", dropped[:5])
+        if missing:
+            print("Exemplos de missing:", missing[:10])
+        if unexpected:
+            print("Exemplos de unexpected:", unexpected[:10])
+    except Exception as e:
+        print(f"Falha ao carregar pesos de {path}: {e}")
+
+_try_load_weights(model, MODEL_FILE)
 model.eval()
 
 # ================================