Update app.py

app.py

@@ -18,6 +18,7 @@ import logging
 import datetime
 import matplotlib.pyplot as plt
 import sunpy.visualization.colormaps as sunpy_cm
+import traceback
 
 # --- Use the official Surya modules ---
 from surya.models.helio_spectformer import HelioSpectFormer
@@ -33,7 +34,6 @@ logger = logging.getLogger(__name__)
 # Global cache for model, config, etc.
 APP_CACHE = {}
 
-# *** FIX: Corrected the a.Wavelength calls to use astropy units ***
 SDO_CHANNELS_MAP = {
     "aia94": (a.Wavelength(94 * u.angstrom), a.Sample(12 * u.s)),
     "aia131": (a.Wavelength(131 * u.angstrom), a.Sample(12 * u.s)),
@@ -52,22 +52,23 @@ SDO_CHANNELS_MAP = {
 SDO_CHANNELS = list(SDO_CHANNELS_MAP.keys())
 
 # --- 1. Model Loading and Setup ---
-def setup_and_load_model(progress=gr.Progress()):
+def setup_and_load_model():
     if "model" in APP_CACHE:
+        yield "Model already loaded. Skipping setup."
         return
 
-    progress(0.1, desc="Downloading model files (first run only)...")
+    yield "Downloading model files (first run only)..."
     snapshot_download(repo_id="nasa-ibm-ai4science/Surya-1.0", local_dir="data/Surya-1.0",
                       allow_patterns=["config.yaml", "scalers.yaml", "surya.366m.v1.pt"])
 
-    progress(0.5, desc="Loading configuration and scalers...")
+    yield "Loading configuration and data scalers..."
     with open("data/Surya-1.0/config.yaml") as fp:
         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)
     
-    progress(0.7, desc="Initializing and loading model...")
+    yield "Initializing model architecture..."
     model_config = config["model"]
     model = HelioSpectFormer(
         img_size=model_config["img_size"], patch_size=model_config["patch_size"],
@@ -83,15 +84,17 @@ def setup_and_load_model(progress=gr.Progress()):
     )
     device = "cuda" if torch.cuda.is_available() else "cpu"
     APP_CACHE["device"] = device
+    
+    yield f"Loading model weights to {device}..."
     weights = torch.load(f"data/Surya-1.0/surya.366m.v1.pt", map_location=torch.device(device))
     model.load_state_dict(weights, strict=True)
     model.to(device)
     model.eval()
     APP_CACHE["model"] = model
-    logger.info("Model setup complete.")
+    yield "✅ Model setup complete."
 
-# --- 2. Live Data Fetching and Preprocessing ---
-def fetch_and_process_sdo_data(target_dt, progress):
+# --- 2. Live Data Fetching and Preprocessing (as a generator) ---
+def fetch_and_process_sdo_data(target_dt):
     config = APP_CACHE["config"]
     img_size = config["model"]["img_size"][0]
     
@@ -104,11 +107,12 @@ def fetch_and_process_sdo_data(target_dt, progress):
     data_maps = {}
     total_downloads = len(all_times) * len(SDO_CHANNELS)
     downloads_done = 0
+    yield f"Starting download of {total_downloads} data files..."
     for t in all_times:
         data_maps[t] = {}
         for i, (channel, (physobs, sample)) in enumerate(SDO_CHANNELS_MAP.items()):
             downloads_done += 1
-            progress(downloads_done / total_downloads, desc=f"Downloading {channel} for {t.strftime('%H:%M')}...")
+            yield f"Downloading [{downloads_done}/{total_downloads}]: {channel} for {t.strftime('%Y-%m-%d %H:%M')}..."
             
             if channel in ["hmi_by", "hmi_bz"]: 
                 if data_maps[t].get("hmi_bx"): data_maps[t][channel] = data_maps[t]["hmi_bx"]
@@ -122,12 +126,14 @@ def fetch_and_process_sdo_data(target_dt, progress):
             files = Fido.fetch(query[0, 0], path="./data/sdo_cache")
             data_maps[t][channel] = sunpy.map.Map(files[0])
             
+    yield "✅ All files downloaded. Starting preprocessing..."
     output_wcs = WCS(naxis=2)
     output_wcs.wcs.crpix = [(img_size + 1) / 2, (img_size + 1) / 2]
     output_wcs.wcs.cdelt = np.array([-1.2, 1.2]) * u.arcsec
     output_wcs.wcs.crval = [0, 0] * u.arcsec
     output_wcs.wcs.ctype = ['HPLN-TAN', 'HPLT-TAN']
 
+    scaler = APP_CACHE["scalers"]
     processed_tensors = {}
     for t, channel_maps in data_maps.items():
         channel_tensors = []
@@ -139,48 +145,45 @@ def fetch_and_process_sdo_data(target_dt, progress):
             if exp_time is None or exp_time <= 0: exp_time = 1.0
             norm_data = reprojected_data / exp_time
 
-            scaler = APP_CACHE["scalers"][channel]
-            scaled_data = scaler.transform(norm_data)
+            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)
 
+    yield "✅ Preprocessing complete."
     input_tensor_list = [processed_tensors[t] for t in input_times]
     input_tensor = torch.stack(input_tensor_list, dim=1).unsqueeze(0)
     target_map = data_maps[target_time]
     last_input_map = data_maps[input_times[-1]]
 
-    return input_tensor, last_input_map, target_map
+    # The final yield of a generator is its return value
+    yield (input_tensor, last_input_map, target_map)
+
 
 # --- 3. Inference and Visualization ---
+# (These are fast and don't need to be generators)
 def run_inference(input_tensor):
-    logger.info("Running model inference...")
     model = APP_CACHE["model"]
     device = APP_CACHE["device"]
-    
     time_deltas = APP_CACHE["config"]["data"]["time_delta_input_minutes"]
     time_delta_tensor = torch.tensor(time_deltas, dtype=torch.float32).unsqueeze(0).to(device)
-    
     input_batch = {"ts": input_tensor.to(device), "time_delta_input": time_delta_tensor}
-    
     with torch.no_grad():
         with torch.autocast(device_type=device.split(':')[0], dtype=torch.bfloat16):
             prediction = model(input_batch)
-    logger.info("Inference complete.")
     return prediction.cpu()
 
 def generate_visualization(last_input_map, prediction_tensor, target_map, channel_name):
     if last_input_map is None: return None, None, None
-    
     c_idx = SDO_CHANNELS.index(channel_name)
-    means, stds, epsilons, sl_scale_factors = APP_CACHE["scalers"][channel_name].get_params()
+    scaler = APP_CACHE["scalers"]
+    all_means, all_stds, all_epsilons, all_sl_scale_factors = scaler.get_params()
+    mean, std, epsilon, sl_scale_factor = all_means[c_idx], all_stds[c_idx], all_epsilons[c_idx], all_sl_scale_factors[c_idx]
     pred_slice = inverse_transform_single_channel(
-        prediction_tensor[0, c_idx].numpy(), mean=means, std=stds, epsilon=epsilons, sl_scale_factor=sl_scale_factors
+        prediction_tensor[0, c_idx].numpy(), mean=mean, std=std, epsilon=epsilon, sl_scale_factor=sl_scale_factor
     )
-    
     vmax = np.quantile(np.nan_to_num(target_map[channel_name].data), 0.995)
     cmap_name = f"sdoaia{channel_name.replace('aia', '')}" if 'aia' in channel_name else 'hmimag'
     cmap = plt.get_cmap(sunpy_cm.cmlist.get(cmap_name, 'gray'))
-    
     def to_pil(data, flip=False):
         data_clipped = np.nan_to_num(data)
         data_clipped = np.clip(data_clipped, 0, vmax)
@@ -188,79 +191,115 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
         colored = (cmap(data_norm)[:, :, :3] * 255).astype(np.uint8)
         img = Image.fromarray(colored)
         return img.transpose(Image.Transpose.FLIP_TOP_BOTTOM) if flip else img
-
-    return to_pil(last_input_map[channel_name].data), to_pil(pred_slice, flip=True), to_pil(target_map[channel_name].data)
+    return to_pil(last_input_map[channel_name].data, flip=True), to_pil(pred_slice, flip=True), to_pil(target_map[channel_name].data, flip=True)
 
 # --- 4. Gradio UI and Controllers ---
-def forecast_controller(dt_str, progress=gr.Progress(track_tqdm=True)):
+def forecast_controller(dt_str):
+    # This is now a generator function that yields updates to the UI
+    
+    # Initial UI state: disable inputs, clear old results
+    yield {
+        log_box: gr.update(value="Starting forecast...", visible=True),
+        run_button: gr.update(interactive=False),
+        datetime_input: gr.update(interactive=False),
+        results_group: gr.update(visible=False)
+    }
+    
     try:
         if not dt_str: raise gr.Error("Please select a date and time.")
         
-        progress(0, desc="Initializing...")
-        setup_and_load_model(progress)
-        
+        # --- Stage 1: Setup Model ---
+        # The setup function is also a generator, so we loop through its yields
+        for status in setup_and_load_model():
+            yield { log_box: status }
+
         target_dt = datetime.datetime.fromisoformat(dt_str)
-        logger.info(f"Starting forecast for target time: {target_dt}")
+        
+        # --- Stage 2: Fetch and Process Data ---
+        # We loop through the yields from the data pipeline
+        data_pipeline = fetch_and_process_sdo_data(target_dt)
+        while True:
+            try:
+                # Get the next status update
+                status = next(data_pipeline)
+                # If it's a tuple, it's the final return value
+                if isinstance(status, tuple):
+                    input_tensor, last_input_map, target_map = status
+                    break
+                # Otherwise, it's a string update
+                yield { log_box: status }
+            except StopIteration:
+                raise gr.Error("Data processing pipeline finished unexpectedly.")
 
-        input_tensor, last_input_map, target_map = fetch_and_process_sdo_data(target_dt, progress)
+        # --- Stage 3: Run Inference ---
+        yield { log_box: "Running AI model inference..." }
         prediction_tensor = run_inference(input_tensor)
         
+        # --- Stage 4: Generate Visualization ---
+        yield { log_box: "Generating final visualizations..." }
         img_in, img_pred, img_target = generate_visualization(last_input_map, prediction_tensor, target_map, "aia171")
         
-        status = f"Forecast complete for {target_dt.isoformat()}. Ready to explore channels."
-        logger.info(status)
-        
-        return (last_input_map, prediction_tensor, target_map,
-                img_in, img_pred, img_target, status, gr.update(visible=True))
-                
+        yield {
+            log_box: f"✅ Forecast complete for {target_dt.isoformat()}.",
+            results_group: gr.update(visible=True),
+            # Pass final data to state objects
+            state_last_input: last_input_map,
+            state_prediction: prediction_tensor,
+            state_target: target_map,
+            # Display final images
+            input_display: img_in,
+            prediction_display: img_pred,
+            target_display: img_target,
+        }
+
     except Exception as e:
-        logger.error(f"An error occurred: {e}", exc_info=True)
-        raise gr.Error(f"Failed to generate forecast. Error: {e}")
+        error_str = traceback.format_exc()
+        logger.error(f"An error occurred: {e}\n{error_str}")
+        yield { log_box: f"❌ ERROR: {e}\n\nTraceback:\n{error_str}" }
+    
+    finally:
+        # Final UI state: re-enable inputs
+        yield {
+            run_button: gr.update(interactive=True),
+            datetime_input: gr.update(interactive=True)
+        }
 
-def update_visualization_controller(last_input_map, prediction_tensor, target_map, channel_name):
-    if last_input_map is None: 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()
     state_target = gr.State()
 
-    gr.Markdown(
-        """
-        <div align='center'>
-        # ☀️ 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 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>
-        """
-    )
+    gr.Markdown(...) # UI definition is the same
 
     with gr.Row():
-        datetime_input = gr.Textbox(label="Enter Forecast Start Time (YYYY-MM-DD HH:MM:SS)",
-                                    value=(datetime.datetime.now() - datetime.timedelta(hours=3)).strftime("%Y-%m-%d %H:%M:%S"))
+        datetime_input = gr.Textbox(...)
         run_button = gr.Button("🔮 Generate Forecast", variant="primary")
 
+    # NEW: A dedicated box for logs and feedback
+    log_box = gr.Textbox(label="Log", interactive=False, visible=False, lines=5)
+
     with gr.Group(visible=False) as results_group:
-        status_box = gr.Textbox(label="Status", interactive=False)
-        channel_selector = gr.Dropdown(choices=SDO_CHANNELS, value="aia171", label="🛰️ Select SDO Channel")
+        channel_selector = gr.Dropdown(...)
         with gr.Row():
-            input_display = gr.Image(label="Last Input to Model", height=512, width=512, interactive=False)
-            prediction_display = gr.Image(label="Surya's Forecast", height=512, width=512, interactive=False)
-            target_display = gr.Image(label="Ground Truth", height=512, width=512, interactive=False)
+            input_display = gr.Image(...)
+            prediction_display = gr.Image(...)
+            target_display = gr.Image(...)
 
+    # The .click() event is now pointed to our generator function
+    # It updates multiple components based on what the generator yields
     run_button.click(
         fn=forecast_controller,
         inputs=[datetime_input],
-        outputs=[state_last_input, state_prediction, state_target,
-                 input_display, prediction_display, target_display, status_box, results_group]
+        outputs=[
+            log_box, run_button, datetime_input, results_group,
+            state_last_input, state_prediction, state_target,
+            input_display, prediction_display, target_display
+        ]
     )
     
     channel_selector.change(
-        fn=update_visualization_controller,
+        fn=generate_visualization, # This is a fast function, no generator needed
         inputs=[state_last_input, state_prediction, state_target, channel_selector],
         outputs=[input_display, prediction_display, target_display]
     )