app.py

import gradio as gr
import torch
from torch import nn
from einops import rearrange
from PIL import Image
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"

# 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):
        print("Baixando pesos do Surya-1.0...")
        r = requests.get(MODEL_URL)
        with open(MODEL_FILE, "wb") as f:
            f.write(r.content)
        print("Download concluído!")

download_model()

# ================================
# 2. Colar aqui a classe HelioSpectFormer
# ================================
# Copie todo o conteúdo que você me enviou da HelioSpectFormer aqui
# ⚠️ Substitua a seção abaixo pelo código real do repo
from surya.models.helio_spectformer import HelioSpectFormer
 # se você tiver a pasta surya local

# ================================
# 3. Instanciar o modelo com parâmetros padrão
# ================================
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
)

# 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()

# ================================
# 4. Função de inferência para heatmap
# ================================
def infer_solar_image_heatmap(img):
    # 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))}

    # Inferência
    with torch.no_grad():
        outputs = model(batch)

    # Pegar o embedding da saída
    emb = outputs.squeeze().numpy()
    heatmap = emb - emb.min()
    heatmap /= heatmap.max() + 1e-8

    # Criar figura do heatmap
    plt.imshow(heatmap, cmap='hot')
    plt.axis('off')
    plt.tight_layout()
    return plt.gcf()

# ================================
# 5. 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",
    description="Upload de imagem solar → visualize heatmap gerado pelo Surya-1.0"
)

interface.launch()