""" HuggingFace Spaces — Gradio demo for lsr-lang. Models are downloaded from singhanshuman/lsr-lang-models on first run. Compatible with Gradio 4.x and 6.x. """ import io import os from pathlib import Path import gradio as gr import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch from huggingface_hub import hf_hub_download from PIL import Image DEVICE = torch.device("cpu") Z_DIM = 4 ACTION_DIM = 4 MODEL_REPO = "singhanshuman/lsr-lang-models" from models.apm import APM from models.clip_vae import ClipVAE from models.lsr import LSR from models.vae import VAE def _get(filename: str) -> str: return hf_hub_download(repo_id=MODEL_REPO, filename=filename) def _load(ckpt: str, model: torch.nn.Module) -> torch.nn.Module: model.load_state_dict(torch.load(ckpt, map_location=DEVICE, weights_only=True)) model.eval() return model print("Loading models…") vae = _load(_get("vae_best.pt"), VAE(z_dim=Z_DIM)) clip_vae = _load(_get("clip_vae_best.pt"), ClipVAE(z_dim=Z_DIM)) apm = _load(_get("apm_best.pt"), APM(z_dim=Z_DIM, action_dim=ACTION_DIM)) text_emb = torch.tensor(np.load(_get("text_emb.npy")), dtype=torch.float32) lsr = LSR.load(_get("lsr_graph.pkl")) print("Ready.") # ── helpers ─────────────────────────────────────────────────────────────────── def _pil_to_tensor(img: Image.Image) -> torch.Tensor: img = img.convert("RGB").resize((64, 64)) return torch.tensor(np.array(img), dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.0 def _tensor_to_pil(t: torch.Tensor) -> Image.Image: arr = (t.squeeze().permute(1, 2, 0).clamp(0, 1).detach().numpy() * 255).astype(np.uint8) return Image.fromarray(arr) def _encode(img: Image.Image, use_clip: bool, lang: str) -> np.ndarray: x = _pil_to_tensor(img) with torch.no_grad(): if use_clip: t_emb = text_emb.clone() if lang.strip(): from models.clip_vae import encode_text t_emb = encode_text([lang], DEVICE) _, mu, _ = clip_vae.encode(x, t_emb.expand(1, -1)) else: _, mu, _ = vae.encode(x) return mu.squeeze().numpy() def _decode(z: np.ndarray, use_clip: bool) -> Image.Image: zt = torch.tensor(z, dtype=torch.float32).unsqueeze(0) model = clip_vae if use_clip else vae with torch.no_grad(): return _tensor_to_pil(model.decode(zt)) def _latent_plot(latent_path: np.ndarray) -> Image.Image: if lsr.latents is None: return None all_z = lsr.latents if all_z.shape[1] > 2: from sklearn.decomposition import PCA pca = PCA(n_components=2).fit(all_z) coords = pca.transform(all_z) path_coords = pca.transform(latent_path) else: coords = all_z path_coords = latent_path fig, ax = plt.subplots(figsize=(5, 5)) ax.scatter(coords[:, 0], coords[:, 1], s=5, alpha=0.35, c="lightgray") ax.plot(path_coords[:, 0], path_coords[:, 1], "r-o", ms=7, lw=2, label="plan") ax.scatter(*path_coords[0], c="limegreen", s=100, zorder=5, label="start") ax.scatter(*path_coords[-1], c="royalblue", s=100, zorder=5, label="goal") ax.legend(fontsize=9) ax.grid(True, alpha=0.25) ax.set_title("Latent space + planned path") plt.tight_layout() buf = io.BytesIO() fig.savefig(buf, format="png", dpi=120) plt.close(fig) buf.seek(0) return Image.open(buf).copy() # ── inference ───────────────────────────────────────────────────────────────── def run_plan(start_np, goal_np, lang, model_choice): use_clip = model_choice == "CLIP-VAE" if start_np is None or goal_np is None: return [], None, "Please upload both a start and a goal image." start_pil = Image.fromarray(start_np) goal_pil = Image.fromarray(goal_np) z_start = _encode(start_pil, use_clip, lang) z_goal = _encode(goal_pil, use_clip, lang) result = lsr.plan(z_start, z_goal) if result is None: return [], None, "No path found in the latent roadmap." path, latent_path = result plan_imgs = [start_pil.resize((64, 64))] for z in latent_path: plan_imgs.append(_decode(z, use_clip)) plan_imgs.append(goal_pil.resize((64, 64))) action_lines = [] with torch.no_grad(): for i in range(len(latent_path) - 1): zi = torch.tensor(latent_path[i], dtype=torch.float32).unsqueeze(0) zj = torch.tensor(latent_path[i + 1], dtype=torch.float32).unsqueeze(0) a = apm(zi, zj).squeeze().tolist() vals = a if isinstance(a, list) else [a] action_lines.append(f"Step {i+1}: [{', '.join(f'{v:.3f}' for v in vals)}]") info = ( f"Model : {'CLIP-VAE' if use_clip else 'Baseline VAE'}\n" f"Path : {len(path)} nodes\n\n" "Predicted actions:\n" + "\n".join(action_lines) ) return plan_imgs, _latent_plot(latent_path), info # ── UI ──────────────────────────────────────────────────────────────────────── with gr.Blocks(title="lsr-lang — Visual Action Planning") as demo: gr.Markdown( "## lsr-lang — Latent Space Roadmap with Language Conditioning\n" "Upload **start** and **goal** images from a box-stacking scene. " "The system plans a visual path through latent space and predicts robot actions.\n\n" "**[GitHub](https://github.com/anshuman-dev/lsr-lang)** · " "**[Paper — LSR-v2 (IEEE T-RO 2023)](https://arxiv.org/abs/2103.02554)**" ) with gr.Row(): with gr.Column(scale=1): inp_start = gr.Image(label="Start image", type="numpy") inp_goal = gr.Image(label="Goal image", type="numpy") inp_lang = gr.Textbox( label="Language goal (optional, CLIP-VAE only)", placeholder="stack the red box on the blue box", ) inp_model = gr.Radio( choices=["Baseline VAE", "CLIP-VAE"], value="Baseline VAE", label="Model", ) btn = gr.Button("Plan →", variant="primary") with gr.Column(scale=2): out_gallery = gr.Gallery(label="Visual plan (start → plan → goal)", columns=12) out_latent = gr.Image(label="Latent space + path") out_info = gr.Textbox(label="Predicted actions", lines=8) btn.click( fn=run_plan, inputs=[inp_start, inp_goal, inp_lang, inp_model], outputs=[out_gallery, out_latent, out_info], ) if __name__ == "__main__": demo.launch()