Implement saved output loading for examples

app.py

@@ -12,6 +12,7 @@ from src.smc.inference import (
     SMCGradInferenceConfig,
     FTInferenceConfig,
 )
+from run_examples import get_out_if_exists
 
 GALLERY_HEIGHT = "224px"
 
@@ -46,6 +47,49 @@ def _format_inference_output(out) -> str:
     except Exception:
         return "Could not parse inference output"
 
+def try_load_saved_outputs(prompt):
+    """
+    Check for saved outputs for the given prompt for each method and return
+    (pretrained_gallery, pretrained_info, smc_gallery, smc_info, ft_gallery, ft_info).
+
+    If no saved output exists for a method, returns an empty gallery and
+    \"No saved output\" for info for that method.
+    """
+    try:
+        # Pretrained
+        pre_cfg = PretrainedInferenceConfig(prompt=prompt)
+        pre_out = get_out_if_exists("pretrained", pre_cfg)
+        if pre_out is not None:
+            pre_gallery = pre_out.images
+            pre_info = _format_inference_output(pre_out)
+        else:
+            pre_gallery, pre_info = [], "No saved output"
+
+        # SMC-grad
+        smc_cfg = SMCGradInferenceConfig(prompt=prompt)
+        smc_out = get_out_if_exists("smc_grad", smc_cfg)
+        if smc_out is not None:
+            smc_gallery = smc_out.images
+            smc_info = _format_inference_output(smc_out)
+        else:
+            smc_gallery, smc_info = [], "No saved output"
+
+        # FT
+        ft_cfg = FTInferenceConfig(prompt=prompt)
+        ft_out = get_out_if_exists("ft", ft_cfg)
+        if ft_out is not None:
+            ft_gallery = ft_out.images
+            ft_info = _format_inference_output(ft_out)
+        else:
+            ft_gallery, ft_info = [], "No saved output"
+
+        return pre_gallery, pre_info, smc_gallery, smc_info, ft_gallery, ft_info
+
+    except Exception as e:
+        # Don't crash the UI; print the traceback and return empty placeholders
+        traceback.print_exc()
+        return [], "Error checking saved outputs", [], "Error checking saved outputs", [], "Error checking saved outputs"
+
 
 # --- Per-method runner functions ---
 def run_pretrained_ui(prompt, pretrained_negative_prompt, pretrained_CFG, pretrained_steps):
@@ -150,7 +194,7 @@ with gr.Blocks() as demo:
         prompt = gr.Textbox(label="Prompt", placeholder="Enter prompt here", value=examples[0], lines=1)
         run_button = gr.Button("Run", variant="primary")
 
-    gr.Examples(examples=examples, inputs=prompt)
+    examples_widget = gr.Examples(examples=examples, inputs=prompt)
 
     # --- Pretrained method row ---
     with gr.Row():
@@ -308,6 +352,24 @@ with gr.Blocks() as demo:
         inputs=[prompt, ft_negative_prompt, ft_CFG, ft_steps],
         outputs=[ft_gallery, ft_info],
     )
+    
+    # Trigger when an example is selected
+    examples_widget.load_input_event.then(
+        fn=try_load_saved_outputs,
+        inputs=[prompt],
+        outputs=[
+            pretrained_gallery, pretrained_info,
+            smc_grad_gallery, smc_grad_info,
+            ft_gallery, ft_info,
+        ],
+    )
+    
+    # Trigger once on page load for the initial prompt value (so example[0] loads on startup)
+    demo.load(
+        fn=try_load_saved_outputs,
+        inputs=[prompt],
+        outputs=[pretrained_gallery, pretrained_info, smc_grad_gallery, smc_grad_info, ft_gallery, ft_info],
+    )
 
 # Enable Gradio queue to allow parallel execution of multiple handlers. Set concurrency
 # to 2 (one per method) — increase if you add more methods.