Update app.py

app.py

@@ -11,7 +11,7 @@ import sunpy.net.attrs as a
 from sunpy.net import Fido
 from astropy.wcs import WCS
 import astropy.units as u
-from reproject import reproject_interp # Correct import statement
+from reproject import reproject_interp
 import os
 import warnings
 import logging
@@ -21,7 +21,9 @@ import sunpy.visualization.colormaps as sunpy_cm
 
 # --- Use the official Surya modules ---
 from surya.models.helio_spectformer import HelioSpectFormer
-from surya.utils.data import build_scalers, inverse_transform_single_channel
+from surya.utils.data import build_scalers
+# *** FIX: Corrected import location for the inverse transform function ***
+from surya.datasets.helio import inverse_transform_single_channel
 
 # --- Configuration ---
 warnings.filterwarnings("ignore", category=UserWarning, module='sunpy')
@@ -106,20 +108,18 @@ def fetch_and_process_sdo_data(target_dt, progress):
         for i, (channel, (physobs, sample)) in enumerate(SDO_CHANNELS_MAP.items()):
             progress(downloads_done / total_downloads, desc=f"Downloading {channel} for {t.strftime('%H:%M')}...")
             
-            instrument = a.Instrument.hmi if "hmi" in channel else a.Instrument.aia
             if channel in ["hmi_by", "hmi_bz"]: 
                 if data_maps[t].get("hmi_bx"): data_maps[t][channel] = data_maps[t]["hmi_bx"]
                 continue
 
             time_attr = a.Time(t - datetime.timedelta(minutes=10), t + datetime.timedelta(minutes=10))
-            search_query = [time_attr, physobs, sample]
-            # AIA and HMI queries are slightly different
+            search_query_list = [time_attr, physobs, sample]
             if "aia" in channel:
-                search_query.append(a.Instrument.aia)
+                search_query_list.append(a.Instrument.aia)
             else:
-                search_query.append(a.Instrument.hmi)
+                search_query_list.append(a.Instrument.hmi)
 
-            query = Fido.search(*search_query)
+            query = Fido.search(*search_query_list)
             
             if not query: raise ValueError(f"No data found for {channel} at {t}")
             files = Fido.fetch(query[0, 0], path="./data/sdo_cache")
@@ -183,11 +183,10 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
     
     c_idx = SDO_CHANNELS.index(channel_name)
     
-    # Process Prediction
-    means, stds, epsilons, sl_scale_factors = APP_CACHE["scalers"][SDO_CHANNELS[0]].get_params()
+    means, stds, epsilons, sl_scale_factors = APP_CACHE["scalers"][SDO_CHANNELS[c_idx]].get_params()
     pred_slice = inverse_transform_single_channel(
         prediction_tensor[0, c_idx].numpy(),
-        mean=means[c_idx], std=stds[c_idx], epsilon=epsilons[c_idx], sl_scale_factor=sl_scale_factors[c_idx]
+        mean=means, std=stds, epsilon=epsilons, sl_scale_factor=sl_scale_factors
     )
     
     vmax = np.quantile(np.nan_to_num(target_map[channel_name].data), 0.995)
@@ -237,7 +236,6 @@ def update_visualization_controller(last_input_map, prediction_tensor, target_ma
         return None, None, None
     return generate_visualization(last_input_map, prediction_tensor, target_map, channel_name)
 
-
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     state_last_input = gr.State()
     state_prediction = gr.State()
@@ -249,7 +247,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         # ☀️ Surya: Live Forecast Demo ☀️
         ### Generate a real forecast for any recent date using NASA's Heliophysics Model.
         **Instructions:**
-        1. Pick a date and time (at least 1 hour in the past).
+        1. Pick a date and time (at least 3 hours in the past).
         2. Click 'Generate Forecast'. **This will be slow (5-15 minutes) as it downloads live data.**
         3. Once complete, select different channels to explore the multi-spectrum forecast.
         </div>