preload

app.py

@@ -1,4 +1,7 @@
+# pip install -U gradio transformers pillow matplotlib
+
 import io
+from functools import lru_cache
 from typing import Optional
 
 import gradio as gr
@@ -19,15 +22,18 @@ def _fig_to_pil(fig) -> Image.Image:
     buf.seek(0)
     return Image.open(buf).convert("RGB")
 
+@lru_cache(maxsize=64)
+def get_viz(model_id: str) -> ImageVisualizer:
+    return ImageVisualizer(model_id)
+
 def _run(model_id: str, image: Optional[Image.Image], use_sample: bool, add_grid: bool):
-    viz = ImageVisualizer(model_id)
+    viz = get_viz(model_id)
 
     captured = []
     orig_show = plt.show
 
     def _capture_show(*_, **__):
-        fig = plt.gcf()
-        captured.append(fig)
+        captured.append(plt.gcf())
 
     try:
         plt.show = _capture_show
@@ -35,9 +41,10 @@ def _run(model_id: str, image: Optional[Image.Image], use_sample: bool, add_grid
     finally:
         plt.show = orig_show
 
-    imgs = [_fig_to_pil(fig) for fig in captured] if captured else []
+    left_img = _fig_to_pil(captured[0]) if len(captured) >= 1 else None
+    right_img = _fig_to_pil(captured[1]) if len(captured) >= 2 else None
     prompt_preview = viz.default_message(full_output=False)
-    return imgs, prompt_preview
+    return left_img, right_img, prompt_preview
 
 def _resolve_and_run(model_pick, custom_model, image, use_sample, add_grid):
     model_id = (custom_model or "").strip() or (model_pick or "").strip()
@@ -46,54 +53,49 @@ def _resolve_and_run(model_pick, custom_model, image, use_sample, add_grid):
     return _run(model_id, image, use_sample, add_grid)
 
 def _on_image_change(_):
-    return False  # uncheck "use sample" when a custom image is set
+    return False  # uncheck sample toggle
+
+def _preload_models():
+    for mid in MODELS:
+        try:
+            get_viz(mid)
+        except Exception:
+            pass
 
 theme = gr.themes.Soft(primary_hue="orange", neutral_hue="gray")
 
 with gr.Blocks(title="Transformers Processor Visualizer", theme=theme) as demo:
-    gr.Markdown("### Visualize what a processor feeds a vision–text model (uses the existing `ImageVisualizer`).")
+    gr.Markdown("## Visualize what a processor feeds a vision–text model")
 
     with gr.Row():
-        # LEFT: clickable models + custom field
-        with gr.Column(scale=1, min_width=260):
-            model_pick = gr.Radio(
-                label="Models",
-                choices=MODELS,
-                value=MODELS[0],
-                interactive=True,
-            )
-            custom_model = gr.Textbox(
-                label="Or type a model id",
-                placeholder="owner/repo",
-                lines=1,
-            )
-        # RIGHT: controls + outputs
+        with gr.Column(scale=1, min_width=280):
+            model_pick = gr.Radio(label="Models", choices=MODELS, value=MODELS[0], interactive=True)
+            custom_model = gr.Textbox(label="Or type a model id", placeholder="owner/repo", lines=1)
         with gr.Column(scale=3):
             with gr.Row():
                 add_grid = gr.Checkbox(label="Show patch grid", value=True)
                 use_sample = gr.Checkbox(label="Use HF logo sample", value=True)
-            image = gr.Image(
-                label="Upload custom image",
-                type="pil",
-                height=140,
-                sources=["upload"],
-            )
-            gr.Markdown("### Render")
-            gallery = gr.Gallery(label="Processor output", columns=[1], height=900)
+            image = gr.Image(label="Upload custom image", type="pil", height=140, sources=["upload"])
+            gr.Markdown("## Output")
+
+            with gr.Row():
+                left_output = gr.Image(label="Processor output", type="pil", height=900)
+                right_output = gr.Image(label="Global image (if any)", type="pil", height=900)
+
             prompt = gr.Textbox(label="Compact chat template preview", lines=2)
 
-    # Reactive updates: change model, toggle options, upload -> update immediately
-    model_pick.change(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt])
-    custom_model.submit(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt])
-    add_grid.change(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt])
-    use_sample.change(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt])
+    # reactive updates
+    model_pick.change(_resolve_and_run, [model_pick, custom_model, image, use_sample, add_grid], [left_output, right_output, prompt])
+    custom_model.submit(_resolve_and_run, [model_pick, custom_model, image, use_sample, add_grid], [left_output, right_output, prompt])
+    add_grid.change(_resolve_and_run, [model_pick, custom_model, image, use_sample, add_grid], [left_output, right_output, prompt])
+    use_sample.change(_resolve_and_run, [model_pick, custom_model, image, use_sample, add_grid], [left_output, right_output, prompt])
     image.change(_on_image_change, inputs=image, outputs=use_sample).then(
-        _resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt]
+        _resolve_and_run, [model_pick, custom_model, image, use_sample, add_grid], [left_output, right_output, prompt]
     )
 
-    # Initial render so there is output before any interaction
-    demo.load(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt])
+    # preload models into cache and render once
+    demo.load(_preload_models, [], [])
+    demo.load(_resolve_and_run, [model_pick, custom_model, image, use_sample, add_grid], [left_output, right_output, prompt])
 
 if __name__ == "__main__":
     demo.launch()
-