Upload web/app.py with huggingface_hub

web/app.py

@@ -64,6 +64,20 @@ NU = cfg["generator"]["num_uv"]
 _predict(jnp.zeros(NU * NU * 3))
 print(f"Model loaded and JIT-compiled. Grid: {NU}x{NU}")
 
+# Load VAE model for diversity sampling
+VAE_PATH = os.path.join(ROOT, "data", f"variational_formfinder_variational_{TASK}.eqx")
+vae_model = None
+if os.path.exists(VAE_PATH):
+    try:
+        vae_cfg_path = os.path.join(ROOT, "scripts", f"variational_{TASK}.yml")
+        with open(vae_cfg_path) as f:
+            vae_cfg = yaml.load(f, Loader=yaml.FullLoader)
+        vae_skeleton = build_neural_model("variational_formfinder", vae_cfg, gen, mk)
+        vae_model = load_model(VAE_PATH, vae_skeleton)
+        print("VAE model loaded for diversity sampling.")
+    except Exception as e:
+        print(f"VAE not loaded: {e}")
+
 # Static topology (edges, boundary vertices, faces) - sent once
 EDGES = np.array(list(mesh.edges())).tolist()
 BOUNDARY = sorted(set(mesh.vertices_on_boundary()))
@@ -156,6 +170,53 @@ def predict(req: PredictRequest):
     })
 
 
+class DiversityRequest(BaseModel):
+    c1_z: float = 3.0
+    c2_x: float = 0.0
+    c2_z: float = 1.5
+    c3_y: float = 0.0
+    n_samples: int = 5
+
+
+@app.post("/api/diversity")
+def sample_diversity(req: DiversityRequest):
+    """Sample diverse equilibrium solutions from the VAE."""
+    if vae_model is None:
+        return JSONResponse({"error": "VAE model not available"}, status_code=404)
+
+    transform = jnp.array([
+        [0.0, 0.0, req.c1_z],
+        [req.c2_x, 0.0, req.c2_z],
+        [0.0, req.c3_y, 0.0],
+        [0.0, 0.0, 0.0],
+    ])
+
+    xyz_target = gen.evaluate_points(transform)
+    key = jrn.PRNGKey(int(time.time()) % 10000)
+    x_hats, qs = vae_model.sample(xyz_target, structure, key, num_samples=req.n_samples)
+
+    samples = []
+    for i in range(req.n_samples):
+        pred_np = np.array(x_hats[i]).reshape(-1, 3)
+        q_np = np.array(qs[i]).flatten()
+        samples.append({
+            "predicted": pred_np.tolist(),
+            "q": q_np.tolist(),
+        })
+
+    return JSONResponse({
+        "samples": samples,
+        "n_samples": req.n_samples,
+        "has_vae": True,
+    })
+
+
+@app.get("/api/has_vae")
+def has_vae():
+    """Check if VAE model is available."""
+    return {"has_vae": vae_model is not None}
+
+
 # Serve static files
 STATIC_DIR = os.path.join(os.path.dirname(__file__), "static")
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")