checkpoint viewer ready?

README.md

@@ -1,12 +1,44 @@
 ---
-title: Jaxgmg Hf
-emoji: 🐢
-colorFrom: yellow
-colorTo: green
+title: JaxGMG Action Probs Viewer
+emoji: 🎮
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
 sdk_version: 6.2.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# JaxGMG Action Probabilities Viewer
+
+Visualize action probability heatmaps from trained RL agents from the [JaxGMG collection](https://huggingface.co/collections/davidquarel/jaxgmg).
+
+## Features
+- Browse 14 model repositories from the JaxGMG collection
+- Searchable checkpoint dropdown
+- Interactive video player with:
+  - Play/pause, frame-by-frame navigation
+  - Click-drag to pan, scroll to zoom
+  - Fit to view / actual size buttons
+  - Fullscreen mode
+  - Adjustable FPS
+  - Keyboard shortcuts (Space, Arrows, F, 1, Esc)
+
+## Local Development
+
+```bash
+# Clone the model repo (optional, for faster local testing)
+git clone https://huggingface.co/davidquarel/jaxgmg_3phase_seed ../jaxgmg_3phase_seed
+cd ../jaxgmg_3phase_seed
+git lfs pull --include="*/action_probs.tar.gz"
+
+# Run with local data
+JAXGMG_USE_LOCAL=1 python app.py
+
+# Test without Gradio (generates HTML file)
+JAXGMG_USE_LOCAL=1 JAXGMG_TEST=1 python app.py
+```
+
+## Environment Variables
+- `JAXGMG_USE_LOCAL=1` - Use local cloned repos instead of downloading from HuggingFace
+- `JAXGMG_TEST=1` - Generate test HTML file instead of launching Gradio server

app.py

@@ -0,0 +1,230 @@
+import os
+import tarfile
+from functools import lru_cache
+from pathlib import Path
+
+import gradio as gr
+from huggingface_hub import HfApi, hf_hub_download
+from vidshow import vidshow_from_png_bytes
+
+JAXGMG_REPOS = [
+    "davidquarel/jaxgmg_3phase_seed",
+    "davidquarel/jaxgmg_3phase_bf16",
+    "davidquarel/jaxgmg_al_sweep",
+    "davidquarel/jaxgmg_al_sweep2",
+    "davidquarel/jaxgmg_al_g_sweep",
+    "davidquarel/jaxgmg_smoke",
+    "davidquarel/jaxgmg_al_1e0",
+    "davidquarel/jaxgmg_checkpoints",
+    "davidquarel/jaxgmg_open_alpha0_gamma_sweep",
+    "davidquarel/jaxgmg_ckpt_df",
+    "davidquarel/jaxgmg_ckpt_zip",
+    "davidquarel/jaxgmg_ckpt_pt_OLD",
+    "davidquarel/jaxgmg_test",
+    "davidquarel/matt_jaxgmg_ckpt",
+]
+
+DEFAULT_REPO = "davidquarel/jaxgmg_3phase_seed"
+LOCAL_BASE_PATH = Path(__file__).parent.parent
+USE_LOCAL = os.environ.get("JAXGMG_USE_LOCAL", "").lower() in ("1", "true", "yes")
+
+_hf_api: HfApi | None = None
+
+
+def get_hf_api() -> HfApi:
+    global _hf_api
+    if _hf_api is None:
+        _hf_api = HfApi()
+    return _hf_api
+
+
+def get_local_path(repo_id: str) -> Path:
+    return LOCAL_BASE_PATH / repo_id.split("/")[-1]
+
+
+@lru_cache(maxsize=32)
+def get_checkpoint_list(repo_id: str) -> tuple[str, ...]:
+    local_path = get_local_path(repo_id)
+    if USE_LOCAL and local_path.exists():
+        checkpoints = sorted(
+            d.name for d in local_path.iterdir()
+            if d.is_dir() and not d.name.startswith(".")
+        )
+    else:
+        files = get_hf_api().list_repo_files(repo_id, repo_type="model")
+        checkpoints = sorted(set(
+            f.split("/")[0] for f in files
+            if "/" in f and f.endswith("action_probs.tar.gz")
+        ))
+    return tuple(checkpoints)
+
+
+def update_checkpoints(repo_id: str):
+    checkpoints = list(get_checkpoint_list(repo_id))
+    return gr.Dropdown(
+        choices=checkpoints,
+        value=checkpoints[0] if checkpoints else None,
+    )
+
+
+def load_action_probs(repo_id: str, checkpoint: str, progress=gr.Progress()) -> str:
+    if not checkpoint:
+        return "<div style='padding:40px;text-align:center;color:#f88;'>Please select a checkpoint</div>"
+
+    progress(0.05, desc="Locating file...")
+
+    local_path = get_local_path(repo_id)
+    local_file = local_path / checkpoint / "action_probs.tar.gz"
+
+    if USE_LOCAL and local_file.exists() and local_file.stat().st_size > 1000:
+        tar_path = str(local_file)
+    else:
+        progress(0.1, desc="Downloading from HuggingFace...")
+        tar_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=f"{checkpoint}/action_probs.tar.gz",
+            repo_type="model",
+        )
+
+    progress(0.2, desc="Reading archive...")
+
+    png_bytes_list: list[bytes] = []
+    with tarfile.open(tar_path, "r:gz") as tar:
+        members = sorted(
+            (m for m in tar.getmembers() if m.name.endswith(".png")),
+            key=lambda m: m.name
+        )
+        total = len(members)
+        last_update = 0
+        for i, member in enumerate(members):
+            f = tar.extractfile(member)
+            if f is not None:
+                png_bytes_list.append(f.read())
+            if i - last_update >= 100:
+                progress(0.2 + 0.6 * (i / total), desc=f"Reading ({i}/{total})...")
+                last_update = i
+
+    progress(0.85, desc="Building player...")
+
+    title = f"{repo_id.split('/')[-1]} / {checkpoint}"
+    html = vidshow_from_png_bytes(png_bytes_list, fps=12, loop=True, title=title)
+
+    progress(1.0, desc="Done!")
+    return html
+
+
+def create_demo():
+    initial_checkpoints = list(get_checkpoint_list(DEFAULT_REPO))
+
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """
+            # 🎮 JaxGMG Action Probabilities Viewer
+            
+            Visualize action probability heatmaps from trained RL agents.
+            Use the slider to scrub through frames, or press play.
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                repo_dropdown = gr.Dropdown(
+                    choices=JAXGMG_REPOS,
+                    value=DEFAULT_REPO,
+                    label="Repository",
+                    filterable=True,
+                    info="Choose from the JaxGMG collection"
+                )
+                checkpoint_dropdown = gr.Dropdown(
+                    choices=initial_checkpoints,
+                    value=initial_checkpoints[0] if initial_checkpoints else None,
+                    label="Checkpoint",
+                    filterable=True,
+                    info="Type to search"
+                )
+                load_btn = gr.Button("Load", variant="primary", size="lg")
+
+                gr.Markdown(
+                    f"**Source:** {'Local' if USE_LOCAL else 'HuggingFace'} · "
+                    f"[Collection](https://huggingface.co/collections/davidquarel/jaxgmg)"
+                )
+
+        output_html = gr.HTML(
+            value="<div style='padding:40px;text-align:center;color:#888;'>Select a checkpoint and click Load</div>"
+        )
+
+        repo_dropdown.change(
+            fn=update_checkpoints,
+            inputs=[repo_dropdown],
+            outputs=[checkpoint_dropdown],
+        )
+
+        load_btn.click(
+            fn=load_action_probs,
+            inputs=[repo_dropdown, checkpoint_dropdown],
+            outputs=[output_html],
+        )
+
+    return demo
+
+
+def test_load():
+    """Test loading without Gradio - run with JAXGMG_TEST=1"""
+    print("Testing load...")
+    checkpoints = get_checkpoint_list(DEFAULT_REPO)
+    print(f"Found {len(checkpoints)} checkpoints")
+    checkpoint = checkpoints[0]
+    print(f"Loading {checkpoint}...")
+    
+    local_path = get_local_path(DEFAULT_REPO)
+    local_file = local_path / checkpoint / "action_probs.tar.gz"
+    
+    if USE_LOCAL and local_file.exists() and local_file.stat().st_size > 1000:
+        tar_path = str(local_file)
+        print(f"Using local: {tar_path}")
+    else:
+        print("Downloading from HuggingFace...")
+        tar_path = hf_hub_download(
+            repo_id=DEFAULT_REPO,
+            filename=f"{checkpoint}/action_probs.tar.gz",
+            repo_type="model",
+        )
+    
+    print("Reading archive...")
+    png_bytes_list: list[bytes] = []
+    with tarfile.open(tar_path, "r:gz") as tar:
+        members = sorted(
+            (m for m in tar.getmembers() if m.name.endswith(".png")),
+            key=lambda m: m.name
+        )
+        total = len(members)
+        print(f"Found {total} images")
+        for i, member in enumerate(members):
+            f = tar.extractfile(member)
+            if f is not None:
+                png_bytes_list.append(f.read())
+            if (i + 1) % 100 == 0:
+                print(f"  Read {i + 1}/{total}...")
+    
+    print(f"Loaded {len(png_bytes_list)} frames")
+    
+    print("Creating player...")
+    title = f"{DEFAULT_REPO.split('/')[-1]} / {checkpoint}"
+    html = vidshow_from_png_bytes(png_bytes_list, fps=12, title=title)
+    
+    # Save to local tmp directory
+    tmp_dir = Path(__file__).parent / "tmp"
+    tmp_dir.mkdir(exist_ok=True)
+    html_path = tmp_dir / "test_vidshow.html"
+    with open(html_path, "w") as f:
+        f.write(f"<!DOCTYPE html><html><body style='margin:0;background:#0f0f1a;'>{html}</body></html>")
+    print(f"Saved to {html_path}")
+    print("Open in browser to test!")
+
+
+if __name__ == "__main__":
+    if os.environ.get("JAXGMG_TEST", "").lower() in ("1", "true", "yes"):
+        test_load()
+    else:
+        demo = create_demo()
+        demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+gradio>=4.0.0
+huggingface_hub>=0.20.0
+numpy>=1.24.0
+Pillow>=10.0.0

vidshow.py

@@ -0,0 +1,267 @@
+import base64
+import html as html_module
+import json
+from concurrent.futures import ThreadPoolExecutor
+
+
+def vidshow_from_png_bytes(
+    png_bytes_list: list[bytes],
+    *,
+    fps: int = 12,
+    loop: bool = True,
+    title: str | None = None,
+) -> str:
+    """Optimized vidshow that takes raw PNG bytes directly. Returns iframe with embedded player."""
+
+    def encode_one(png_bytes: bytes) -> str:
+        return "data:image/png;base64," + base64.b64encode(png_bytes).decode("ascii")
+
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        urls = list(pool.map(encode_one, png_bytes_list))
+
+    N = len(urls)
+    frames_json = json.dumps(urls)
+    title_escaped = html_module.escape(title) if title else "Video"
+    loop_js = "true" if loop else "false"
+
+    # Build standalone HTML document for the iframe
+    inner_html = f"""<!DOCTYPE html>
+<html>
+<head>
+<meta charset="utf-8">
+<style>
+*{{box-sizing:border-box;}}
+html,body{{margin:0;padding:0;height:100%;}}
+body{{
+  padding:12px;
+  background:#1a1a2e;
+  font-family:sans-serif;
+  display:flex;
+  flex-direction:column;
+  height:100%;
+}}
+.title{{font-weight:600;color:#e0e0e0;margin-bottom:8px;font-size:1.1em;text-align:center;flex-shrink:0;}}
+.img-container{{
+  flex:1;
+  min-height:0;
+  overflow:hidden;
+  border-radius:8px;
+  box-shadow:0 4px 20px rgba(0,0,0,.4);
+  cursor:grab;
+  background:#0f0f1a;
+  position:relative;
+}}
+.img-container.dragging{{cursor:grabbing;}}
+img{{
+  position:absolute;
+  top:0;left:0;
+  image-rendering:pixelated;
+  user-select:none;
+  -webkit-user-drag:none;
+  transform-origin:0 0;
+}}
+.controls{{display:flex;gap:8px;align-items:center;margin-top:12px;flex-shrink:0;}}
+button{{padding:6px 12px;border-radius:6px;border:none;background:#4a4a6a;color:#fff;cursor:pointer;font-size:1em;}}
+button:hover{{background:#5a5a7a;}}
+button.primary{{background:#6366f1;}}
+button.primary:hover{{background:#7376f2;}}
+input[type=range]{{flex:1;accent-color:#6366f1;}}
+.label{{min-width:7ch;text-align:right;font-variant-numeric:tabular-nums;color:#e0e0e0;}}
+.fps-row{{display:flex;align-items:center;gap:8px;margin-top:10px;flex-shrink:0;}}
+.fps-label{{font-size:0.9em;color:#a0a0a0;}}
+.fps-value{{min-width:3ch;text-align:right;color:#e0e0e0;}}
+.hint{{font-size:0.8em;color:#666;margin-top:8px;text-align:center;flex-shrink:0;}}
+</style>
+</head>
+<body>
+<div class="title">{title_escaped}</div>
+<div class="img-container" id="imgContainer">
+  <img id="img" src="{urls[0]}">
+</div>
+<div class="controls">
+  <button id="back">◀</button>
+  <button id="toggle" class="primary">⏯️</button>
+  <button id="fwd">▶</button>
+  <button id="fit" title="Fit to view">⊡</button>
+  <button id="actual" title="Actual size">1:1</button>
+  <button id="fullscreen" title="Fullscreen (Esc to exit)">⛶</button>
+  <input type="range" id="slider" min="0" max="{N-1}" value="0" step="1">
+  <span id="label" class="label">1/{N}</span>
+</div>
+<div class="fps-row">
+  <span class="fps-label">FPS:</span>
+  <input type="range" id="fps" min="1" max="60" value="{fps}" step="1">
+  <span id="fpsLabel" class="fps-value">{fps}</span>
+</div>
+<div class="hint">Drag to pan • Scroll to zoom • F=fit • 1=actual size • Esc=exit fullscreen</div>
+<script>
+var frames={frames_json};
+var loopEnabled={loop_js};
+var playing=false,timer=null,fps={fps},idx=0;
+
+var img=document.getElementById("img");
+var imgContainer=document.getElementById("imgContainer");
+var slider=document.getElementById("slider");
+var label=document.getElementById("label");
+var toggle=document.getElementById("toggle");
+var fpsSlider=document.getElementById("fps");
+var fpsLabelEl=document.getElementById("fpsLabel");
+var fitBtn=document.getElementById("fit");
+var actualBtn=document.getElementById("actual");
+var fullscreenBtn=document.getElementById("fullscreen");
+
+// Pan/zoom state
+var isDragging=false;
+var startX=0,startY=0;
+var panX=0,panY=0;
+var scale=1;
+var imgW=0,imgH=0;
+
+function updateTransform(){{
+  img.style.transform="translate("+panX+"px,"+panY+"px) scale("+scale+")";
+}}
+
+function fitToView(){{
+  var cw=imgContainer.clientWidth;
+  var ch=imgContainer.clientHeight;
+  if(imgW===0||imgH===0)return;
+  var scaleX=cw/imgW;
+  var scaleY=ch/imgH;
+  scale=Math.min(scaleX,scaleY,1);
+  panX=(cw-imgW*scale)/2;
+  panY=(ch-imgH*scale)/2;
+  updateTransform();
+}}
+
+function actualSize(){{
+  var cw=imgContainer.clientWidth;
+  var ch=imgContainer.clientHeight;
+  scale=1;
+  panX=(cw-imgW)/2;
+  panY=(ch-imgH)/2;
+  updateTransform();
+}}
+
+var initialLoad=true;
+img.onload=function(){{
+  imgW=img.naturalWidth;
+  imgH=img.naturalHeight;
+  if(initialLoad){{
+    fitToView();
+    initialLoad=false;
+  }}
+}};
+
+window.addEventListener("resize",fitToView);
+
+// Mouse drag for panning
+imgContainer.addEventListener("mousedown",function(e){{
+  if(e.button!==0)return;
+  isDragging=true;
+  startX=e.clientX-panX;
+  startY=e.clientY-panY;
+  imgContainer.classList.add("dragging");
+  e.preventDefault();
+}});
+
+document.addEventListener("mousemove",function(e){{
+  if(!isDragging)return;
+  panX=e.clientX-startX;
+  panY=e.clientY-startY;
+  updateTransform();
+}});
+
+document.addEventListener("mouseup",function(){{
+  isDragging=false;
+  imgContainer.classList.remove("dragging");
+}});
+
+// Scroll to zoom
+imgContainer.addEventListener("wheel",function(e){{
+  e.preventDefault();
+  var delta=e.deltaY>0?0.9:1.1;
+  var newScale=Math.max(0.1,Math.min(10,scale*delta));
+  
+  var rect=imgContainer.getBoundingClientRect();
+  var mx=e.clientX-rect.left;
+  var my=e.clientY-rect.top;
+  
+  panX=mx-(mx-panX)*(newScale/scale);
+  panY=my-(my-panY)*(newScale/scale);
+  scale=newScale;
+  
+  updateTransform();
+}},{{passive:false}});
+
+fitBtn.onclick=fitToView;
+actualBtn.onclick=actualSize;
+
+function toggleFullscreen(){{
+  if(!document.fullscreenElement){{
+    document.documentElement.requestFullscreen().then(function(){{
+      setTimeout(fitToView,100);
+    }});
+  }}else{{
+    document.exitFullscreen();
+  }}
+}}
+fullscreenBtn.onclick=toggleFullscreen;
+document.addEventListener("fullscreenchange",function(){{
+  setTimeout(fitToView,100);
+}});
+
+function show(i){{
+  i=Math.max(0,Math.min(frames.length-1,i|0));
+  idx=i;
+  img.src=frames[i];
+  slider.value=i;
+  label.textContent=(i+1)+"/"+frames.length;
+}}
+
+function step(d){{
+  var i=idx+d;
+  if(i>=frames.length)i=loopEnabled?0:frames.length-1;
+  if(i<0)i=loopEnabled?frames.length-1:0;
+  show(i);
+}}
+
+function play(){{
+  if(playing)return;
+  playing=true;
+  toggle.textContent="⏸️";
+  timer=setInterval(function(){{step(1);}},1000/Math.max(1,fps));
+}}
+
+function pause(){{
+  if(!playing)return;
+  playing=false;
+  toggle.textContent="⏯️";
+  clearInterval(timer);
+  timer=null;
+}}
+
+slider.oninput=function(e){{show(parseInt(e.target.value));}};
+document.getElementById("back").onclick=function(){{pause();step(-1);}};
+document.getElementById("fwd").onclick=function(){{pause();step(1);}};
+toggle.onclick=function(){{if(playing)pause();else play();}};
+fpsSlider.oninput=function(){{
+  fps=parseInt(fpsSlider.value);
+  fpsLabelEl.textContent=fps;
+  if(playing){{pause();play();}}
+}};
+document.onkeydown=function(e){{
+  if(e.key===" "){{e.preventDefault();if(playing)pause();else play();}}
+  else if(e.key==="ArrowRight")step(1);
+  else if(e.key==="ArrowLeft")step(-1);
+  else if(e.key==="f"||e.key==="F")fitToView();
+  else if(e.key==="1")actualSize();
+}};
+</script>
+</body>
+</html>"""
+
+    # Escape the inner HTML for use in srcdoc attribute
+    srcdoc_escaped = html_module.escape(inner_html)
+    
+    # Wrap in iframe - srcdoc allows scripts to run
+    return f'<iframe srcdoc="{srcdoc_escaped}" style="width:100%;height:90vh;border:none;border-radius:12px;"></iframe>'