Update app.py

app.py

@@ -2,8 +2,9 @@ import torch
 import tempfile
 import os
 import shutil
-import surya  # Import to register custom 'surya' model_type with transformers
-from transformers import AutoModelForImageSegmentation, AutoProcessor
+import surya  # Import to register custom modules
+from surya.models.surya_model import SuryaModel  # Custom model class
+from surya.processors import SuryaProcessor  # Custom processor
 from peft import PeftModel  # For LoRA adapter
 from sunpy.net import Fido, attrs as a
 import astropy.units as u
@@ -28,19 +29,19 @@ def cleanup_temp():
             shutil.rmtree(dir_path, ignore_errors=True)
 cleanup_temp()  # Run once at start
 
-# Surya model setup (base + LoRA adapter; registration via 'import surya')
+# Surya model setup (custom class + LoRA adapter)
 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")
 
-# Load base model (now recognized after surya import)
-base_model = AutoModelForImageSegmentation.from_pretrained(BASE_MODEL_ID).to(device)
+# Load base model using custom SuryaModel class
+base_model = SuryaModel.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)
+# Load custom processor
+processor = SuryaProcessor.from_pretrained(BASE_MODEL_ID)
 
 # Historical observed values for May 2024 Gannon Storm
 HISTORICAL_DENSITY_INCREASE = 6.0  # Up to 6x at 400 km