Update app.py

app.py

@@ -16,7 +16,8 @@ from io import BytesIO
 import re
 
 from surya.models.helio_spectformer import HelioSpectFormer
-from surya.utils.data import build_scalers
+# *** FIX: Import ScalerCollection and correct the import for inverse_transform_single_channel ***
+from surya.utils.data import build_scalers, ScalerCollection
 from surya.datasets.helio import inverse_transform_single_channel
 
 warnings.filterwarnings("ignore", category=UserWarning, module='sunpy')
@@ -48,7 +49,10 @@ def setup_and_load_model():
         config = yaml.safe_load(fp)
     APP_CACHE["config"] = config
     scalers_info = yaml.safe_load(open("data/Surya-1.0/scalers.yaml", "r"))
-    APP_CACHE["scalers"] = build_scalers(info=scalers_info)
+    
+    # *** FIX: Create the ScalerCollection object from the dictionary returned by build_scalers ***
+    scalers_dict = build_scalers(info=scalers_info)
+    APP_CACHE["scalers"] = ScalerCollection(scalers_dict, config["data"]["sdo_channels"])
     
     yield "Initializing model architecture..."
     model_config = config["model"]
@@ -79,7 +83,7 @@ def find_nearest_browse_image_url(channel, target_dt):
     url_code = CHANNEL_TO_URL_CODE[channel]
     base_url = "https://sdo.gsfc.nasa.gov/assets/img/browse"
     
-    for i in range(2): # Try today, then yesterday
+    for i in range(2):
         dt_to_try = target_dt - datetime.timedelta(days=i)
         dir_url = dt_to_try.strftime(f"{base_url}/%Y/%m/%d/")
         
@@ -150,7 +154,7 @@ def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
             img_resized = img.resize((img_size, img_size), Image.Resampling.LANCZOS)
             norm_data = np.array(img_resized, dtype=np.float32)
 
-            scaled_data = scaler.transform(norm_data.reshape(-1, 1), c_idx=i).reshape(norm_data.shape)
+            scaled_data = scaler.transform(norm_data, c_idx=i)
             channel_tensors.append(torch.from_numpy(scaled_data.astype(np.float32)))
         processed_tensors[t] = torch.stack(channel_tensors)