Update app.py

app.py

@@ -3,6 +3,7 @@ import tempfile
 import os
 import shutil
 from transformers import AutoModelForImageSegmentation, AutoProcessor
+from peft import PeftModel  # For LoRA adapter
 from sunpy.net import Fido, attrs as a
 import astropy.units as u
 import numpy as np
@@ -26,11 +27,19 @@ def cleanup_temp():
             shutil.rmtree(dir_path, ignore_errors=True)
 cleanup_temp()  # Run once at start
 
-# Surya model setup (loads into /tmp cache)
-MODEL_ID = "nasa-ibm-ai4science/ar_segmentation_surya"
+# Surya model setup (base + LoRA adapter; now works after repo install)
+BASE_MODEL_ID = "nasa-ibm-ai4science/Surya-1.0"
+ADAPTER_MODEL_ID = "nasa-ibm-ai4science/ar_segmentation_surya"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = AutoModelForImageSegmentation.from_pretrained(MODEL_ID).to(device)
-processor = AutoProcessor.from_pretrained(MODEL_ID)
+
+# Load base model (custom "surya" type now registered via repo install)
+base_model = AutoModelForImageSegmentation.from_pretrained(BASE_MODEL_ID).to(device)
+
+# Load LoRA adapter for AR segmentation
+model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID).to(device)
+
+# Load processor (from base; assumes standard image processor)
+processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)
 
 # Historical observed values for May 2024 Gannon Storm
 HISTORICAL_DENSITY_INCREASE = 6.0  # Up to 6x at 400 km