Update app.py

app.py

@@ -1,4 +1,4 @@
-# Save this file as app.py in the root of the cloned Surya repository
+# Save this file as in the root of the cloned Surya repository
 
 import gradio as gr
 import torch
@@ -155,12 +155,10 @@ def fetch_and_process_sdo_data(target_dt):
     target_map = data_maps[target_time]
     last_input_map = data_maps[input_times[-1]]
 
-    # 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):
     model = APP_CACHE["model"]
     device = APP_CACHE["device"]
@@ -195,9 +193,6 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
 
 # --- 4. Gradio UI and Controllers ---
 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),
@@ -208,45 +203,34 @@ def forecast_controller(dt_str):
     try:
         if not dt_str: raise gr.Error("Please select a date and time.")
         
-        # --- 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)
         
-        # --- 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.")
 
-        # --- 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")
         
         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,
@@ -258,36 +242,49 @@ def forecast_controller(dt_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)
         }
 
-
+# --- 5. Gradio UI Definition ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     state_last_input = gr.State()
     state_prediction = gr.State()
     state_target = gr.State()
 
-    gr.Markdown(...) # UI definition is the same
+    # *** FIX: Replaced all '...' with complete UI component definitions ***
+    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>
+        """
+    )
 
     with gr.Row():
-        datetime_input = gr.Textbox(...)
+        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")
+        )
         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)
+    log_box = gr.Textbox(label="Log", interactive=False, visible=False, lines=5, max_lines=10)
 
     with gr.Group(visible=False) as results_group:
-        channel_selector = gr.Dropdown(...)
+        channel_selector = gr.Dropdown(
+            choices=SDO_CHANNELS, value="aia171", label="🛰️ Select SDO Channel to Visualize"
+        )
         with gr.Row():
-            input_display = gr.Image(...)
-            prediction_display = gr.Image(...)
-            target_display = gr.Image(...)
+            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)
 
-    # 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],
@@ -299,7 +296,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
     
     channel_selector.change(
-        fn=generate_visualization, # This is a fast function, no generator needed
+        fn=generate_visualization,
         inputs=[state_last_input, state_prediction, state_target, channel_selector],
         outputs=[input_display, prediction_display, target_display]
     )