Switch the demo to the gradio GUI

Run demo.py --with-gui gradio (Dockerfile CMD); add gradio + scipy to
requirements; refresh README. Sync the optgs source + demo.py to the current
upstream-merged tree (the gradio GUI: live decoder-render stream + interactive
Model3D, with the Z-up->Y-up splat reorientation and the thread-safe
optimize_iter). Prebuilt wheels and submodules are unchanged.

Dockerfile

@@ -1,7 +1,7 @@
 # Learn2Splat — interactive demo for a Hugging Face Space (Docker SDK, GPU).
 #
 # Installs the optgs package + prebuilt CUDA-extension wheels, then runs
-# demo.py's viser GUI: SfM-initialize a COLMAP scene and refine the Gaussians
+# demo.py's gradio GUI: SfM-initialize a COLMAP scene and refine the Gaussians
 # with the learned optimizer live in the browser.
 #
 # The CUDA extensions are NOT compiled here — the HF Docker builder runs out of
@@ -53,15 +53,6 @@ RUN pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 \
 COPY --chown=user:user requirements.txt .
 RUN pip install -r requirements.txt
 
-# HF serves the Space over HTTP/2, and the WebSocket subprotocol viser uses to
-# announce its client version doesn't survive the proxy — so viser's server
-# reads the client version as "unknown" and rejects the connection (the GUI
-# then hangs on "connecting"). Client and server are the same viser build
-# here, so treat an undeterminable client version as a match, not a reject.
-RUN VISER_INFRA="$(python -c 'import viser.infra._infra as m; print(m.__file__)')" \
- && sed -i 's/client_version_str = "unknown"/client_version_str = viser.__version__/' "$VISER_INFRA" \
- && grep -q 'client_version_str = viser.__version__' "$VISER_INFRA"
-
 # Prebuilt CUDA-extension wheels — gsplat, nerfacc, pycolmap, fused-ssim,
 # simple-knn, pointops, fused_knn_attn. Built on a matching machine (see
 # DEPLOY.md) so the HF builder never compiles CUDA and never OOMs.
@@ -72,11 +63,10 @@ RUN pip install --no-deps ./wheels/*.whl
 COPY --chown=user:user . .
 RUN pip install --no-build-isolation --no-deps -e .
 
-# viser serves the GUI here — must equal app_port in README.md.
+# gradio serves the GUI here — must equal app_port in README.md.
 EXPOSE 7860
 
-# server mode: the optgs decoder renders frames on the GPU and viser streams
-# them as images — no multi-MB splat-geometry transfer to the browser and no
-# client-side WebGL, which is far more robust behind HF's HTTP/2 proxy.
-# viser binds 0.0.0.0 by default.
-CMD ["python", "demo.py", "--with-gui", "server", "--gui-port", "7860"]
+# gradio mode: the optgs decoder renders the live optimization on the GPU and
+# gradio streams the frames as images (column 2); the finished splats load into
+# an interactive Model3D viewer (column 3). demo.py binds 0.0.0.0 on this port.
+CMD ["python", "demo.py", "--with-gui", "gradio", "--gui-port", "7860"]

README.md

@@ -14,11 +14,13 @@ short_description: Interactive demo of the Learn2Splat learned 3DGS optimizer
 A learned optimizer for 3D Gaussian Splatting. This Space SfM-initializes a
 COLMAP scene and refines the Gaussians live in your browser: pick the
 Learn2Splat optimizer (dense or sparse checkpoint) or a 3DGS Adam baseline,
-press **Start**, and watch the splats converge.
+press **Start**, and watch the decoder render converge — then explore the
+finished splats in an interactive 3D viewer.
 
-Runs `demo.py --with-gui client` from the
+Runs `demo.py --with-gui gradio` from the
 [Learn2Splat repository](https://github.com/autonomousvision/learn2splat);
-the splats are drawn by viser's in-browser WebGL renderer.
+the optimization is rendered on the GPU and streamed by gradio, and the
+result loads into a `Model3D` splat viewer.
 
 > Requires GPU hardware. The demo holds two checkpoints in VRAM at once —
 > an A10G (24 GB) is recommended.

demo.py

@@ -27,6 +27,7 @@ Usage (run from the repo root, with ``optgs`` importable):
     python demo.py                    # headless: dense + sparse checkpoints + an Adam baseline
     python demo.py --with-gui server  # interactive viser GUI (frames rendered by the decoder)
     python demo.py --with-gui client  # interactive viser GUI (viser's WebGL splat renderer)
+    python demo.py --with-gui gradio  # interactive gradio GUI (streamed renders + Model3D splats)
 
 The demo scene and the checkpoints are fetched from the Hugging Face Hub on
 first run (cached under ./data and ./checkpoints). A CUDA device is required.
@@ -118,11 +119,13 @@ class Config:
     seed: int = 42
 
     # --- Interactive GUI ---
-    # Launch a viser GUI instead of the headless comparison. "server" renders
-    # frames with the optgs decoder; "client" uses viser's built-in WebGL
-    # Gaussian-splat renderer. Unset = headless run.
-    with_gui: Optional[Literal["client", "server"]] = None
-    # Port for the viser GUI web server (--with-gui only).
+    # Launch an interactive GUI instead of the headless comparison. viser:
+    # "server" renders frames with the optgs decoder, "client" uses viser's
+    # built-in WebGL Gaussian-splat renderer. "gradio" runs a browser GUI
+    # (decoder renders streamed live + an interactive Model3D splat viewer for
+    # the result). Unset = headless run.
+    with_gui: Optional[Literal["client", "server", "gradio"]] = None
+    # Port for the GUI web server (--with-gui only).
     gui_port: int = 8080
 
     # --- OptGS learned optimizer ---
@@ -600,6 +603,373 @@ def run_gui(
         console.print("\n[yellow]GUI stopped.[/]")
 
 
+def run_gradio_gui(
+    instances: dict,
+    gaussians: Gaussians,
+    train_bv: dict,
+    cfg: Config,
+    device: torch.device,
+) -> None:
+    """Interactive gradio GUI — a browser port of :func:`run_gui` (viser).
+
+    gradio can't stream the camera back to Python, so there is no free-camera
+    server rendering. Instead the optimization is *watched* as a streamed
+    decoder render from a chosen training view (``gr.Image``, refreshed every
+    step), and the finished scene is handed to an interactive ``gr.Model3D``
+    splat viewer (orbit / zoom in the browser). The controls mirror the viser
+    GUI: pick the optimizer (Learn2Splat dense/sparse or a 3DGS Adam baseline),
+    set the step budget / view-minibatch size / sampling strategy, Start, Reset.
+
+    ``instances`` maps "dense"/"sparse" to their initialized ``OptGS``.
+    """
+    import gc
+
+    import gradio as gr
+
+    from optgs.experimental.api.integration.config_bridge import build_adam_baseline
+    from optgs.misc.image_io import prep_image
+    from optgs.model.ply_export import save_gaussian_ply
+
+    # Optimizer dropdown label -> (instances key, swap in a 3DGS Adam baseline);
+    # mirrors run_gui's OPTIONS.
+    OPTIONS: Dict[str, Tuple[str, bool]] = {
+        "Learn2Splat (dense)": ("dense", False),
+        "Learn2Splat (sparse)": ("sparse", False),
+        "Adam (3DGS)": ("dense", True),
+    }
+
+    n_train_views = int(train_bv["image"].shape[1])
+    h_full, w_full = train_bv["image"].shape[3], train_bv["image"].shape[4]
+    init_gaussians = gaussians.clone()  # pristine copy; each Start re-inits from it
+
+    # Shared state (single-GPU, single-session demo, like run_gui's globals): the
+    # Gaussians currently shown, the OptGS rendering them, and a counter for
+    # unique PLY filenames (so Model3D reloads instead of serving a stale cache).
+    holder = {"current": init_gaussians, "active": instances["dense"], "ply": 0}
+
+    ply_dir = os.path.join(cfg.result_dir, "gradio")
+    os.makedirs(ply_dir, exist_ok=True)
+
+    @torch.no_grad()
+    def render_decoder(inst, g: Gaussians, view_idx: float, height: float) -> np.ndarray:
+        """Decoder-render Gaussians ``g`` from training view ``view_idx``.
+
+        Normalized intrinsics make the render resolution-independent; the width
+        is derived from ``height`` at the training views' aspect ratio.
+        """
+        h = int(height)
+        w = max(1, round(h * w_full / h_full))
+        sl = slice(int(view_idx), int(view_idx) + 1)
+        out = inst.decoder.forward(
+            g,
+            train_bv["extrinsics"][:, sl],
+            train_bv["intrinsics"][:, sl],
+            train_bv["near"][:, sl],
+            train_bv["far"][:, sl],
+            image_shape=(h, w),
+        )
+        return prep_image(out.color[0, 0])  # [H, W, 3] uint8
+
+    def reorient_for_viewer(g: Gaussians) -> Gaussians:
+        """Reorient a copy of ``g`` from the COLMAP world (this scene is Z-up)
+        into the Y-up frame the gradio Model3D shows upright.
+
+        gradio/Babylon flips the loaded splats' Y (``scaling.y *= -1``), so
+        exporting through the reflection E(p)=(x,-z,-y) makes the *displayed*
+        scene N(p)=(x,z,-y) — the world's Z-up mapped onto Babylon's Y-up. E's
+        point-reflection part leaves the covariance unchanged; its proper-rotation
+        part R_p=-E rotates the splat orientations to match.
+        """
+        from scipy.spatial.transform import Rotation as Rsp
+
+        g = g.clone()
+        m = g.means[0]
+        g.means = torch.stack([m[:, 0], -m[:, 2], -m[:, 1]], dim=1)[None]  # E
+        q = F.normalize(g.rotations_unnorm[0], dim=-1).detach().cpu().numpy()  # xyzw
+        R_p = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
+        q = (Rsp.from_matrix(R_p) * Rsp.from_quat(q)).as_quat()  # R_p @ R, xyzw
+        g.rotations_unnorm = torch.from_numpy(q).to(g.means)[None]
+        return g
+
+    def export_ply(g: Gaussians) -> str:
+        """Write ``g`` (reoriented to Y-up) to a fresh PLY path for the viewer."""
+        from pathlib import Path
+
+        holder["ply"] += 1
+        path = os.path.join(ply_dir, f"result_{holder['ply']}.ply")
+        save_gaussian_ply(reorient_for_viewer(g), save_path=Path(path))
+        return path
+
+    def start(optimizer_label, max_steps, batch_size, strategy, view_idx, height):
+        """Generator: run the picked optimizer, streaming a decoder render each
+        step, then load the finished splats into the Model3D viewer."""
+        name, use_adam = OPTIONS[optimizer_label]
+        inst = instances[name]
+        holder["active"] = inst
+        # Re-init from the pristine copy so repeated Starts share one start point.
+        inst.initialize_from_tensors(init_gaussians.clone(), train_bv)
+        inst.num_refine = int(max_steps)
+        inst.opt_batch_size = min(int(batch_size), n_train_views)
+        inst.opt_batch_strategy = strategy
+        opt = build_adam_baseline(inst.num_refine).to(device) if use_adam else None
+
+        try:
+            # Disable Start, hide the viewer, keep the placeholder while running.
+            yield (
+                render_decoder(inst, init_gaussians, view_idx, height),
+                f"**optimizing** — step 0/{inst.num_refine} — "
+                f"{init_gaussians.means.shape[1]} Gaussians",
+                gr.update(visible=False, value=None),  # hide the viewer
+                gr.update(interactive=False),          # disable Start
+                gr.update(visible=True),               # keep the placeholder
+            )
+
+            g = init_gaussians
+            for step, g in inst.optimize_iter(optimizer=opt):
+                holder["current"] = g
+                yield (
+                    render_decoder(inst, g, view_idx, height),
+                    f"**optimizing** — step {step + 1}/{inst.num_refine} — "
+                    f"{g.means.shape[1]} Gaussians",
+                    gr.update(), gr.update(), gr.update(),  # no change mid-run
+                )
+
+            yield (
+                render_decoder(inst, g, view_idx, height),
+                f"**done** — {inst.num_refine} steps — {g.means.shape[1]} Gaussians",
+                gr.update(visible=True, value=export_ply(g)),  # reveal the viewer
+                gr.update(interactive=True),                   # re-enable Start
+                gr.update(visible=False),                      # hide the placeholder
+            )
+        finally:
+            # Free the run's CUDA work (also runs if the user hits Stop mid-run),
+            # so GPU memory doesn't accumulate across repeated runs.
+            opt = None
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def reset(view_idx, height):
+        """Restore the initialization: re-render it, hide the viewer.
+
+        No CUDA cleanup here — ``start``'s ``finally`` already frees each run's
+        work; reset only re-renders the init.
+        """
+        holder["current"] = init_gaussians
+        holder["active"] = instances["dense"]
+        return (
+            render_decoder(instances["dense"], init_gaussians, view_idx, height),
+            "**Initialized.** Pick a method, then Start.",
+            gr.update(visible=False, value=None),  # hide the viewer
+            gr.update(interactive=True),           # enable Start
+            gr.update(visible=True),               # show the placeholder
+        )
+
+    def rerender(view_idx, height):
+        """Re-render the current Gaussians (preview view / height changed)."""
+        return render_decoder(holder["active"], holder["current"], view_idx, height)
+
+    initial_img = render_decoder(instances["dense"], init_gaussians, 0, 540)
+
+    # Open the interactive viewer on the same vantage as the live render's
+    # default preview (view 0). The viewer shows splats in the reoriented frame
+    # N(p)=(x,z,-y) (see reorient_for_viewer); the orbit camera sits at the
+    # training view's position and looks at the scene centroid. babylon_camera
+    # maps the world camera into N and inverts Babylon's ArcRotateCamera position
+    # formula rel=(r·cosα·sinβ, r·cosβ, r·sinα·sinβ) into (alpha°, beta°, radius).
+    def babylon_camera(c2w: np.ndarray, centroid: np.ndarray) -> tuple:
+        p = c2w[:3, 3] - centroid
+        rel = np.array([p[0], p[2], -p[1]], dtype=np.float64)  # N applied to (cam - centroid)
+        radius = float(np.linalg.norm(rel)) or 1e-3
+        beta = float(np.degrees(np.arccos(np.clip(rel[1] / radius, -1.0, 1.0))))
+        alpha = float(np.degrees(np.arctan2(rel[2], rel[0])))
+        return (alpha, beta, radius)
+
+    means0 = init_gaussians.means[0].detach().float().cpu().numpy()
+    centroid0 = (means0.min(0) + means0.max(0)) / 2.0
+    cam0_c2w = train_bv["extrinsics"][0, 0].detach().float().cpu().numpy()
+    init_camera = babylon_camera(cam0_c2w, centroid0)
+
+    # --- Visual style: lifted from the Learn2Splat project page
+    # (https://naamapearl.github.io/learn2splat/) — plum accent (#B04080) with an
+    # indigo hover, a light slate canvas with white cards, Source Serif 4
+    # headings / Inter body / JetBrains Mono labels. ---
+    plum = gr.themes.Color(
+        c50="#faf0f6", c100="#f5e6ef", c200="#eccadd", c300="#dda3c2",
+        c400="#c66ba0", c500="#b04080", c600="#9b3570", c700="#7f2a5b",
+        c800="#6a2550", c900="#581f43", c950="#350e26", name="plum",
+    )
+    slate = gr.themes.Color(
+        c50="#f7f8fc", c100="#eef0f8", c200="#e2e5ef", c300="#c8cde0",
+        c400="#8890b0", c500="#5a6080", c600="#454b6b", c700="#343a56",
+        c800="#262b42", c900="#1a1d2e", c950="#0f1120", name="slate",
+    )
+    theme = gr.themes.Soft(
+        primary_hue=plum,
+        neutral_hue=slate,
+        font=["Inter", "system-ui", "sans-serif"],
+        font_mono=["JetBrains Mono", "ui-monospace", "monospace"],
+    ).set(
+        body_background_fill="#f7f8fc",
+        body_text_color="#1a1d2e",
+        block_background_fill="#ffffff",
+        block_border_color="#e2e5ef",
+        block_radius="12px",
+        block_label_text_color="#5a6080",
+        block_title_text_color="#1a1d2e",
+        border_color_primary="#e2e5ef",
+        input_background_fill="#ffffff",
+        button_primary_background_fill="#b04080",
+        button_primary_background_fill_hover="#3d50c0",
+        button_primary_text_color="#ffffff",
+        button_secondary_background_fill="#ffffff",
+        button_secondary_border_color="#c8cde0",
+        button_secondary_text_color="#1a1d2e",
+        slider_color="#b04080",
+        link_text_color="#b04080",
+    )
+    # Source Serif 4 / Inter / JetBrains Mono, loaded into <head> like the page.
+    fonts_head = (
+        '<link rel="preconnect" href="https://fonts.googleapis.com">'
+        '<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>'
+        '<link rel="stylesheet" href="https://fonts.googleapis.com/css2?'
+        "family=Source+Serif+4:opsz,wght@8..60,400;8..60,600&"
+        "family=Inter:wght@300;400;500;600&"
+        'family=JetBrains+Mono:wght@400;500&display=swap">'
+    )
+    css = """
+    .gradio-container { max-width: 1580px !important; margin: 0 auto !important; }
+    #l2s-hero { background: linear-gradient(180deg,#f3f4fb 0%,#f7f8fc 100%);
+        border:1px solid #e2e5ef; border-radius:14px; padding:26px 30px; margin-bottom:6px; }
+    #l2s-hero .eyebrow { font-family:'JetBrains Mono',monospace; font-size:12px;
+        letter-spacing:.16em; text-transform:uppercase; color:#b04080; font-weight:500;
+        display:inline-flex; align-items:center; gap:8px; }
+    #l2s-hero .eyebrow .dot { width:6px; height:6px; border-radius:50%; background:#b04080; }
+    #l2s-hero h1 { font-family:'Source Serif 4',Georgia,serif; font-weight:600; color:#1a1d2e;
+        font-size:clamp(1.55rem,3vw,2.25rem); line-height:1.2; margin:13px 0 10px; }
+    #l2s-hero p { color:#5a6080; font-size:15px; line-height:1.6; margin:0; max-width:820px; }
+    #l2s-hero a { color:#b04080; text-decoration:none; border-bottom:1px solid #e6cdd9;
+        white-space:nowrap; }
+    .l2s-eyebrow span { font-family:'JetBrains Mono',monospace; font-size:11px;
+        letter-spacing:.15em; text-transform:uppercase; color:#8890b0; font-weight:500; }
+    #l2s-start button, #l2s-reset button { font-family:'JetBrains Mono',monospace;
+        letter-spacing:.03em; font-weight:500; }
+    #l2s-status { background:#f7f8fc; border:1px solid #e2e5ef; border-left:3px solid #b04080;
+        border-radius:9px; padding:4px 14px; }
+    #l2s-status p { color:#5a6080; font-size:14px; margin:8px 0; }
+    #l2s-status strong { color:#1a1d2e; }
+    .l2s-ph { display:flex; align-items:center; justify-content:center; text-align:center;
+        height:420px; border:1px dashed #c8cde0; border-radius:12px; background:#fbfcff;
+        color:#8890b0; font-family:'JetBrains Mono',monospace; font-size:12.5px;
+        letter-spacing:.04em; line-height:1.7; }
+    footer { display:none !important; }
+    """
+    hero_html = (
+        "<div class='eyebrow'><span class='dot'></span>Learn2Splat · Interactive demo</div>"
+        "<h1>Extending the Horizon of Learned 3DGS Optimization</h1>"
+        "<p>SfM-initialize a COLMAP scene, then refine the Gaussians with the "
+        "meta-learned optimizer — pick a method, press <b>Start</b>, and watch the "
+        "decoder render converge. The finished splats load in the interactive 3D "
+        "viewer. <a href='https://naamapearl.github.io/learn2splat/' target='_blank' "
+        "rel='noopener'>Project page&nbsp;↗</a></p>"
+    )
+
+    with gr.Blocks(
+        title="Learn2Splat — Demo", theme=theme, css=css, head=fonts_head,
+        analytics_enabled=False,
+    ) as ui:
+        gr.HTML(hero_html, elem_id="l2s-hero")
+        with gr.Row(equal_height=False):
+            # Column 1 — controls.
+            with gr.Column(scale=3, min_width=300):
+                with gr.Group():
+                    gr.HTML("<div class='l2s-eyebrow'><span>Optimizer</span></div>")
+                    optimizer_dd = gr.Dropdown(
+                        list(OPTIONS), value=next(iter(OPTIONS)), label="Method"
+                    )
+                    with gr.Row():
+                        max_steps_input = gr.Number(
+                            value=cfg.max_steps, minimum=1, maximum=1000, step=1,
+                            precision=0, label="Max steps",
+                        )
+                        batch_size_input = gr.Number(
+                            value=min(cfg.opt_batch_size, n_train_views),
+                            minimum=1, maximum=n_train_views, step=1, precision=0,
+                            label="Opt batch size",
+                        )
+                    strategy_dd = gr.Dropdown(
+                        ["random", "sequential", "fps"],
+                        value=cfg.opt_batch_strategy, label="Batch strategy",
+                    )
+                with gr.Group():
+                    gr.HTML("<div class='l2s-eyebrow'><span>Preview</span></div>")
+                    view_slider = gr.Slider(
+                        0, n_train_views - 1, value=0, step=1, label="Preview view"
+                    )
+                    height_slider = gr.Slider(
+                        240, 1080, value=540, step=60, label="Render height"
+                    )
+                with gr.Row():
+                    start_btn = gr.Button(
+                        "Start optimization", variant="primary",
+                        elem_id="l2s-start", scale=2,
+                    )
+                    reset_btn = gr.Button(
+                        "Reset", variant="secondary", elem_id="l2s-reset", scale=1
+                    )
+                status_md = gr.Markdown(
+                    "**Initialized.** Pick a method, then Start.", elem_id="l2s-status"
+                )
+            # Column 2 — live decoder render (streamed during optimization).
+            with gr.Column(scale=5, min_width=380):
+                image_out = gr.Image(
+                    value=initial_img, label="Optimizer · live",
+                    height=540, format="jpeg", interactive=False,
+                )
+            # Column 3 — interactive splats (hidden until a run finishes; a
+            # placeholder holds the column so the 3-up layout stays balanced).
+            with gr.Column(scale=5, min_width=380):
+                model3d_out = gr.Model3D(
+                    label="Refined splats · interactive", height=540,
+                    visible=False, camera_position=init_camera,
+                )
+                placeholder = gr.HTML(
+                    "<div class='l2s-ph'>The interactive 3D splats<br>"
+                    "appear here once a run finishes.</div>"
+                )
+
+        start_inputs = [
+            optimizer_dd, max_steps_input, batch_size_input, strategy_dd,
+            view_slider, height_slider,
+        ]
+        gui_outputs = [image_out, status_md, model3d_out, start_btn, placeholder]
+        # One shared GPU lane (concurrency_id) so Start / Reset / preview re-renders
+        # never run on the GPU at the same time — overlapping runs were the path to
+        # runaway VRAM growth.
+        start_btn.click(
+            start, inputs=start_inputs, outputs=gui_outputs, concurrency_id="gpu"
+        )
+        reset_btn.click(
+            reset, inputs=[view_slider, height_slider], outputs=gui_outputs,
+            concurrency_id="gpu",
+        )
+        # Slider release (not change) — re-render once when the user lets go.
+        view_slider.release(
+            rerender, [view_slider, height_slider], image_out, concurrency_id="gpu"
+        )
+        height_slider.release(
+            rerender, [view_slider, height_slider], image_out, concurrency_id="gpu"
+        )
+
+    console.print(
+        f"[green]✓[/] gradio GUI on port [cyan]{cfg.gui_port}[/]"
+        f" — forward the port over SSH and open the printed URL"
+    )
+    ui.queue(default_concurrency_limit=1).launch(
+        server_name="0.0.0.0", server_port=cfg.gui_port, share=False,
+        show_error=True,
+    )
+
+
 def main(cfg: Config) -> None:
     # Fetch the demo scene on first run, before anything else touches it.
     ensure_data(cfg.data_dir)
@@ -663,7 +1033,10 @@ def main(cfg: Config) -> None:
         for inst in instances.values():
             inst.initialize_from_tensors(gaussians, train_bv)
 
-        run_gui(instances, gaussians, train_bv, cfg, device, dtype)
+        if cfg.with_gui == "gradio":
+            run_gradio_gui(instances, gaussians, train_bv, cfg, device)
+        else:
+            run_gui(instances, gaussians, train_bv, cfg, device, dtype)
         return
 
     val_c2w, val_Ks, val_images = collect_cameras(dataset, val_idx)

optgs/config.py

@@ -1,29 +1,34 @@
-import importlib
+"""Typed configuration schema and Hydra-to-dataclass loading.
+
+`RootCfg` is the top-level config dataclass (dataset, scene_trainer, meta_trainer, loss, checkpointing,
+...); `load_typed_root_config` converts the raw Hydra `DictConfig` into it via dacite, after running
+`config_migrate` so older checkpoints' configs load. `__post_init__` also resolves the output directory
+for `mode=test`. The per-component schemas live next to their components (e.g.
+`scene_trainer_cfg.py`, `meta_trainer_cfg.py`).
+"""
+
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal, Optional, Type, TypeVar, Any, Callable
+from typing import Literal, Optional, Type, TypeVar
 
-import hydra
 import torch
 from dacite import Config, from_dict, UnionMatchError
-from hydra.core.global_hydra import GlobalHydra
 from hydra.core.hydra_config import HydraConfig
-from hydra.types import RunMode
 from omegaconf import DictConfig
 from omegaconf import OmegaConf
 from pytorch_lightning.strategies import DDPStrategy, FSDPStrategy
 
 from .config_migrate import migrate, CURRENT_CFG_VERSION
 from .dataset.data_module import DataLoaderCfg, DatasetCfg
-from .global_cfg import set_cfg
 from .loss import LossCfgWrapper
 from .misc.io import CustomPath
 from .misc.io import cyan, read_omega_cfg
 from .misc.checkpointing import find_latest_ckpt
-from .misc.hf_ckpt import maybe_resolve_hf_ref
-from .paths import CKPT_DIR, RESULTS_DIR
-from .scene_trainer.scene_trainer_cfg import SceneTrainerCfg, MetaOptimizerCfg, TestCfg, TrainCfg
+from .misc.hf_ckpt import hf_sibling_config, is_hf_ref, maybe_resolve_hf_ref
+from .paths import CKPT_DIR, RESULTS_DIR, DEBUG
+from .scene_trainer.scene_trainer_cfg import SceneTrainerCfg
+from .meta_trainer.meta_trainer_cfg import MetaOptimizerCfg, TestCfg, TrainCfg
 
 
 # In order to extract filename or dirname from a path in the config
@@ -58,15 +63,28 @@ class CheckpointingCfg:
     resume_update_module: str | None
     load_existing_cfg: bool
 
+    # Checkpoints whose sibling config.yaml carries the model architecture that
+    # mode=test rebuilds from (the optimizer/initializer the Gaussians were
+    # trained with).
+    _ARCH_CHECKPOINTS = ("pretrained_model", "pretrained_optimizer", "pretrained_initializer")
+
     def __post_init__(self):
         # Resolve any Hugging Face Hub references (hf://org/repo/file[@rev]) to
         # local cached paths so all downstream torch.load calls work unchanged.
+        # Rebuilding the model architecture needs both the .ckpt and its sibling
+        # config.yaml, but hf_hub_download pulls only the .ckpt, so fetch the
+        # sibling config.yaml too. It lands in the same checkpoints/ layout as
+        # the resolved .ckpt, so _find_config_for_checkpoint later discovers it
+        # on the local path.
         for attr in ("pretrained_model", "pretrained_optimizer", "pretrained_initializer",
                      "pretrained_monodepth", "pretrained_mvdepth", "pretrained_depth",
                      "pretrained_scale_predictor", "pretrained_depth_teacher",
                      "resume_update_module"):
-            resolved = maybe_resolve_hf_ref(getattr(self, attr))
-            if resolved != getattr(self, attr):
+            ref = getattr(self, attr)
+            if attr in self._ARCH_CHECKPOINTS and is_hf_ref(ref):
+                hf_sibling_config(ref)
+            resolved = maybe_resolve_hf_ref(ref)
+            if resolved != ref:
                 setattr(self, attr, resolved)
 
         for attr in ("pretrained_model", "pretrained_optimizer", "pretrained_initializer"):
@@ -90,6 +108,7 @@ class MetaTrainerCfg:
     eval_index: str | None
     limit_test_batches: int | float
     limit_train_batches: int | float
+    meta_optimizer: MetaOptimizerCfg
     test: TestCfg
     train: TrainCfg
 
@@ -109,9 +128,9 @@ class MetaTrainerCfg:
                         return has_trainable
 
                     dist_strategy = FSDPStrategy(auto_wrap_policy=only_wrap_trainable)
-        if self.train.use_replay_buffer:
-            # When resampling from the replay buffer,
-            # we don't project the condition_features to state, so the update_proj is not used
+        if self.train.use_ckpt_buffer:
+            # When resampling from the ckpt buffer,
+            # we don't project the condition_features to state, so the state_proj is not used
             dist_strategy = "ddp_find_unused_parameters_true"
         return dist_strategy
 
@@ -123,14 +142,13 @@ class RootCfg:
     dataset: DatasetCfg
     data_loader: DataLoaderCfg
     scene_trainer: SceneTrainerCfg
-    meta_optimizer: MetaOptimizerCfg  ## TODO Naama: should we move under meta trainer config?
     checkpointing: CheckpointingCfg
     meta_trainer: MetaTrainerCfg
     loss: list[LossCfgWrapper]
     seed: int
     use_plugins: bool
     output_dir: str
-    version: int | None
+    version: int | float | None  # point versions like 2.1 are floats
     debug_cfg: bool
 
     def __post_init__(self):
@@ -202,12 +220,6 @@ TYPE_HOOKS = {
 T = TypeVar("T")
 
 
-def get_class_by_path(path: str):
-    module_path, class_name = path.rsplit('.', 1)
-    module = importlib.import_module(module_path)
-    return getattr(module, class_name)
-
-
 def _diagnose_union_error(e: UnionMatchError, data: dict, dacite_config: Config) -> str:
     """Try each union member individually and report per-member errors."""
     import dataclasses
@@ -273,32 +285,6 @@ def separate_loss_cfg_wrappers(joined: dict) -> list[LossCfgWrapper]:
     ]
 
 
-def universal_target_hook(cfg: dict, _: Type) -> Any:
-    """Generic hook to construct config objects from `__target__`."""
-    if not isinstance(cfg, dict):
-        return None
-    if "__target__" not in cfg:
-        return None  # Let decite handle it
-
-    cfg_copy = deepcopy(cfg)  # avoid mutating original
-    target = cfg_copy.pop("__target__")
-
-    if isinstance(target, str):
-        target_type = get_class_by_path(target)
-    else:
-        target_type = target
-
-    # Use recursive loading with known additional hooks
-    return load_typed_config(
-        DictConfig(cfg_copy),
-        target_type,
-    )
-
-
-def make_target_hook_for_type(t: Type) -> Callable:
-    return lambda cfg: universal_target_hook(cfg, t)
-
-
 def load_typed_root_config(cfg: DictConfig) -> RootCfg:
     # scene_trainer/scene_optimizer=none loads a full dict from none.yaml;
     # dacite can't match that dict to the None arm of SceneOptimizerCfg | None.
@@ -316,30 +302,17 @@ def load_typed_root_config(cfg: DictConfig) -> RootCfg:
     )
 
 
-def should_run(cfg_dict):
-    if cfg_dict.mode == "test":
-        if cfg_dict.meta_trainer.test.skip_if_outputs_exist:
-            output_dir = cfg_dict.output_dir
-            if not output_dir.exists():
-                return True
-            metrics_path_pattern = output_dir / "metrics" / "target_*_psnr.json"
-            metric_paths = list(metrics_path_pattern.parent.glob(metrics_path_pattern.name))
-            if len(metric_paths) > 0:
-                print(cyan(f"Test metrics already exist at {metric_paths}."))
-                return False
-    return True
-
-
 def setup_cfg(cfg_dict):
     # Get the original config from the output directory, when testing or resuming.
     cfg_dict = merge_config_from_file(cfg_dict)
     eval_cfg = get_eval_cfg(cfg_dict)
     cfg = load_typed_root_config(cfg_dict)
-    # Set global cfg object.
-    set_cfg(cfg_dict)
     # Set up the output directory.
     setup_output_dir(cfg, cfg_dict)
-    return cfg, cfg_dict, eval_cfg  # TODO Naama: why do we need both cfg and cfg_dict?
+    # cfg is the typed config, used for attribute access throughout the code. cfg_dict is the
+    # raw OmegaConf tree: wandb logs it in full, the eval-config and migration helpers operate
+    # on it as nested dicts, and it holds fields outside the typed schema (e.g. profiling).
+    return cfg, cfg_dict, eval_cfg
 
 
 def flatten_wandb(cfg):
@@ -375,6 +348,11 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
     print(cyan(f"Re-applying {len(raw_overrides)} CLI overrides onto merged config."))
     OmegaConf.set_struct(merged_cfg, False)
 
+    # Per-key merge decisions are noisy on every run; show them only under a debugger.
+    def _trace(msg: str) -> None:
+        if DEBUG:
+            print(cyan(msg))
+
     # Architecture subtrees: CLI group default fills in *new* fields only;
     # checkpoint values win for fields that already exist.
     ARCH_KEYS = {"scene_optimizer", "scene_initializer"}
@@ -390,7 +368,7 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
         if cli_val is None:
             # No direct config path — e.g. +experiment=re10k is a defaults-list override
             # whose effect is already baked into orig_cli_cfg; nothing to apply.
-            print(cyan(f"  Skipping '{key}' (no direct config path in cli)"))
+            _trace(f"  Skipping '{key}' (no direct config path in cli)")
             continue
 
         # For architecture group overrides: fill in missing fields from CLI defaults
@@ -401,7 +379,7 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
             dotkey_parts = set(dotkey.split("."))
             if dotkey_parts & CLI_WINS_SUBKEYS:
                 OmegaConf.update(merged_cfg, dotkey, cli_val, merge=False)
-                print(cyan(f"  '{dotkey}': replace from cli (CLI wins)"))
+                _trace(f"  '{dotkey}': replace from cli (CLI wins)")
                 continue
 
             existing_val = OmegaConf.select(merged_cfg, dotkey, default=None)
@@ -414,15 +392,15 @@ def _apply_cli_overrides(merged_cfg: DictConfig, orig_cli_cfg: DictConfig, raw_o
                     if cli_subval is not None:
                         OmegaConf.set_struct(new_val, False)
                         OmegaConf.update(new_val, subkey, cli_subval, merge=False)
-                        print(cyan(f"  '{dotkey}.{subkey}': CLI override applied (CLI wins)"))
+                        _trace(f"  '{dotkey}.{subkey}': CLI override applied (CLI wins)")
                 OmegaConf.update(merged_cfg, dotkey, new_val, merge=False)
-                print(cyan(f"  '{dotkey}': fill-missing from cli (checkpoint values preserved)"))
+                _trace(f"  '{dotkey}': fill-missing from cli (checkpoint values preserved)")
                 continue
 
         # Group overrides and complex values replace the whole subtree;
         # scalars are merged so sibling keys are preserved.
         replace = is_group_override
-        print(cyan(f"  '{dotkey}': {'replace' if replace else 'update'} from cli"))
+        _trace(f"  '{dotkey}': {'replace' if replace else 'update'} from cli")
         OmegaConf.update(merged_cfg, dotkey, cli_val, merge=not replace)
 
     OmegaConf.set_struct(merged_cfg, True)
@@ -548,12 +526,16 @@ def _merge_test_mode(
 ) -> tuple[DictConfig, DictConfig]:
     """
     Test mode: CLI config is the base for all settings (dataset, test flags, etc.).
-    Only optimizer and initializer *architecture* are patched in from checkpoint configs.
+    The optimizer and initializer *architecture* and the decoder are patched in from the
+    checkpoint config, so the model is rebuilt and rendered the way it was trained. An explicit
+    CLI override still wins for any of these via _apply_cli_overrides.
 
     Initializer source priority:
       1. separate initializer_config_path  (pretrained_initializer ckpt with a config file)
-      2. main loaded_cfg                   (optimizer checkpoint's bundled initializer)
-      3. CLI config as-is                  (pretrained_initializer set but has no config file)
+      2. full pretrained_model checkpoint  (it bundles the initializer weights, so its
+                                            architecture must match the checkpoint config)
+      3. CLI config as-is                  (pretrained_optimizer alone carries no initializer
+                                            weights; or pretrained_initializer has no config file)
 
     Returns (merged_cfg, orig_cli_cfg); orig_cli_cfg is the snapshot taken before any
     checkpoint patches so that _apply_cli_overrides can restore explicit CLI values.
@@ -572,17 +554,28 @@ def _merge_test_mode(
         print(cyan("Test mode: patching scene_trainer.scene_optimizer from checkpoint config."))
         OmegaConf.update(merged_cfg, "scene_trainer.scene_optimizer", optimizer_subcfg, merge=True)
 
+    # Patch decoder from checkpoint so rendering uses the rasterizer the Gaussians were trained
+    # for. Merge over the CLI decoder (checkpoint values win) so any field the checkpoint lacks
+    # is filled from the CLI default -- old checkpoints predate some DecoderCfg fields and there
+    # is no decoder migration. An explicit CLI decoder override still wins via _apply_cli_overrides.
+    decoder_subcfg = OmegaConf.select(loaded_cfg, "scene_trainer.decoder", default=None)
+    if decoder_subcfg is not None:
+        print(cyan("Test mode: patching scene_trainer.decoder from checkpoint config (CLI fills missing fields)."))
+        OmegaConf.update(merged_cfg, "scene_trainer.decoder", decoder_subcfg, merge=True)
+
     # Patch initializer architecture (priority order above)
     if initializer_config_path is not None and initializer_config_path.exists():
         _patch_scene_initializer(merged_cfg, initializer_config_path, context="Test mode")
     elif pretrained_initializer is None:
-        pass
-        # TODO Naama
-        # No explicit initializer checkpoint — fall back to the optimizer checkpoint's initializer
-        # initializer_subcfg = OmegaConf.select(loaded_cfg, "scene_trainer.scene_initializer", default=None)
-        # if initializer_subcfg is not None:
-            # print(cyan("Test mode: patching scene_trainer.scene_initializer from checkpoint config."))
-            # OmegaConf.update(merged_cfg, "scene_trainer.scene_initializer", initializer_subcfg, merge=True)
+        # No separate initializer checkpoint. A full pretrained_model bundles the initializer
+        # weights (load_full_model), so its architecture must come from the checkpoint config.
+        # A pretrained_optimizer checkpoint carries no initializer weights, so the CLI
+        # scene_initializer config is kept as-is.
+        if cli_cfg.checkpointing.pretrained_model is not None:
+            initializer_subcfg = OmegaConf.select(loaded_cfg, "scene_trainer.scene_initializer", default=None)
+            if initializer_subcfg is not None:
+                print(cyan("Test mode: patching scene_trainer.scene_initializer from full-model checkpoint config."))
+                OmegaConf.update(merged_cfg, "scene_trainer.scene_initializer", initializer_subcfg, merge=True)
     else:
         print(cyan("pretrained_initializer set but has no config file; using CLI scene_initializer config."))
 
@@ -677,14 +670,8 @@ def setup_output_dir(cfg, cfg_dict):
         output_dir = CustomPath(output_dir)
         output_dir.mkdir(exist_ok=True, parents=True)
 
-    if HydraConfig.get().mode == RunMode.MULTIRUN and output_dir == "placeholder":
-        # Hack to overcome multirun issues
-        # TODO Naama, need to move to post_init of cfg
-        output_dir = CustomPath(hydra.core.hydra_config.HydraConfig.get()["run"]["dir"])
-        print(cyan(f"Multirun detected, setting output_dir to {CustomPath(output_dir):link}"))
-        # save checkoint path to a file for debugging
-        ckpt_path = cfg.checkpointing.pretrained_model or cfg.checkpointing.pretrained_optimizer
-        (output_dir / "ckpt_dir.txt").write_text(str(ckpt_path))
+    # Single point where output_dir is finalized. Downstream code reads cfg.output_dir; the
+    # value is also written into cfg_dict so the saved config.yaml records it. Keep both in sync.
     cfg_dict.output_dir = output_dir
     cfg.output_dir = output_dir
     output_dir.mkdir(exist_ok=True, parents=True)

optgs/config/dataset/dl3dv.yaml

@@ -4,19 +4,16 @@ defaults:
 
 name: dl3dv
 roots: [datasets/dl3dv]
-make_baseline_1: false
 augment: true
 
 
 image_shape: [270, 480]
 
-baseline_epsilon: 1e-3
 max_fov: 100.0
 
 skip_bad_shape: true
 near: -1.
 far: -1.
-baseline_scale_bounds: false
 shuffle_val: true
 test_len: -1
 test_chunk_interval: 1
@@ -27,35 +24,11 @@ train_times_per_scene: 1
 test_times_per_scene: 1
 ori_image_shape: [270, 480]
 overfit_max_views: 148
-use_index_to_load_chunk: false
 
-mix_tartanair: false
 no_mix_test_set: true
-load_depth: false
-center_pose: false
-
-pose_align_first_view: false
-
-scale_extrinsics: 1.
-metric_scale_align_dl3dv: false
 
 # view filtering
 min_views: 0
 max_views: 0
-highres: false
-
-# mix re10k & dl3dv
-mix_re10k: false
-re10k_min_view_dist: 40
-re10k_max_view_dist: 300
-
-# load remaining context views
-load_remain_context: false
-num_remain_context: 8
-
-# random crop in training
-random_crop: false
-min_size: null
-max_size: null
 
-index_name: index.json
+index_name: index.json

optgs/config/dataset/re10k.yaml

@@ -4,24 +4,16 @@ defaults:
 
 name: re10k
 roots: [datasets/re10k]
-make_baseline_1: false
 augment: true
 
 image_shape: [180, 320]
 highres: false
 
-baseline_epsilon: 1e-3
 max_fov: 100.0
 
 skip_bad_shape: true
 near: -1.
 far: -1.
-baseline_scale_bounds: true
 shuffle_val: true
 test_len: -1
-test_chunk_interval: 1
-
-use_index_to_load_chunk: false
-
-average_pose: false
-center_pose: false
+test_chunk_interval: 1

optgs/config/experiment/re10k_unified.yaml

@@ -63,16 +63,11 @@ dataset:
   roots: [datasets/re10k]
   near: 0.01
   far: 100.
-  baseline_scale_bounds: false
-  make_baseline_1: false
   train_times_per_scene: 1
   highres: false
   scannet: false
   tartanair: false
   load_depth: false
-  pose_align_first_view: false
-  scale_extrinsics: 1.
-  load_remain_context: false
   pose_align_middle_view: false
   overfit_to_scene: null
   opencv_pose_format: false

optgs/config/experiment/train_dl3dv.yaml

@@ -21,7 +21,6 @@ meta_trainer:
   max_steps: 50_000
   val_check_interval: 0.25
   train:
-    l1_loss: true
     depth_smooth_loss_weight: 0.0
   test:
     eval_time_skip_steps: 0
@@ -31,6 +30,8 @@ meta_trainer:
 
 # lpips loss
 loss:
+  mse:
+    l1_loss: true
   lpips:
     apply_after_step: 0
     weight: 0.5
@@ -40,8 +41,6 @@ dataset:
   roots: [ datasets/dl3dv-480p-chunks ]
   near: 0.01
   far: 200.
-  min_size: [ 384,512 ]
-  max_size: [ 512,960 ]
   image_shape: [ 256, 448 ]
   view_sampler:
     num_context_views: 8

optgs/config/experiment/train_l2s_sparse_dl3dv.yaml

@@ -2,8 +2,8 @@
 
 defaults:
   - train_dl3dv
-  - override /meta_trainer/train/replay_buffer_cfg: default
-  - override /loss: [ mse, lpips, deltas ]
+  - override /meta_trainer/train/ckpt_buffer_cfg: default
+  - override /loss: [ mse, lpips, deltas, stability ]
 
 loss:
   mse:
@@ -23,19 +23,24 @@ meta_trainer:
   train:
     loss_on_input_views: true
     loss_on_input_views_num: 4
-    use_replay_buffer: true
+    use_ckpt_buffer: true
+  meta_optimizer:
+    lr: 1e-4
+    lr_monodepth: 0.0
 
 scene_trainer:
   train_scene_opt: true
-  num_update_steps: 4
+  num_update_steps: 2000
+  opt_batch_size: 8
   train_max_refine: 6
   train_min_refine: 1
-
-meta_optimizer:
-  lr: 1e-4
-  lr_monodepth: 0.0
-
+  # Decoder settings for the learn2splat paper weights (the generic gsplat default is
+  # antialiased / eps2d=0.3; these weights are trained with classic rasterization).
+  decoder:
+    rasterize_mode: classic
+    eps2d: 1e-4
 
 checkpointing:
-  pretrained_initializer: checkpoints/optgs/unified-dl3dv-8views/init/checkpoints/epoch_20-step_100000.ckpt  # resplat inititalizer
+  # ReSplat initializer
+  pretrained_initializer: hf://autonomousvision/learn2splat/resplat_init/checkpoints/epoch_20-step_100000.ckpt
   no_strict_load: false

optgs/config/experiment/train_l2s_sparse_dl3dv_no_delta.yaml

@@ -2,7 +2,7 @@
 
 defaults:
   - train_dl3dv
-  - override /meta_trainer/train/replay_buffer_cfg: default
+  - override /meta_trainer/train/ckpt_buffer_cfg: default
   - override /loss: [ mse, lpips ]
 
 loss:
@@ -17,7 +17,10 @@ meta_trainer:
   train:
     loss_on_input_views: true
     loss_on_input_views_num: 4
-    use_replay_buffer: true
+    use_ckpt_buffer: true
+  meta_optimizer:
+    lr: 1e-4
+    lr_monodepth: 0.0
 
 scene_trainer:
   train_scene_opt: true
@@ -25,11 +28,6 @@ scene_trainer:
   train_max_refine: 6
   train_min_refine: 1
 
-meta_optimizer:
-  lr: 1e-4
-  lr_monodepth: 0.0
-
-
 checkpointing:
   pretrained_initializer: checkpoints/optgs/unified-dl3dv-8views/init/checkpoints/epoch_20-step_100000.ckpt  # resplat inititalizer
   no_strict_load: false

optgs/config/experiment/train_l2s_sparse_dl3dv_no_loss.yaml

@@ -2,7 +2,7 @@
 
 defaults:
   - train_dl3dv
-  - override /meta_trainer/train/replay_buffer_cfg: default
+  - override /meta_trainer/train/ckpt_buffer_cfg: default
   - override /loss: [ mse, lpips ]
 
 loss:
@@ -17,7 +17,10 @@ meta_trainer:
   train:
     loss_on_input_views: true
     loss_on_input_views_num: 4
-    use_replay_buffer: true
+    use_ckpt_buffer: true
+  meta_optimizer:
+    lr: 1e-4
+    lr_monodepth: 0.0
 
 scene_trainer:
   train_scene_opt: true
@@ -25,11 +28,6 @@ scene_trainer:
   train_max_refine: 6
   train_min_refine: 1
 
-meta_optimizer:
-  lr: 1e-4
-  lr_monodepth: 0.0
-
-
 checkpointing:
   pretrained_initializer: checkpoints/optgs/unified-dl3dv-8views/init/checkpoints/epoch_20-step_100000.ckpt  # resplat inititalizer
   no_strict_load: false

optgs/config/loss/lpips.yaml

@@ -2,3 +2,4 @@ lpips:
   weight: 0.05
   apply_after_step: 150_000
   perceptual_loss: false
+  half_res: false

optgs/config/loss/mse.yaml

@@ -1,2 +1,4 @@
 mse:
   weight: 1.0
+  l1_loss: false
+  clamp_large_error: 0.0

optgs/config/loss/sgd.yaml

@@ -1,2 +1,4 @@
 sgd:
-  weight: 1.0
+  weight: 1.0
+  l1_loss: false
+  clamp_large_error: 0.0

optgs/config/loss/stability.yaml

@@ -1,2 +1,4 @@
 stability:
   weight: 1.0
+  # Also penalize inputs optimized on per-step view subsets (not fully tested in training).
+  subset_aware: false

optgs/config/main.yaml

@@ -5,7 +5,7 @@ defaults:
   - scene_trainer/scene_optimizer: null
   - scene_trainer/decoder: gsplat
   - meta_trainer/test/postprocessing: none
-  - meta_trainer/train/replay_buffer_cfg: none
+  - meta_trainer/train/ckpt_buffer_cfg: none
 
 wandb:
   project: placeholder
@@ -34,15 +34,6 @@ data_loader:
     batch_size: 1
     seed: 3456
 
-meta_optimizer:
-  lr: 2.e-4
-  lr_monodepth: 2.e-6
-  lr_depth: 0.
-  warm_up_steps: 2000
-  weight_decay: 0.01
-  warm_up_ratio: 0.01
-  adamw_8bit: false
-
 checkpointing:
   load: null
   every_n_train_steps: 1000
@@ -75,25 +66,30 @@ meta_trainer:
   limit_test_batches: 1.0
   limit_train_batches: 1.0
   num_nodes: 1
+  meta_optimizer:
+    lr: 2.e-4
+    lr_monodepth: 2.e-6
+    lr_depth: 0.
+    warm_up_steps: 2000
+    weight_decay: 0.01
+    warm_up_ratio: 0.01
+    adamw_8bit: false
   train:
     depth_mode: null
     extended_visualization: false
     print_log_every_n_steps: 100
     eval_model_every_n_val: 2  # quantitative evaluation every n val
     eval_data_length: 999999
-    eval_deterministic: false
     eval_time_skip_steps: 3
     eval_save_model: true
-    l1_loss: false
     intermediate_loss_weight: 0.9
     no_viz_video: false
     eval_depth: false
-    train_ignore_large_loss: 0.
     no_log_projections: true
     no_log_video: true
     depth_loss_weight: 0.
     log_depth_loss: true
-    depth_smooth_loss_weight: 0.01
+    depth_smooth_loss_weight: 0.  # only meaningful when train_scene_init=True
     depth_smooth_loss_nonorm: false
     depth_smooth_loss_weight_nvs: 0. # for novel views
     monodepth_loss_weight: 0. # for monocular depth loss
@@ -103,21 +99,17 @@ meta_trainer:
     render_depth_loss_weight: 0.
     viz_render_depth: false
     use_gt_depth_range: false
-    depth_range_from_disparity: false
-    max_disparity: 128.
-    min_disparity: 4.
     loss_on_input_views: false
     loss_on_target_views: true
     loss_on_input_views_num: 1
     loss_on_target_views_num: -1
     train_window_size: null
-    half_res_lpips_loss: false
     viz_depth_separate: false
     # L2 weight decay on Gaussian properties (meta-loss)
     scale_l2_loss_weight: 0.
     sh_l2_loss_weight: 0.
     opacity_l2_loss_weight: 0.
-    use_replay_buffer: false
+    use_ckpt_buffer: false
   test:
     output_path: null
     compute_scores: true
@@ -130,22 +122,22 @@ meta_trainer:
     save_gt_image: false
     save_render_depth: false
     save_gt_depth: false
+    save_init_pred_depth: false
     save_error_image: false
     save_video: false
-    save_video_fixed_view: false
-    save_video_fixed_view_index: 0
-    save_video_fixed_view_duplicate: 0
-    save_video_fixed_iteration: false
-    save_video_fixed_iteration_indices: null
-    save_video_fixed_iteration_render_fixed_view: false
-    save_video_combined: false
-    save_video_combined_iterations: null
-    save_video_combined_fixed_iteration_length: 50
+    save_video_optim: false
+    save_video_view_index: 0
+    save_video_frame_repeat: 1
+    save_video_orbit: false
+    save_video_orbit_steps: null
+    save_video_orbit_with_optim: false
+    save_video_optim_orbit: false
+    save_video_optim_orbit_steps: null
+    save_video_orbit_span: 50
     save_gaussian: false
     save_poses: false
     save_cameras_json: true
     save_cameras_npz: true
-    save_point_cloud: false
     render_chunk_size: null
     dec_chunk_size: null
     stablize_camera: false

optgs/config/meta_trainer/train/{replay_buffer_cfg → ckpt_buffer_cfg}/default.yaml

@@ -3,10 +3,10 @@ sample_batch_size: 1
 sample_prob: 0.7
 insert_prob: 0.7
 return_prob: 0.99
-simulate_ahead: true
-simulate_ahead_min_steps: 1
-simulate_ahead_max_steps: 50
-simulate_ahead_grow: 10000
+rollout: true
+rollout_min_steps: 1
+rollout_max_steps: 50
+rollout_grow: 10000
 max_t: null
 push_only_if_not_full: false
 remove_strategy_when_full: oldest

optgs/config/meta_trainer/train/{replay_buffer_cfg → ckpt_buffer_cfg}/none.yaml

@@ -3,10 +3,10 @@ sample_batch_size: 1
 sample_prob: 0.0
 insert_prob: 0.0
 return_prob: 0.0
-simulate_ahead: false
-simulate_ahead_min_steps: 0
-simulate_ahead_max_steps: 0
-simulate_ahead_grow: 0
+rollout: false
+rollout_min_steps: 0
+rollout_max_steps: 0
+rollout_grow: 0
 max_t: null
 push_only_if_not_full: false
 remove_strategy_when_full: oldest

optgs/config/scene_trainer/scene_initializer/edgs.yaml

@@ -7,4 +7,5 @@ sh_degree: 3
 init_opacity: 0.5
 scaling_factor: 0.5
 roma_model_type: outdoors
-sample_init_gaussians: -1
+sample_init_gaussians: -1
+cache_dir: cache/edgs

optgs/config/scene_trainer/scene_initializer/resplat_v1.yaml

@@ -71,19 +71,13 @@ init_pt_with_mv_attn_lowres: false
 pt_head_channels: null
 pt_head_concat_img: false
 pt_head_conv: false
-multi_scale_pt: false
 attn_proj_channels: 64
-fps_num_samples: null
 knn_samples: 16
 post_norm: false
 no_rpe: true
 no_knn_attn: false
 num_blocks: 4
-pt_downsample: 0
-fps_agg_func: attn
-subsample_method: fps
 add_pt_residual: true
-pt_pred_residual_position: false
 
 # freeze depth
 freeze_depth: false
@@ -99,10 +93,6 @@ bilinear_upsample_depth: false
 no_upsample_depth: false
 return_lowres_depth: false
 
-# lvsm gaussian regressor
-lvsm_gaussian_regressor: false
-lvsm_layers: 6
-
 # latent gaussian instead of pixel aligned gaussians
 latent_gs: true
 latent_downsample: 4

optgs/config/scene_trainer/scene_initializer/resplat_v2.yaml

@@ -77,19 +77,13 @@ init_pt_with_mv_attn_lowres: true
 pt_head_channels: null
 pt_head_concat_img: false
 pt_head_conv: false
-multi_scale_pt: false
 attn_proj_channels: 64
-fps_num_samples: null
 knn_samples: 16
 post_norm: false
 no_rpe: true
 no_knn_attn: false
 num_blocks: 6
-pt_downsample: 0
-fps_agg_func: attn
-subsample_method: fps
 add_pt_residual: true
-pt_pred_residual_position: false
 
 # freeze depth
 freeze_depth: false
@@ -105,10 +99,6 @@ bilinear_upsample_depth: false
 no_upsample_depth: false
 return_lowres_depth: false
 
-# lvsm gaussian regressor
-lvsm_gaussian_regressor: false
-lvsm_layers: 6
-
 # latent gaussian instead of pixel aligned gaussians
 latent_gs: true
 latent_downsample: 4

optgs/config/scene_trainer/scene_optimizer/3dgs.yaml

@@ -1,22 +1,11 @@
 defaults:
   - base
   - override refiner: default
+  - override lr_scheduler: expon
 
 name: adam
 
 # Adam optimizer
 betas: [0.9, 0.999]
 eps: 1e-15
-weight_decay: 0.0
-
-# learning rates (gsplat)
-base_lr: 1
-means_lr_init: 1.6e-4
-means_lr_final: 1.6e-6
-means_lr_delay_mult: 1.0
-means_lr_max_steps: 30000  # should be equal to total optimization steps
-scales_lr: 5e-3
-rotations_lr: 1e-3
-opacities_lr: 5e-2
-sh0s_lr: 2.5e-3
-shNs_lr: 1.25e-4
+weight_decay: 0.0

optgs/config/scene_trainer/scene_optimizer/3dgs_star.yaml

@@ -1,6 +1,7 @@
 defaults:
   - base
   - override refiner: default
+  - override lr_scheduler: expon
 
 name: adam
 
@@ -9,14 +10,8 @@ betas: [0.99, 0.999]
 eps: 1e-15
 weight_decay: 0.0
 
-base_lr: 5
-# 3dgs defaults
-means_lr_init: 1.6e-4  # Setting same as final to have constant LR
-means_lr_final: 1.6e-4
-means_lr_delay_mult: 0.0
-means_lr_max_steps: 30000  # should be equal to total optimization steps
-scales_lr: 5e-3
-rotations_lr: 1e-3
-opacities_lr: 5e-2
-sh0s_lr: 2.5e-3
-shNs_lr: 1.25e-4
+# 5x base LR; means held constant (lr_final == _means, so means init 5*1.6e-4 == final 5*1.6e-4).
+lr_scheduler:
+  lr_data:
+    _base: 5
+  lr_final: 1.6e-4

optgs/config/scene_trainer/scene_optimizer/base.yaml

@@ -6,12 +6,12 @@ defaults:
 input_gradients_chunk_size: -1
 
 # iterative refine
-no_refine_mean: false
-no_refine_scale: false
-no_refine_rotation: false
-no_refine_opacity: false
-no_refine_sh0: false
-no_refine_shN: false
+freeze_mean: false
+freeze_scale: false
+freeze_rotation: false
+freeze_opacity: false
+freeze_sh0: false
+freeze_shN: false
 
 zero_state_on_densify: false
 

optgs/config/scene_trainer/scene_optimizer/knn_based.yaml

@@ -5,19 +5,18 @@ name: knn_based
 
 # iterative refine
 no_render_error: false
-refine_sh_only: false
-num_basic_refine_blocks: 4
-num_refine_blocks: 1
+num_basic_blocks: 4
+num_blocks: 1
 concat_init_state: false
 replace_init_state: false
 state_channels: 256
-refine_block_rmsnorm: false
-refine_block_layernorm: false
+block_rmsnorm: false
+block_layernorm: false
 pt_qk_norm: false
 norm_pt_block: false
-refine_gaussian_multiple: 1
-refine_residual_init_state: false
-clamp_refine_max_scale: 3.0
+delta_gaussian_multiple: 1
+residual_init_state: false
+clamp_max_scale: 3.0
 clamp_min_scale: 1e-6
 clamp_min_raw_scales: -1e10
 clamp_max_raw_scales: 1e10
@@ -30,15 +29,14 @@ clamp_min_shs: -1e10
 clamp_max_shs: 1e10
 clamp_shs_soft: false
 
-update_attn_proj_channels: 64
-update_no_knn_attn: false
-update_no_tran_block_norm: false
-update_tran_block_act: gelu
+attn_proj_channels: 64
+no_knn_attn: false
+no_tran_block_norm: false
+tran_block_act: gelu
 multi_gaussian_scale_smaller: false
 init_gaussian_multiple: 1
-refine_condition_pt_feature: true
-reinit_gaussian_when_refine_multiple: false
-refine_same_num_points: false
+condition_pt_feature: true
+same_num_points: false
 input_error_rgb_no_shuffle: false
 input_error_cache_resnet_feature: false
 
@@ -50,20 +48,19 @@ init_state_scale: 0
 pt_heads: 1
 
 # refine with mv attention
-refine_with_mv_attn: false
-refine_with_mv_attn_lowres: false
-refine_no_mv_attn: false
+with_mv_attn: false
+with_mv_attn_lowres: false
+no_mv_attn: false
 mv_attn_conv_with_norm: false
-refine_mv_shuffle_attn: false
-refine_mv_attn_with_pos_enc: false
-refine_shuffle_attn_no_norm: false
-refine_mv_unimatch_attn: false
-refine_knn_samples: 16
-refine_multi_scale_pt: false
+mv_shuffle_attn: false
+mv_attn_with_pos_enc: false
+shuffle_attn_no_norm: false
+mv_unimatch_attn: false
+knn_samples: 16
+max_active_gaussians: 100_000
 
 # KNN
 use_fused_attn: true  # fused KNN gather + attention CUDA kernel (faster, less memory)
-prune_invisible_gaussians: false
 knn_idx_update_every: 1
 
 # inputs
@@ -103,25 +100,17 @@ input_gradient: false
 input_gradient_log: false
 input_gradient_log_clip_deltas: 0.001
 input_gradient_scale: 1.
-gradient_update_scale: 1.
+residual_grad_scale: 1.
 input_gradient_with_ssim_loss: false
 input_gradient_same_loss: false
 input_gradient_loss_reduction: mean
 scale_residual_grads: false
 
-window_local_refine: false
-window_global_refine: false
-window_local_global_refine: false
-
-# sliding window update to save training memory
-update_window_size: 0
-local_gaussian_render: false
-
-train_global_update_only: false
+train_global_only: false
 
 # random size refine
 # update more for low resolution, less for high
-random_update_with_size: false
+random_step_with_size: false
 
 
 use_amp: true
@@ -131,9 +120,6 @@ pt_update_amp: true
 use_checkpointing: false
 recurrent_use_checkpointing: false
 
-# Debugging
-debug_refine_update_module: true
-
 # Normalizing input
 input_gradient_normalize: false
 input_gradient_normalize_type: layer
@@ -146,21 +132,20 @@ predict_state_scale: false
 predict_state_scale_norm: false
 
 # Update head
-update_head_concat_img: false
-update_head_layer_num: 2
-update_head_act: gelu
-update_head_final_act: identity
-update_head_hidden_dim_matches: "input" # rebuttal version. switch to "output" for submission version
+delta_head_concat_img: false
+delta_head_layer_num: 2
+delta_head_act: gelu
+delta_head_final_act: identity
+delta_head_hidden_dim_matches: "input" # rebuttal version. switch to "output" for submission version
 
-update_head_scale_mag: false
-update_head_scalar_scale: false
-update_head_scalar_scale_act: relu
+delta_head_scalar_scale: false
+delta_head_scalar_scale_act: relu
 
 # Per-parameter-group heads (Feature A): separate heads per param group, each with own normalize+scale
-update_head_per_param_heads: false
-update_head_per_param_hidden_dim: 48  # tuned so total params ≈ baseline head (~81K)
-# Per-parameter scalar scales (Feature B): per-group scalar scales (requires update_head_scalar_scale=true)
-update_head_per_param_scales: false
+delta_head_per_param_heads: false
+delta_head_per_param_hidden_dim: 48  # tuned so total params ≈ baseline head (~81K)
+# Per-parameter scalar scales (Feature B): per-group scalar scales (requires delta_head_scalar_scale=true)
+delta_head_per_param_scales: false
 
 opt_scales_before_act: false
 

optgs/config/scene_trainer/scene_optimizer/learn2splat.yaml

@@ -5,13 +5,25 @@ name: l2s
 
 # General optimization settings
 opt_scales_before_act: true
+# Scales are refined in log space, so clamp in log space too.
+clamp_min_raw_scales: -8.0
+clamp_max_raw_scales: 1.0986  # ln(3.0)
 sh_d: 16  # should be the default, but just in case
 
+# SH coefficients are clamped to a fixed range during refinement.
+clamp_min_shs: -2.0
+clamp_max_shs: 2.0
+
+# KNN neighborhood size for the point-transformer attention.
+knn_samples: 4
+
 # Input gradient settings
 input_gradient: true
 input_gradient_normalize: true
 input_gradient_normalize_type: adam
 input_gradient_with_ssim_loss: true
+# Average the input gradient over pixels, then sum over the supervised views.
+input_gradient_loss_reduction: mean_pixels_sum_views
 
 # Freeze zero-grad gaussians ❄️($G_{∇=0})
 update_only_nonzero_grad: true
@@ -20,8 +32,8 @@ update_only_nonzero_grad: true
 predict_state_scale: true
 
 # Delta scale
-update_head_scalar_scale: true
-update_head_scalar_scale_act: relu  # should be the default, but just in case
+delta_head_scalar_scale: true
+delta_head_scalar_scale_act: relu  # should be the default, but just in case
 
 
 

optgs/config/scene_trainer/scene_optimizer/learn2splat_dense.yaml

@@ -0,0 +1,8 @@
+defaults:
+  - learn2splat
+
+# Dense variant (64-view SfM/COLMAP init): the optimizer state is initialized as random
+# features rather than the sparse model's constant zero state.
+init_state_wo_features: true
+init_state_type: random
+init_state_scale: 1.0

optgs/config/scene_trainer/scene_optimizer/lr_scheduler/expon.yaml

@@ -0,0 +1,28 @@
+defaults:
+  - base
+
+name: expon
+
+# 3DGS learning rates (gsplat); means decays, the rest are constant
+lr_data:
+  _base: 1
+  _means: 1.6e-4
+  _scales: 5e-3
+  _quats: 1e-3
+  _opacities: 5e-2
+  _sh0: 2.5e-3
+  _shN: 1.25e-4
+apply_scheduler:
+  _base: true   # gates all params (Bool3DGSCfg.<param> = _base AND _param); must stay true
+  _means: true
+  _scales: false
+  _quats: false
+  _opacities: false
+  _sh0: false
+  _shN: false
+
+# means decay: 1.6e-4 -> 1.6e-6 over the full step budget
+lr_final: 1.6e-6
+lr_delay_steps: 0
+lr_delay_mult: 1.0
+max_steps: 30000  # should be equal to total optimization steps

optgs/config/scene_trainer/scene_optimizer/refiner/mcmc.yaml

@@ -39,7 +39,7 @@ fallback_means_lr: 1.6e-4
 relocate_copy_state: true  # inherit alive Gaussian's optimizer state (better than zeroing)
 
 # Cap (in scale units) applied to the *noise-only* covariance computation. Does NOT change the
-# rendered Gaussian scales. Needed for knn_based whose network saturates clamp_refine_max_scale,
+# rendered Gaussian scales. Needed for knn_based whose network saturates clamp_max_scale,
 # which produces covariances 10²-10⁴× larger than vanilla's and makes the noise overflow the
 # renderer's tile-binning math (silent CUDA OOB downstream). Tune to ~scene_scale / 5.
 noise_scale_cap: 1.0

optgs/config/scene_trainer/scene_optimizer/resplat_v2.yaml

@@ -5,7 +5,7 @@ name: resplat_v2
 input_error: true
 input_error_mv_attn: true
 input_error_add_rgb_feature: true
-refine_knn_samples: 8
+knn_samples: 8
 state_channels: 512
 residual_state: true
-update_head_layer_num: 4
+delta_head_layer_num: 4

optgs/config/scene_trainer/scene_optimizer/sgd.yaml

@@ -1,22 +0,0 @@
-defaults:
-  - base
-  - override refiner: default
-
-name: sgd
-
-# Adam optimizer
-betas: [0.9, 0.999]
-eps: 1e-15
-weight_decay: 0.0
-
-# learning rates (gsplat)
-base_lr: 1
-means_lr_init: 1.6e-4
-means_lr_final: 1e-5
-means_lr_delay_mult: 0.01
-means_lr_max_steps: 30000  # should be equal to total optimization steps
-scales_lr: 5e-3
-rotations_lr: 1e-3
-opacities_lr: 5e-2
-sh0s_lr: 2.5e-3
-shNs_lr: 1.25e-4

optgs/config_migrate.py

@@ -1,7 +1,9 @@
+from math import log
+
 from omegaconf import OmegaConf
 
 
-CURRENT_CFG_VERSION = 2
+CURRENT_CFG_VERSION = 2.5
 
 def migrate(cfg_dict):
     was_omega = not isinstance(cfg_dict, dict)
@@ -27,6 +29,31 @@ def migrate(cfg_dict):
         cfg_dict = migrate_v1_to_v2(cfg_dict)
         version = 2
 
+    if version == 2:
+        print("Migrating config from version 2 to version 2.1 (strip update_/refine_ prefixes)...")
+        cfg_dict = migrate_v2_to_v2_1(cfg_dict)
+        version = 2.1
+
+    if version == 2.1:
+        print("Migrating config from version 2.1 to version 2.2 (single-space scale clamp)...")
+        cfg_dict = migrate_v2_1_to_v2_2(cfg_dict)
+        version = 2.2
+
+    if version == 2.2:
+        print("Migrating config from version 2.2 to version 2.3 (replay_buffer -> ckpt_buffer, simulate_ahead -> rollout)...")
+        cfg_dict = migrate_v2_2_to_v2_3(cfg_dict)
+        version = 2.3
+
+    if version == 2.3:
+        print("Migrating config from version 2.3 to version 2.4 (adam flat LRs -> expon lr_scheduler)...")
+        cfg_dict = migrate_v2_3_to_v2_4(cfg_dict)
+        version = 2.4
+
+    if version == 2.4:
+        print("Migrating config from version 2.4 to version 2.5 (meta_optimizer moved under meta_trainer)...")
+        cfg_dict = migrate_v2_4_to_v2_5(cfg_dict)
+        version = 2.5
+
     if version != CURRENT_CFG_VERSION:
         raise ValueError(f"Unsupported config version: {version}")
 
@@ -43,16 +70,95 @@ def migrate(cfg_dict):
         si["name"] = "resplat_v1"
 
     # Strip stale postprocessing fields from old checkpoint configs
-    pp = cfg_container.get("meta_trainer", {}).get("test", {}).get("postprocessing", None)
+    te = cfg_container.get("meta_trainer", {}).get("test", {})
+    pp = te.get("postprocessing", None) if isinstance(te, dict) else None
     if isinstance(pp, dict):
         pp.pop("__target__", None)
         pp.pop("enabled", None)
         pp.pop("lr", None)
 
+    # save_point_cloud removed (the flag was never read by any code path).
+    if isinstance(te, dict):
+        te.pop("save_point_cloud", None)
+
+    # Video flags renamed (fixed_view->optim, fixed_iteration->orbit, combined->optim_orbit).
+    # Test flags come from the current run's config, so just drop the stale keys from old configs.
+    for k in ("save_video_fixed_view", "save_video_fixed_view_index", "save_video_fixed_view_duplicate",
+              "save_video_fixed_iteration", "save_video_fixed_iteration_indices",
+              "save_video_fixed_iteration_render_fixed_view", "save_video_combined",
+              "save_video_combined_iterations", "save_video_combined_fixed_iteration_length"):
+        if isinstance(te, dict):
+            te.pop(k, None)
+
     # Strip stale foundationstereo fields (encoder removed)
     si.pop("foundationstereo", None)
     si.pop("fstereo_num_refine", None)
 
+    # Strip removed optimizer sliding-window fields (feature removed)
+    for k in ("window_size", "update_window_size", "local_gaussian_render",
+              "window_local_refine", "window_global_refine", "window_local_global_refine"):
+        so.pop(k, None)
+
+    # delta_head_scale_mag removed (experimental). Drop it under all historical names.
+    for k in ("delta_head_scale_mag", "update_head_scale_mag", "refine_output_scale_mag"):
+        so.pop(k, None)
+
+    # reinit_gaussian_when_multiple removed (reinit path was never implemented). Drop it
+    # under all historical names (reinit_gaussian_when_refine_multiple is the pre-2.1 name).
+    for k in ("reinit_gaussian_when_multiple", "reinit_gaussian_when_refine_multiple"):
+        so.pop(k, None)
+
+    # The initializer has no point-downsampling path, so pt_downsample,
+    # fps_agg_func and subsample_method are not valid config keys; drop them.
+    # multi_scale_pt and fps_num_samples are likewise not config keys (the point
+    # transformer is always PlainPointTransformer). subsample_method also appears
+    # on old optimizer configs, so strip it from there too.
+    for k in ("multi_scale_pt", "refine_multi_scale_pt", "subsample_method"):
+        so.pop(k, None)
+    for k in ("multi_scale_pt", "fps_num_samples", "pt_downsample", "fps_agg_func", "subsample_method"):
+        si.pop(k, None)
+
+    # The lvsm gaussian-regressor path was removed (disabled in every config), so
+    # lvsm_gaussian_regressor and lvsm_layers are no longer valid initializer keys.
+    for k in ("lvsm_gaussian_regressor", "lvsm_layers"):
+        si.pop(k, None)
+
+    # pt_pred_residual_position removed (residual-position prediction was disabled in every config).
+    si.pop("pt_pred_residual_position", None)
+
+    # sh_only removed: it was "freeze everything but SH", now expressed via per-group freeze_*.
+    # (refine_sh_only is the pre-2.1 name.)
+    sh_only = so.pop("sh_only", None)
+    if sh_only is None:
+        sh_only = so.pop("refine_sh_only", None)
+    if sh_only:
+        for fk in ("freeze_mean", "freeze_scale", "freeze_rotation", "freeze_opacity"):
+            so[fk] = True
+
+    # l1_loss / train_ignore_large_loss / half_res_lpips_loss moved off train into the per-loss
+    # cfgs: mse/sgd take l1_loss + clamp_large_error, lpips takes half_res.
+    train = cfg_container.get("meta_trainer", {}).get("train", {})
+
+    # depth_range_from_disparity removed (disabled in every config); its only consumers,
+    # max_disparity and min_disparity, are removed with it.
+    for k in ("depth_range_from_disparity", "max_disparity", "min_disparity"):
+        train.pop(k, None)
+
+    old_l1 = train.pop("l1_loss", None)
+    old_clamp = train.pop("train_ignore_large_loss", None)
+    old_half_res = train.pop("half_res_lpips_loss", None)
+    for loss_wrapper in cfg_container.get("loss", []) or []:
+        if not isinstance(loss_wrapper, dict):
+            continue
+        for name, lcfg in loss_wrapper.items():
+            if not isinstance(lcfg, dict):
+                continue
+            if name in ("mse", "sgd"):
+                lcfg.setdefault("l1_loss", old_l1 if old_l1 is not None else False)
+                lcfg.setdefault("clamp_large_error", old_clamp if old_clamp is not None else 0.0)
+            elif name == "lpips":
+                lcfg.setdefault("half_res", old_half_res if old_half_res is not None else False)
+
     if was_omega:
         return OmegaConf.create(cfg_container)
     return cfg_container
@@ -74,6 +180,183 @@ def migrate_v1_to_v2(cfg_dict):
     return cfg
 
 
+def migrate_v2_4_to_v2_5(cfg_dict):
+    """Migration from v2.4 to v2.5: move top-level 'meta_optimizer' under 'meta_trainer'."""
+    cfg = OmegaConf.to_container(cfg_dict, resolve=False) if not isinstance(cfg_dict, dict) else dict(cfg_dict)
+
+    if "meta_optimizer" in cfg:
+        meta_trainer = cfg.setdefault("meta_trainer", {})
+        if "meta_optimizer" not in meta_trainer:
+            meta_trainer["meta_optimizer"] = cfg.pop("meta_optimizer")
+
+    cfg["version"] = 2.5
+    return cfg
+
+
+def migrate_v2_to_v2_1(cfg_dict):
+    """
+    Migration from v2 to v2.1: strip redundant update_/refine_ prefixes from optimizer
+    config fields (the class is already named Optimizer), rename no_refine_* -> freeze_*,
+    and update_head_* -> delta_head_*.
+    """
+    cfg = OmegaConf.to_container(cfg_dict, resolve=False) if not isinstance(cfg_dict, dict) else dict(cfg_dict)
+
+    so = cfg.get("scene_trainer", {}).get("scene_optimizer", {})
+
+    RENAME_MAP = {
+        # refine_
+        "num_basic_refine_blocks": "num_basic_blocks",
+        "num_refine_blocks": "num_blocks",
+        "refine_block_rmsnorm": "block_rmsnorm",
+        "refine_block_layernorm": "block_layernorm",
+        "refine_gaussian_multiple": "delta_gaussian_multiple",
+        "refine_residual_init_state": "residual_init_state",
+        "clamp_refine_max_scale": "clamp_max_scale",
+        "refine_condition_pt_feature": "condition_pt_feature",
+        "refine_same_num_points": "same_num_points",
+        "refine_knn_samples": "knn_samples",
+        "refine_with_mv_attn": "with_mv_attn",
+        "refine_with_mv_attn_lowres": "with_mv_attn_lowres",
+        "refine_no_mv_attn": "no_mv_attn",
+        "refine_mv_shuffle_attn": "mv_shuffle_attn",
+        "refine_mv_attn_with_pos_enc": "mv_attn_with_pos_enc",
+        "refine_shuffle_attn_no_norm": "shuffle_attn_no_norm",
+        "refine_mv_unimatch_attn": "mv_unimatch_attn",
+        # no_refine_ -> freeze_
+        "no_refine_mean": "freeze_mean",
+        "no_refine_scale": "freeze_scale",
+        "no_refine_rotation": "freeze_rotation",
+        "no_refine_opacity": "freeze_opacity",
+        "no_refine_sh0": "freeze_sh0",
+        "no_refine_shN": "freeze_shN",
+        # other update_
+        "update_attn_proj_channels": "attn_proj_channels",
+        "update_no_knn_attn": "no_knn_attn",
+        "update_no_tran_block_norm": "no_tran_block_norm",
+        "update_tran_block_act": "tran_block_act",
+        "train_global_update_only": "train_global_only",
+        "random_update_with_size": "random_step_with_size",
+        "gradient_update_scale": "residual_grad_scale",
+        # update_head_ -> delta_head_
+        "update_head_layer_num": "delta_head_layer_num",
+        "update_head_concat_img": "delta_head_concat_img",
+        "update_head_act": "delta_head_act",
+        "update_head_final_act": "delta_head_final_act",
+        "update_head_hidden_dim_matches": "delta_head_hidden_dim_matches",
+        "update_head_scalar_scale": "delta_head_scalar_scale",
+        "update_head_scalar_scale_act": "delta_head_scalar_scale_act",
+        "update_head_per_param_heads": "delta_head_per_param_heads",
+        "update_head_per_param_hidden_dim": "delta_head_per_param_hidden_dim",
+        "update_head_per_param_scales": "delta_head_per_param_scales",
+    }
+
+    for old, new in RENAME_MAP.items():
+        if old in so:
+            so[new] = so.pop(old)
+
+    cfg["version"] = 2.1
+    return cfg
+
+
+def migrate_v2_1_to_v2_2(cfg_dict):
+    """
+    Migration from v2.1 to v2.2: scales are now clamped in a single space.
+
+    When opt_scales_before_act is set, scales are refined in log space, and the clamp is now
+    applied there (raw limits) both when unpacking and in the per-step update. Earlier configs
+    instead clamped once in activation space (clamp_min/max_scale) during unpack and once in log
+    space (clamp_min/max_raw_scales) during the update, so the scale was effectively bounded by
+    the intersection of the two. Fold that intersection into the raw limits so old checkpoints
+    keep the same effective bounds under the new single-space clamp.
+    """
+    cfg = OmegaConf.to_container(cfg_dict, resolve=False) if not isinstance(cfg_dict, dict) else dict(cfg_dict)
+
+    so = cfg.get("scene_trainer", {}).get("scene_optimizer", {})
+    if so.get("opt_scales_before_act", False):
+        min_scale = float(so.get("clamp_min_scale", 1e-6))
+        max_scale = float(so.get("clamp_max_scale", 3.0))
+        min_raw = float(so.get("clamp_min_raw_scales", -1e10))
+        max_raw = float(so.get("clamp_max_raw_scales", 1e10))
+        so["clamp_min_raw_scales"] = max(log(min_scale), min_raw)
+        so["clamp_max_raw_scales"] = min(log(max_scale), max_raw)
+
+    cfg["version"] = 2.2
+    return cfg
+
+
+def migrate_v2_2_to_v2_3(cfg_dict):
+    """
+    Migration from v2.2 to v2.3: the replay-buffer feature is renamed to ckpt_buffer and
+    simulate_ahead to rollout (use_replay_buffer -> use_ckpt_buffer, replay_buffer_cfg ->
+    ckpt_buffer_cfg, and the simulate_ahead* fields inside it -> rollout*).
+    """
+    cfg = OmegaConf.to_container(cfg_dict, resolve=False) if not isinstance(cfg_dict, dict) else dict(cfg_dict)
+
+    train = cfg.get("meta_trainer", {}).get("train", {})
+    if "use_replay_buffer" in train:
+        train["use_ckpt_buffer"] = train.pop("use_replay_buffer")
+    if "replay_buffer_cfg" in train:
+        train["ckpt_buffer_cfg"] = train.pop("replay_buffer_cfg")
+
+    buffer_cfg = train.get("ckpt_buffer_cfg", {})
+    if isinstance(buffer_cfg, dict):
+        ROLLOUT_RENAME_MAP = {
+            "simulate_ahead": "rollout",
+            "simulate_ahead_min_steps": "rollout_min_steps",
+            "simulate_ahead_max_steps": "rollout_max_steps",
+            "simulate_ahead_grow": "rollout_grow",
+        }
+        for old, new in ROLLOUT_RENAME_MAP.items():
+            if old in buffer_cfg:
+                buffer_cfg[new] = buffer_cfg.pop(old)
+
+    cfg["version"] = 2.3
+    return cfg
+
+
+def migrate_v2_3_to_v2_4(cfg_dict):
+    """
+    Migration from v2.3 to v2.4: the Adam baseline (name adam/sgd) no longer carries flat per-param
+    LR fields. They move into the inherited lr_scheduler cfg (name "expon"): per-param values into
+    lr_data, the means decay into the expon scheduler params, with apply_scheduler enabling the
+    schedule on means only.
+
+    The expon scheduler scales the whole curve by lr_data._base (= old base_lr) and the means LR is
+    additionally scaled by scene extent at runtime, so the bounds reproduce the old behavior:
+    old means LR = base_lr * scene_scale * expon(means_lr_init -> means_lr_final).
+    """
+    cfg = OmegaConf.to_container(cfg_dict, resolve=False) if not isinstance(cfg_dict, dict) else dict(cfg_dict)
+
+    so = cfg.get("scene_trainer", {}).get("scene_optimizer", {})
+    if so.get("name") in ("adam", "sgd") and "means_lr_init" in so:
+        base = so.pop("base_lr", 1)
+        lr_scheduler = so.get("lr_scheduler", {}) or {}
+        lr_scheduler["name"] = "expon"
+        lr_scheduler["lr_data"] = {
+            "_base": base,
+            "_means": so.pop("means_lr_init"),
+            "_scales": so.pop("scales_lr"),
+            "_quats": so.pop("rotations_lr"),
+            "_opacities": so.pop("opacities_lr"),
+            "_sh0": so.pop("sh0s_lr"),
+            "_shN": so.pop("shNs_lr"),
+        }
+        # _base gates all params (Bool3DGSCfg.<param> = _base AND _param), so it must stay True;
+        # the schedule is enabled per-param via the individual flags (means only).
+        lr_scheduler["apply_scheduler"] = {
+            "_base": True, "_means": True, "_scales": False,
+            "_quats": False, "_opacities": False, "_sh0": False, "_shN": False,
+        }
+        lr_scheduler["lr_final"] = so.pop("means_lr_final")
+        lr_scheduler["lr_delay_steps"] = 0
+        lr_scheduler["lr_delay_mult"] = so.pop("means_lr_delay_mult", 1.0)
+        lr_scheduler["max_steps"] = so.pop("means_lr_max_steps", 30000)
+        so["lr_scheduler"] = lr_scheduler
+
+    cfg["version"] = 2.4
+    return cfg
+
+
 def migrate_v0_to_v1(cfg):
     """
     Migration from submission v0 (refine_*) to rebuttal v1 (input_error_*).
@@ -159,7 +442,6 @@ def migrate_v0_to_v1(cfg):
 
         # update head
         "final_head_act": "update_head_final_act",
-        "refine_output_scale_mag": "update_head_scale_mag",
         "scalar_scale_out": "update_head_scalar_scale",
         "scalar_scale_out_act": "update_head_scalar_scale_act",
 

optgs/dataset/camera_datasets/camera.py

@@ -10,8 +10,6 @@ import os
 import json
 
 from optgs.geometry.projection import get_fov, get_projection_matrix
-from optgs.visualization.camera_trajectory.wobble import generate_wobble_transformation
-from optgs.visualization.camera_trajectory.interpolation import interpolate_extrinsics, interpolate_intrinsics
 
 
 def get_scene_scale(camtoworlds: Float[np.ndarray, "N 4 4"]) -> float:
@@ -143,133 +141,6 @@ class Camera(nn.Module):
                 indent=4,
             )
 
-    @classmethod
-    def load_camera(cls, cam_dir: Path, data_device: torch.device):
-        extrinsics = torch.load(cam_dir / "extrinsics.pt")
-        intrinsics = torch.load(cam_dir / "intrinsics.pt")
-        image = torch.load(cam_dir / "image.pt")
-
-        if (cam_dir / "gt_alpha_mask.pt").exists():
-            gt_alpha_mask = torch.load(cam_dir / "gt_alpha_mask.pt")
-        else:
-            gt_alpha_mask = None
-
-        with open(cam_dir / "cam_info.json", "r") as f:
-            cam_info = json.load(f)
-
-        return cls(
-            colmap_id=cam_info["colmap_id"],
-            extrinsics=extrinsics.to(data_device),
-            intrinsics=intrinsics.to(data_device),
-            image=image.to(data_device),
-            gt_alpha_mask=gt_alpha_mask.to(data_device) if gt_alpha_mask is not None else None,
-            raw_image_shape=tuple(cam_info["raw_image_shape"]),
-            image_name=cam_info["image_name"],
-            uid=cam_info["uid"],
-            near=torch.Tensor([cam_info["near"]]).to(data_device),
-            far=torch.Tensor([cam_info["far"]]).to(data_device),
-            data_device=data_device,
-        ).to(data_device)
-
-
-def generate_cam_params_for_wobble(t: Tensor, cam_a: Camera, cam_b: Camera):
-    origin_a = cam_a.extrinsics[:3, 3]
-    origin_b = cam_b.extrinsics[:3, 3]
-    cam_a_extrinsics = cam_a.extrinsics
-    cam_b_extrinsics = cam_b.extrinsics
-    cam_a_intrinsics = cam_a.intrinsics
-    cam_b_intrinsics = cam_b.intrinsics
-
-    delta = (origin_a - origin_b).norm(dim=-1)
-    
-    tf = generate_wobble_transformation(
-        radius=delta * 0.5,
-        t=t,
-        num_rotations=1,
-        scale_radius_with_t=False,
-    )
-
-    extrinsics = interpolate_extrinsics(
-        initial=cam_a_extrinsics,
-        final=cam_b_extrinsics,
-        t=(t - 2),
-    )
-    intrinsics = interpolate_intrinsics(
-        initial=cam_a_intrinsics,
-        final=cam_b_intrinsics,
-        t=(t - 2),
-    )
-    return extrinsics @ tf, intrinsics
-
-
-def generate_cam_params_for_interpolation(t: Tensor, cam_a: Camera, cam_b: Camera):
-    cam_a_extrinsics = cam_a.extrinsics
-    cam_a_extrinsics_render_view = cam_a.extrinsics_render_view
-    cam_b_extrinsics = cam_b.extrinsics
-    cam_b_extrinsics_render_view = cam_b.extrinsics_render_view
-    cam_a_intrinsics = cam_a.intrinsics
-    cam_a_intrinsics_render_view = cam_a.intrinsics_render_view
-    cam_b_intrinsics = cam_b.intrinsics
-    cam_b_intrinsics_render_view = cam_b.intrinsics_render_view
-
-    extrinsics = interpolate_extrinsics(
-        initial=cam_a_extrinsics,
-        final=cam_b_extrinsics,
-        t=(t - 2),
-    )
-    intrinsics = interpolate_intrinsics(
-        initial=cam_a_intrinsics,
-        final=cam_b_intrinsics,
-        t=(t - 2),
-    )
-    extrinsics_render_view = interpolate_extrinsics(
-        initial=cam_a_extrinsics_render_view,
-        final=cam_b_extrinsics_render_view,
-        t=(t - 2),
-    )
-    intrinsics_render_view = interpolate_intrinsics(
-        initial=cam_a_intrinsics_render_view,
-        final=cam_b_intrinsics_render_view,
-        t=(t - 2),
-    )
-    return extrinsics, intrinsics, extrinsics_render_view, intrinsics_render_view
-
-
-def get_intermediate_cameras(cam_a: Camera, cam_b: Camera, num_frames: int = 150, smooth: bool = False):
-    t = torch.linspace(0, 1, num_frames, dtype=torch.float32, device=cam_a.data_device)
-    if smooth: t = (torch.cos(torch.pi * (t + 1)) + 1) / 2
-    
-    extrinsics, intrinsics, extrinsics_render_view, intrinsics_render_view = (
-        generate_cam_params_for_interpolation(t, cam_a, cam_b)
-    )
-    extrinsics = extrinsics.squeeze(0)
-    intrinsics = intrinsics.squeeze(0)
-    extrinsics_render_view = extrinsics_render_view.squeeze(0)
-    intrinsics_render_view = intrinsics_render_view.squeeze(0)
-
-    cameras = [
-        Camera(
-            colmap_id=cam_a.colmap_id,
-            image_name=f"{cam_a.image_name}_{index:04d}",
-            uid=index,
-            near=cam_a.znear,
-            far=cam_a.zfar,
-            data_device=cam_a.data_device,
-            image=cam_a.original_image,     # These views have no ground truth image but we should never require images for mesh views
-            raw_image_shape=cam_a.raw_image_shape,
-            extrinsics=extrinsics[index],
-            intrinsics=intrinsics[index],
-            extrinsics_render_view=extrinsics_render_view[index],
-            intrinsics_render_view=intrinsics_render_view[index],
-            scale_matrix=cam_a.scale_matrix,
-            trans_matrix=cam_a.trans_matrix,
-            gt_alpha_mask=None
-        )
-        for index in range(num_frames)
-    ]
-    return cameras
-
-
 def patch_shim(cams: list[Camera], patch_size: int) -> list[Camera]:
     new_cams = []
 
@@ -319,18 +190,6 @@ def patch_shim(cams: list[Camera], patch_size: int) -> list[Camera]:
     return new_cams
 
 
-def calculate_cameras_extent(cam_centers: Tensor):
-    avg_cam_center = cam_centers.mean(dim=0, keepdim=True)
-    dist = torch.norm(cam_centers - avg_cam_center, dim=-1, keepdim=True)
-    diagonal = dist.max()
-
-    center = avg_cam_center.flatten()
-    radius = diagonal * 1.1
-
-    translate = -center
-    return translate, radius.item()
-        
-
 def save_cameras(cameras: list[Camera], save_dir: Path):
     os.makedirs(save_dir, exist_ok=True)
 

optgs/dataset/data_module.py

@@ -14,15 +14,15 @@ from .validation_wrapper import ValidationWrapper
 from ..misc.step_tracker import StepTracker
 
 
-def get_data_shim(encoder: nn.Module) -> DataShim:
+def get_data_shim(initializer: nn.Module) -> DataShim:
     """Get functions that modify the batch. It's sometimes necessary to modify batches
     outside the data loader because GPU computations are required to modify the batch or
     because the modification depends on something outside the data loader.
     """
 
     shims: list[DataShim] = []
-    if hasattr(encoder, "get_data_shim"):
-        shims.append(encoder.get_data_shim())
+    if hasattr(initializer, "get_data_shim"):
+        shims.append(initializer.get_data_shim())
 
     def combined_shim(batch):
         for shim in shims:

optgs/dataset/dataset_dl3dv.py

@@ -12,9 +12,6 @@ from jaxtyping import Float, UInt8
 from PIL import Image
 from torch import Tensor
 from torch.utils.data import IterableDataset
-import numpy as np
-import os
-import random
 
 from ..geometry.projection import get_fov
 from .dataset import DatasetCfgCommon
@@ -28,47 +25,24 @@ from .view_sampler import ViewSampler
 class DatasetDL3DVCfg(DatasetCfgCommon):
     name: Literal["dl3dv"]
     roots: list[Path]
-    baseline_epsilon: float
     max_fov: float
-    make_baseline_1: bool
     augment: bool
     test_len: int
     test_chunk_interval: int
     train_times_per_scene: int
     test_times_per_scene: int
     ori_image_shape: list[int]
-    # random crop training
-    random_crop: bool
-    max_size: list[int] | None
-    min_size: list[int] | None
 
     skip_bad_shape: bool = True
     near: float = -1.0
     far: float = -1.0
-    baseline_scale_bounds: bool = True
     shuffle_val: bool = True
     no_mix_test_set: bool = True
-    load_depth: bool = False
     min_views: int = 0
     max_views: int = 0
-    highres: bool = False
     sort_target_index: Optional[bool] = False
     overfit_max_views: Optional[int] = None
     sort_context_index: Optional[bool] = False
-    use_index_to_load_chunk: Optional[bool] = False
-    pose_align_first_view: bool = False  # align the camera pose to the first view
-    scale_extrinsics: float = 1.
-    metric_scale_align_dl3dv: bool = False
-    center_pose: bool = False  # center and normalize the pose by the distance to the center
-
-    # mix re10k & dl3dv
-    mix_re10k: bool = False
-    re10k_min_view_dist: int = 40
-    re10k_max_view_dist: int = 300
-
-    # load remaining context views
-    load_remain_context: bool = False
-    num_remain_context: int = 8
 
     index_name: str = "index.json"
 
@@ -104,14 +78,9 @@ class DatasetDL3DV(IterableDataset):
         self.chunks = []
         for i, root in enumerate(cfg.roots):
             root = root / self.data_stage
-            if self.cfg.use_index_to_load_chunk:
-                with open(root / self.cfg.index_name, "r") as f:
-                    json_dict = json.load(f)
-                root_chunks = sorted(list(set(json_dict.values())))
-            else:
-                root_chunks = sorted(
-                    [path for path in root.iterdir() if path.suffix == ".torch"]
-                )
+            root_chunks = sorted(
+                [path for path in root.iterdir() if path.suffix == ".torch"]
+            )
 
             # mixed data training only evaluate on a single test set
             if cfg.no_mix_test_set and self.data_stage in ['val', 'test'] and i > 0:
@@ -133,17 +102,11 @@ class DatasetDL3DV(IterableDataset):
         if self.stage == "val":
             self.chunks = self.chunks * int(1e6 // len(self.chunks))
 
-        if self.cfg.metric_scale_align_dl3dv:
-            # read invalid scales
-            scale_dir = '/cluster/project/cvg/haofei/datasets/depthsplat/dl3dv_metric_scale_factor'
-            filename = os.path.join(scale_dir, 'dl3dv_invalid.txt')
-            with open(filename, "r") as f:
-                self.invalid_scale_scenes = [line.strip() for line in f]
-
         # Calculate actual number of scenes (keys of index file that their chunks exist).
+        chunk_set = set(self.chunks)  # self.chunks can be huge (val multiplies it) -> set lookup
         num_available_scenes = 0
         for scene, chunk_path in self.index.items():
-            if chunk_path in self.chunks:
+            if chunk_path in chunk_set:
                 num_available_scenes += 1
         self.num_available_scenes = num_available_scenes
 
@@ -151,24 +114,6 @@ class DatasetDL3DV(IterableDataset):
         indices = torch.randperm(len(lst))
         return [lst[x] for x in indices]
 
-    def _get_scale_factor(self, scene: str) -> float:
-        """Get the scale factor for a scene."""
-        if self.cfg.metric_scale_align_dl3dv:
-            scale_dir = '/cluster/project/cvg/haofei/datasets/depthsplat/dl3dv_metric_scale_factor'
-            if self.stage == 'train':
-                folder = scene.split('_')[1]
-            else:
-                folder = scene
-            filename = os.path.join(scale_dir, folder, 'scale_factor.txt')
-
-            if not os.path.exists(filename) or folder in self.invalid_scale_scenes:
-                return self.cfg.scale_extrinsics
-            else:
-                with open(filename, "r") as f:
-                    return float(f.read().strip())
-        else:
-            return self.cfg.scale_extrinsics
-
     def _process_example_to_batch(
             self,
             example: dict,
@@ -176,7 +121,6 @@ class DatasetDL3DV(IterableDataset):
             intrinsics: Tensor,
             context_indices: Tensor,
             target_indices: Tensor,
-            scale_factor: float,
     ) -> Optional[dict]:
         """
         Process an example into a batch dict (original behavior).
@@ -184,19 +128,6 @@ class DatasetDL3DV(IterableDataset):
         """
         scene = example["key"]
 
-        # Load remaining context views if configured
-        if self.cfg.load_remain_context:
-            remaining_indices = get_remaining_indices(
-                context_indices, target_indices, self.cfg.num_remain_context
-            )
-            remain_context_images = [
-                example["images"][index.item()] for index in remaining_indices
-            ]
-            try:
-                remain_context_images = self.convert_images(remain_context_images)
-            except OSError:
-                return None
-
         # Load context images
         context_images = [
             example["images"][index.item()] for index in context_indices
@@ -216,16 +147,7 @@ class DatasetDL3DV(IterableDataset):
             return None
 
         # Validate image shapes
-        if self.cfg.mix_re10k and 'dl3dv' not in scene:
-            if self.cfg.highres:
-                expected_shape = (3, 720, 1280)
-            else:
-                expected_shape = (3, 360, 640)
-        else:
-            expected_shape = tuple([3, *self.cfg.ori_image_shape])
-
-        if self.stage in ['test', 'val'] or 'dl3dv' in scene:
-            expected_shape = tuple([3, *self.cfg.ori_image_shape])
+        expected_shape = tuple([3, *self.cfg.ori_image_shape])
 
         if self.cfg.skip_bad_shape:
             if context_images.shape[1:] != expected_shape or target_images.shape[1:] != expected_shape:
@@ -236,9 +158,6 @@ class DatasetDL3DV(IterableDataset):
                 )
                 return None
 
-            if self.cfg.load_remain_context and remain_context_images.shape[1:] != expected_shape:
-                return None
-
         # Apply pose transformations
         if self.cfg.pose_align_middle_view:
             mid_index = context_indices.shape[0] // 2
@@ -246,12 +165,6 @@ class DatasetDL3DV(IterableDataset):
                 extrinsics[context_indices][mid_index:mid_index + 1], extrinsics
             )
 
-        if self.cfg.pose_align_first_view:
-            extrinsics = camera_normalization(extrinsics[context_indices][0:1], extrinsics)
-
-        if self.cfg.center_pose:
-            extrinsics = center_norm_pose(extrinsics)
-
         # Validate extrinsics
         if any(torch.isnan(torch.det(extrinsics[context_indices][:, :3, :3]))):
             return None
@@ -272,20 +185,6 @@ class DatasetDL3DV(IterableDataset):
         ):
             return None
 
-        if self.cfg.load_remain_context:
-            if any(torch.isnan(torch.det(extrinsics[remaining_indices][:, :3, :3]))):
-                return None
-            if (extrinsics[remaining_indices][:, :3, 3] > 1e3).any():
-                return None
-            if not torch.allclose(
-                    torch.det(extrinsics[remaining_indices][:, :3, :3]),
-                    extrinsics[remaining_indices][:, :3, :3].new_tensor(1)
-            ):
-                return None
-
-        # Apply scale factor
-        extrinsics[:, :3, 3] *= scale_factor
-
         # Build output
         example_out = {
             "context": {
@@ -307,22 +206,12 @@ class DatasetDL3DV(IterableDataset):
             "scene": scene,
         }
 
-        if self.cfg.load_remain_context:
-            example_out["context_remain"] = {
-                "extrinsics": extrinsics[remaining_indices],
-                "intrinsics": intrinsics[remaining_indices],
-                "image": remain_context_images,
-                "near": self.get_bound("near", len(remaining_indices)),
-                "far": self.get_bound("far", len(remaining_indices)),
-                "index": remaining_indices,
-            }
-
         return example_out
 
     def __iter__(self):
         # Chunks must be shuffled here (not inside __init__) for validation to show
         # random chunks.
-        if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")):
+        if self.stage == "train" or (self.cfg.shuffle_val and self.stage == "val"):
             self.chunks = self.shuffle(self.chunks)
 
         # When testing, the data loaders alternate chunks.
@@ -350,7 +239,7 @@ class DatasetDL3DV(IterableDataset):
                 else:
                     chunk = item * len(chunk)
 
-            if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")):
+            if self.stage == "train" or (self.cfg.shuffle_val and self.stage == "val"):
                 chunk = self.shuffle(chunk)
 
             times_per_scene = (
@@ -373,8 +262,6 @@ class DatasetDL3DV(IterableDataset):
                 if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any():
                     continue
 
-                scale_factor = self._get_scale_factor(scene)
-
                 try:
                     extra_kwargs = {}
                     if self.cfg.overfit_to_scene is not None and self.stage != "test":
@@ -383,16 +270,12 @@ class DatasetDL3DV(IterableDataset):
                             else self.cfg.overfit_max_views
                         )
 
-                    is_re10k = self.cfg.mix_re10k and 'dl3dv' not in scene and self.stage == 'train'
-
                     out_data = self.view_sampler.sample(
                         scene,
                         extrinsics,
                         intrinsics,
                         min_context_views=self.cfg.min_views,
                         max_context_views=self.cfg.max_views,
-                        min_view_dist=self.cfg.re10k_min_view_dist if is_re10k else None,
-                        max_view_dist=self.cfg.re10k_max_view_dist if is_re10k else None,
                         **extra_kwargs,
                     )
 
@@ -421,7 +304,7 @@ class DatasetDL3DV(IterableDataset):
                 for context_indices, target_indices in zip(c_list, t_list):
                     example_out = self._process_example_to_batch(
                         example, extrinsics.clone(), intrinsics,
-                        context_indices, target_indices, scale_factor
+                        context_indices, target_indices,
                     )
 
                     if example_out is None:
@@ -440,13 +323,7 @@ class DatasetDL3DV(IterableDataset):
                     if self.cfg.image_shape == list(context_images.shape[2:]):
                         yield example_out
                     else:
-                        if self.stage == "train" and self.cfg.random_crop:
-                            crop_h = random.randint(self.cfg.min_size[0], self.cfg.max_size[0] + 1) // 64 * 64
-                            crop_w = random.randint(self.cfg.min_size[1], self.cfg.max_size[1] + 1) // 64 * 64
-                            crop_size = (crop_h, crop_w)
-                            yield apply_crop_shim(example_out, crop_size)
-                        else:
-                            yield apply_crop_shim(example_out, tuple(self.cfg.image_shape))
+                        yield apply_crop_shim(example_out, tuple(self.cfg.image_shape))
 
     def convert_poses(
             self,
@@ -472,22 +349,16 @@ class DatasetDL3DV(IterableDataset):
         w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4)
 
         if self.cfg.opencv_pose_format:
-            return self.opengl_to_opencv(w2c.inverse()), intrinsics
-        else:
-            return w2c.inverse(), intrinsics
-
-    def opengl_to_opencv(self, c2w):
-        # https://github.com/DL3DV-10K/Dataset/issues/4#issuecomment-2019441741
-        blender2opencv = np.array(
-            [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
-        )
-        blender2opencv = torch.tensor(blender2opencv, dtype=c2w.dtype, device=c2w.device).unsqueeze(0)
-        c2w = torch.matmul(c2w, blender2opencv)
-        c2w[:, 2, :] *= -1
-        c2w = c2w[:, torch.tensor(np.array([1, 0, 2, 3])), :]
-        c2w[:, 0:3, 1:3] *= -1
-
-        return c2w
+            # DL3DV chunk poses are already in OpenCV convention (verified: they match the
+            # COLMAP/OpenCV poses up to a world rotation), so opengl_to_opencv would corrupt
+            # them. Fail loudly instead of silently converting. NOTE: the resplat_v2 baseline
+            # in scripts/testing/_common/run_experiments.sh sets opencv_pose_format=true and
+            # will now hit this — update that path if resplat_v2 is still needed.
+            raise ValueError(
+                "opencv_pose_format=true is not supported for DL3DV: chunk poses are already "
+                "in OpenCV convention and need no conversion."
+            )
+        return w2c.inverse(), intrinsics
 
     def convert_images(
             self,
@@ -575,66 +446,3 @@ def camera_normalization(pivotal_pose: torch.Tensor, poses: torch.Tensor):
     normalized_poses = camera_norm_matrix_manuall @ poses  # [N, 4, 4]
 
     return normalized_poses
-
-
-def center_norm_pose(extrinsics):
-    # extrinsics: [V, 4, 4]
-    cam_centers = extrinsics[:, :3, 3]  # [V, 3]
-    avg_center = cam_centers.mean(dim=0, keepdim=True)  # [1, 3]
-    dist = (cam_centers - avg_center).norm(dim=1, keepdim=True)  # [V, 1]
-    scale = dist.max()
-
-    # translate
-    extrinsics = extrinsics.clone()
-    extrinsics[:, :3, 3] -= avg_center
-    extrinsics[:, :3, 3] /= scale
-
-    return extrinsics
-
-
-def get_remaining_indices(context_indices: torch.Tensor,
-                          target_indices: torch.Tensor,
-                          num_remain_context: int) -> torch.Tensor:
-    """
-    Randomly selects a fixed number of remaining indices in the range [min(context), max(context)],
-    excluding those in context or target. Pads by repeating if not enough remain.
-
-    Args:
-        context_indices (torch.Tensor): 1D tensor of context indices.
-        target_indices (torch.Tensor): 1D tensor of target indices.
-        num_remain_context (int): Number of remaining indices to return.
-
-    Returns:
-        torch.Tensor: 1D tensor of length `num_remain_context`.
-    """
-    if context_indices.numel() == 0:
-        raise ValueError("context_indices must not be empty.")
-
-    min_idx = torch.min(context_indices).item()
-    max_idx = torch.max(context_indices).item()
-
-    full_range = torch.arange(min_idx, max_idx + 1, dtype=torch.long)
-    exclude_indices = torch.cat([context_indices, target_indices])
-    mask = ~torch.isin(full_range, exclude_indices)
-
-    remaining = full_range[mask]
-
-    if remaining.numel() == 0:
-        # Nothing to sample from; repeat the first context index (or any fallback)
-        return context_indices[0].repeat(num_remain_context)
-
-    # return all
-    selected = remaining
-
-    # Randomly sample with or without replacement
-    # if remaining.numel() >= num_remain_context:
-    #     selected = remaining[torch.randperm(remaining.numel())[:num_remain_context]]
-    # else:
-    #     # return all
-    #     selected = remaining
-    #     # # Repeat with wrap-around to pad
-    #     # num_repeat = (num_remain_context + remaining.numel() - 1) // remaining.numel()
-    #     # padded = remaining.repeat(num_repeat)[:num_remain_context]
-    #     # selected = padded[torch.randperm(num_remain_context)]  # Shuffle for randomness
-
-    return selected.sort().values

optgs/dataset/dataset_re10k.py

@@ -8,13 +8,11 @@ from typing import Literal, Optional
 import numpy as np
 import torch
 import torchvision.transforms as tf
-import torch.nn.functional as F
 from einops import rearrange, repeat
 from jaxtyping import Float, UInt8
 from PIL import Image
 from torch import Tensor
 from torch.utils.data import IterableDataset
-import cv2
 
 from ..geometry.projection import get_fov
 from .dataset import DatasetCfgCommon
@@ -22,38 +20,25 @@ from .shims.augmentation_shim import apply_augmentation_shim
 from .shims.crop_shim import apply_crop_shim
 from .data_types import Stage
 from .view_sampler import ViewSampler
-from .dataset_dl3dv import get_remaining_indices
 
 
 @dataclass
 class DatasetRE10kCfg(DatasetCfgCommon):
     name: Literal["re10k"]
     roots: list[Path]
-    baseline_epsilon: float
     max_fov: float
-    make_baseline_1: bool
     augment: bool
     test_len: int
     test_chunk_interval: int
-    average_pose: bool 
     skip_bad_shape: bool = True
     near: float = -1.0
     far: float = -1.0
-    baseline_scale_bounds: bool = True
     shuffle_val: bool = True
     train_times_per_scene: int = 1
     highres: bool = False
     scannet: bool = False
     tartanair: bool = False
-    use_index_to_load_chunk: Optional[bool] = False
     load_depth: bool = False
-    pose_align_first_view: bool = False  # align the camera pose to the first view
-    center_pose: bool = False  # center and normalize the pose by the distance to the center
-
-    scale_extrinsics: float = 1.
-
-    # load remaining context views
-    load_remain_context: bool = False
 
 
 class DatasetRE10k(IterableDataset):
@@ -86,14 +71,9 @@ class DatasetRE10k(IterableDataset):
         self.chunks = []
         for i, root in enumerate(cfg.roots):
             root = root / self.data_stage
-            if self.cfg.use_index_to_load_chunk:
-                with open(root / "index.json", "r") as f:
-                    json_dict = json.load(f)
-                root_chunks = sorted(list(set(json_dict.values())))
-            else:
-                root_chunks = sorted(
-                    [path for path in root.iterdir() if path.suffix == ".torch"]
-                )
+            root_chunks = sorted(
+                [path for path in root.iterdir() if path.suffix == ".torch"]
+            )
 
             self.chunks.extend(root_chunks)
         if self.cfg.overfit_to_scene is not None:
@@ -110,7 +90,7 @@ class DatasetRE10k(IterableDataset):
     def __iter__(self):
         # Chunks must be shuffled here (not inside __init__) for validation to show
         # random chunks.
-        if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")):
+        if self.stage == "train" or (self.cfg.shuffle_val and self.stage == "val"):
             self.chunks = self.shuffle(self.chunks)
 
         # When testing, the data loaders alternate chunks.
@@ -134,7 +114,7 @@ class DatasetRE10k(IterableDataset):
                 assert len(item) == 1
                 chunk = item * len(chunk)
 
-            if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")):
+            if self.stage == "train" or (self.cfg.shuffle_val and self.stage == "val"):
                 chunk = self.shuffle(chunk)
 
             times_per_scene = (
@@ -163,23 +143,6 @@ class DatasetRE10k(IterableDataset):
                 if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any():
                     continue
 
-                # load remaining context views
-                if self.cfg.load_remain_context:
-                    # randomly select fixed number of remaining views such that they can be batched
-                    remaining_indices = get_remaining_indices(context_indices, target_indices, 
-                        0)
-
-                    # Load the images.
-                    remain_context_images = [
-                        example["images"][index.item()] for index in remaining_indices
-                    ]
-
-                    try:
-                        remain_context_images = self.convert_images(remain_context_images)
-                    except OSError:
-                        # some data might be corrupted
-                        continue
-
                 # Load the images.
                 context_images = [
                     example["images"][index.item()] for index in context_indices
@@ -207,12 +170,6 @@ class DatasetRE10k(IterableDataset):
                     )
                     continue
 
-                if self.cfg.load_remain_context:
-                    remain_context_invalid = remain_context_images.shape[1:] != expected_shape
-
-                    if self.cfg.skip_bad_shape and remain_context_invalid:
-                        continue
-
                 # check the extrinsics
                 if any(torch.isnan(torch.det(extrinsics[context_indices][:, :3, :3]))):
                     continue
@@ -220,9 +177,6 @@ class DatasetRE10k(IterableDataset):
                 if any(torch.isnan(torch.det(extrinsics[target_indices][:, :3, :3]))):
                     continue
 
-                if self.cfg.average_pose:
-                    extrinsics = self.preprocess_poses(extrinsics)
-
                 # load depth
                 if self.cfg.load_depth:
                     context_depths = [
@@ -245,15 +199,12 @@ class DatasetRE10k(IterableDataset):
                     else:
                         raise NotImplementedError
 
-                # align pose to the first view
-                if self.cfg.pose_align_first_view:
-                    extrinsics = camera_normalization(extrinsics[context_indices][0:1], extrinsics)
-
-                if self.cfg.center_pose:
-                    extrinsics = center_norm_pose(extrinsics)
-
-                # scale the scene when necessary: only scale the extrinsics
-                extrinsics[:, :3, 3] *= self.cfg.scale_extrinsics
+                # align pose to the middle view
+                if self.cfg.pose_align_middle_view:
+                    mid_index = context_indices.shape[0] // 2
+                    extrinsics = camera_normalization(
+                        extrinsics[context_indices][mid_index:mid_index + 1], extrinsics
+                    )
 
                 example = {
                     "context": {
@@ -275,19 +226,6 @@ class DatasetRE10k(IterableDataset):
                     "scene": scene,
                 }
 
-                if self.cfg.load_remain_context:
-                    example.update({
-                        "context_remain": {
-                            "extrinsics": extrinsics[remaining_indices],
-                            "intrinsics": intrinsics[remaining_indices],
-                            "image": remain_context_images,
-                            "near": self.get_bound("near", len(remaining_indices)),
-                            "far": self.get_bound("far", len(remaining_indices)),
-                            "index": remaining_indices,
-                        }
-                        }
-                    )
-
                 if self.cfg.load_depth:
                     example['context']['depth'] = context_depths
                     example['target']['depth'] = target_depths
@@ -400,65 +338,20 @@ class DatasetRE10k(IterableDataset):
             else len(self.index.keys()) * self.cfg.train_times_per_scene
         )
 
-    def preprocess_poses(
-        self,
-        in_c2ws: torch.Tensor,
-        scene_scale_factor=1.35,
-    ):
-        """
-        Ref: https://github.com/Haian-Jin/LVSM/blob/main/data/dataset_scene.py
-        Preprocess the poses to:
-        1. translate and rotate the scene to align the average camera direction and position
-        2. rescale the whole scene to a fixed scale
-        """
-
-        # Translation and Rotation
-        # align coordinate system (OpenCV coordinate) to the mean camera
-        # center is the average of all camera centers
-        # average direction vectors are computed from all camera direction vectors (average down and forward)
-        center = in_c2ws[:, :3, 3].mean(0)
-        avg_forward = F.normalize(in_c2ws[:, :3, 2].mean(0), dim=-1) # average forward direction (z of opencv camera)
-        avg_down = in_c2ws[:, :3, 1].mean(0) # average down direction (y of opencv camera)
-        avg_right = F.normalize(torch.cross(avg_down, avg_forward, dim=-1), dim=-1) # (x of opencv camera)
-        avg_down = F.normalize(torch.cross(avg_forward, avg_right, dim=-1), dim=-1) # (y of opencv camera)
-
-        avg_pose = torch.eye(4, device=in_c2ws.device) # average c2w matrix
-        avg_pose[:3, :3] = torch.stack([avg_right, avg_down, avg_forward], dim=-1)
-        avg_pose[:3, 3] = center 
-        avg_pose = torch.linalg.inv(avg_pose) # average w2c matrix
-        in_c2ws = avg_pose @ in_c2ws 
-
-
-        # Rescale the whole scene to a fixed scale
-        scene_scale = torch.max(torch.abs(in_c2ws[:, :3, 3]))
-        scene_scale = scene_scale_factor * scene_scale
-
-        in_c2ws[:, :3, 3] /= scene_scale
-
-        return in_c2ws
-
 
 def camera_normalization(pivotal_pose: torch.Tensor, poses: torch.Tensor):
     # [1, 4, 4], [N, 4, 4]
 
-    camera_norm_matrix = torch.inverse(pivotal_pose)
-    
-    # normalize all views
-    poses = torch.bmm(camera_norm_matrix.repeat(poses.shape[0], 1, 1), poses)
-
-    return poses
-
+    # Manually calculate the inverse of SE(3) to avoid numerical issues
+    R = pivotal_pose[:, :3, :3]  # [1, 3, 3]
+    t = pivotal_pose[:, :3, 3:]  # [1, 3, 1]
+    R_inv = R.transpose(-1, -2)  # [1, 3, 3]
+    t_inv = -R_inv @ t  # [1, 3, 1]
+    camera_norm_matrix_manuall = torch.eye(4, dtype=poses.dtype, device=poses.device).unsqueeze(0)  # [1, 4, 4]
+    camera_norm_matrix_manuall[:, :3, :3] = R_inv
+    camera_norm_matrix_manuall[:, :3, 3:] = t_inv
 
-def center_norm_pose(extrinsics):
-    # extrinsics: [V, 4, 4]
-    cam_centers = extrinsics[:, :3, 3]  # [V, 3]
-    avg_center = cam_centers.mean(dim=0, keepdim=True)  # [1, 3]
-    dist = (cam_centers - avg_center).norm(dim=1, keepdim=True)  # [V, 1]
-    scale = dist.max()
-
-    # translate
-    extrinsics = extrinsics.clone()
-    extrinsics[:, :3, 3] -= avg_center
-    extrinsics[:, :3, 3] /= scale
+    # normalize all views
+    normalized_poses = camera_norm_matrix_manuall @ poses  # [N, 4, 4]
 
-    return extrinsics
+    return normalized_poses

optgs/dataset/shims/augmentation_shim.py

@@ -22,25 +22,26 @@ def reflect_views(views: AnyViews) -> AnyViews:
     }
 
 
+def mark_unflipped(views: AnyViews) -> AnyViews:
+    """No-op augmentation that still records x_flipped=False, so the key is
+    always present and downstream code (e.g. BatchedViews.from_dict) need not
+    fall back to a default."""
+    return {**views, "x_flipped": False}
+
+
 def apply_augmentation_shim(
     example: AnyExample,
     generator: torch.Generator | None = None,
 ) -> AnyExample:
-    """Randomly augment the training images."""
-    # Do not augment with 50% chance.
-    if torch.rand(tuple(), generator=generator) < 0.5:
-        return example
-
-    if "context_remain" in example:
-        return {
-            **example,
-            "context": reflect_views(example["context"]),
-            "target": reflect_views(example["target"]),
-            "context_remain": reflect_views(example["context_remain"]),
-        }
-
-    return {
+    """Randomly horizontally-flip the scene (50% chance). Either way, x_flipped
+    is recorded on every view set so the key is always present."""
+    flip = torch.rand(tuple(), generator=generator) >= 0.5
+    transform = reflect_views if flip else mark_unflipped
+    out = {
         **example,
-        "context": reflect_views(example["context"]),
-        "target": reflect_views(example["target"]),
+        "context": transform(example["context"]),
+        "target": transform(example["target"]),
     }
+    if "context_remain" in example:
+        out["context_remain"] = transform(example["context_remain"])
+    return out

optgs/dataset/shims/bounds_shim.py

@@ -1,80 +0,0 @@
-import torch
-from einops import einsum, reduce, repeat
-from jaxtyping import Float
-from torch import Tensor
-
-from ..data_types import BatchedExample
-
-
-def compute_depth_for_disparity(
-    extrinsics: Float[Tensor, "batch view 4 4"],
-    intrinsics: Float[Tensor, "batch view 3 3"],
-    image_shape: tuple[int, int],
-    disparity: float,
-    delta_min: float = 1e-6,  # This prevents motionless scenes from lacking depth.
-) -> Float[Tensor, " batch"]:
-    """Compute the depth at which moving the maximum distance between cameras
-    corresponds to the specified disparity (in pixels).
-    """
-
-    # Use the furthest distance between cameras as the baseline.
-    origins = extrinsics[:, :, :3, 3]
-    deltas = (origins[:, None, :, :] - origins[:, :, None, :]).norm(dim=-1)
-    deltas = deltas.clip(min=delta_min)
-    baselines = reduce(deltas, "b v ov -> b", "max")
-
-    # Compute a single pixel's size at depth 1.
-    h, w = image_shape
-    pixel_size = 1 / torch.tensor((w, h), dtype=torch.float32, device=extrinsics.device)
-    pixel_size = einsum(
-        intrinsics[..., :2, :2].inverse(), pixel_size, "... i j, j -> ... i"
-    )
-
-    # This wouldn't make sense with non-square pixels, but then again, non-square pixels
-    # don't make much sense anyway.
-    mean_pixel_size = reduce(pixel_size, "b v xy -> b", "mean")
-
-    return baselines / (disparity * mean_pixel_size)
-
-
-def apply_bounds_shim(
-    batch: BatchedExample,
-    near_disparity: float,
-    far_disparity: float,
-) -> BatchedExample:
-    """Compute reasonable near and far planes (lower and upper bounds on depth). This
-    assumes that all of an example's views are of roughly the same thing.
-    """
-
-    context = batch["context"]
-    _, cv, _, h, w = context["image"].shape
-
-    # Compute near and far planes using the context views.
-    near = compute_depth_for_disparity(
-        context["extrinsics"],
-        context["intrinsics"],
-        (h, w),
-        near_disparity,
-    )
-    far = compute_depth_for_disparity(
-        context["extrinsics"],
-        context["intrinsics"],
-        (h, w),
-        far_disparity,
-    )
-
-    target = batch["target"]
-    _, tv, _, _, _ = target["image"].shape
-    return {
-        **batch,
-        "context": {
-            **context,
-            "near": repeat(near, "b -> b v", v=cv),
-            "far": repeat(far, "b -> b v", v=cv),
-        },
-        "target": {
-            **target,
-            "near": repeat(near, "b -> b v", v=tv),
-            "far": repeat(far, "b -> b v", v=tv),
-        },
-    }

optgs/dataset/view_sampler/view_sampler.py

@@ -44,10 +44,6 @@ class ViewSampler(ABC, Generic[T]):
         self._all_context_indices = None
         self._all_target_indices = None
 
-    @property
-    def all_context_indices(self) -> Int64[Tensor, " context_view"]:
-        return self._all_context_indices
-
     @property
     def context_indices(self) -> Int64[Tensor, " target_view"]:
         return self._all_context_indices
@@ -70,9 +66,6 @@ class ViewSampler(ABC, Generic[T]):
         else:
             raise RuntimeError("Target indices have already been set.")
 
-    def sample_subset(self, extrinsics, intrinsics, device):
-        pass
-
     @abstractmethod
     def _sample_impl(
             self,
@@ -123,13 +116,3 @@ class ViewSampler(ABC, Generic[T]):
     @property
     def global_step(self) -> int:
         return 0 if self.step_tracker is None else self.step_tracker.get_step()
-
-    def new_instance(self) -> "ViewSampler":
-        """Create a new instance of the same ViewSampler class with the same configuration."""
-        return value(self.__class__)(
-            cfg=self.cfg,
-            stage=self.stage,
-            is_overfitting=self.is_overfitting,
-            cameras_are_circular=self.cameras_are_circular,
-            step_tracker=self.step_tracker,
-        )

optgs/evaluation/metric_computer.py

@@ -1,115 +0,0 @@
-import os
-from pathlib import Path
-
-import torch
-from pytorch_lightning import LightningModule
-from ..misc.console import metrics_table
-from ..misc.image_io import load_image, save_image
-from ..visualization.annotation import add_label
-from ..visualization.layout import add_border, hcat
-from .evaluation_cfg import EvaluationCfg
-from .metrics import compute_lpips, compute_psnr, compute_ssim
-
-
-class MetricComputer(LightningModule):
-    cfg: EvaluationCfg
-
-    def __init__(self, cfg: EvaluationCfg) -> None:
-        super().__init__()
-        self.cfg = cfg
-
-    def test_step(self, batch, batch_idx):
-        scene = batch["scene"][0]
-        b, cv, _, _, _ = batch["context"]["image"].shape
-        assert b == 1 and cv == 2
-        _, v, _, _, _ = batch["target"]["image"].shape
-
-        # Skip scenes.
-        for method in self.cfg.methods:
-            if not (method.path / scene).exists():
-                print(f'Skipping "{scene}".')
-                return
-
-        # Load the images.
-        all_images = {}
-        try:
-            for method in self.cfg.methods:
-                images = [
-                    load_image(method.path / scene / f"color/{index.item():0>6}.png")
-                    for index in batch["target"]["index"][0]
-                ]
-                all_images[method.key] = torch.stack(images).to(self.device)
-        except FileNotFoundError:
-            print(f'Skipping "{scene}".')
-            return
-
-        # Compute metrics.
-        all_metrics = {}
-        rgb_gt = batch["target"]["image"][0]
-        for key, images in all_images.items():
-            alex_lpips, vgg_lpips = compute_lpips(rgb_gt, images)
-            all_metrics = {
-                **all_metrics,
-                f"alex_lpips_{key}": alex_lpips,
-                f"vgg_lpips_{key}": vgg_lpips,
-                f"ssim_{key}": compute_ssim(rgb_gt, images),
-                f"psnr_{key}": compute_psnr(rgb_gt, images),
-            }
-        self.log_dict(all_metrics)
-        self.print_preview_metrics(all_metrics)
-
-        # Skip the rest if no side-by-side is needed.
-        if self.cfg.side_by_side_path is None:
-            return
-
-        # Create side-by-side.
-        scene_key = f"{batch_idx:0>6}_{scene}"
-        for i in range(v):
-            true_index = batch["target"]["index"][0, i]
-            row = [add_label(batch["target"]["image"][0, i], "Ground Truth")]
-            for method in self.cfg.methods:
-                image = all_images[method.key][i]
-                image = add_label(image, method.name)
-                row.append(image)
-            start_frame = batch["target"]["index"][0, 0]
-            end_frame = batch["target"]["index"][0, -1]
-            label = f"Scene {batch['scene'][0]} (frames {start_frame} to {end_frame})"
-            row = add_border(add_label(hcat(*row), label, font_size=16))
-            save_image(
-                row,
-                self.cfg.side_by_side_path / scene_key / f"{true_index:0>6}.png",
-            )
-
-        # Create an animation.
-        if self.cfg.animate_side_by_side:
-            (self.cfg.side_by_side_path / "videos").mkdir(exist_ok=True, parents=True)
-            command = (
-                'ffmpeg -y -framerate 30 -pattern_type glob -i "*.png"  -c:v libx264 '
-                '-pix_fmt yuv420p -vf "pad=ceil(iw/2)*2:ceil(ih/2)*2"'
-            )
-            os.system(
-                f"cd {self.cfg.side_by_side_path / scene_key} && {command} "
-                f"{Path.cwd()}/{self.cfg.side_by_side_path}/videos/{scene_key}.mp4"
-            )
-
-    def print_preview_metrics(self, metrics: dict[str, float]) -> None:
-        if getattr(self, "running_metrics", None) is None:
-            self.running_metrics = metrics
-            self.running_metric_steps = 1
-        else:
-            s = self.running_metric_steps
-            self.running_metrics = {
-                k: ((s * v) + metrics[k]) / (s + 1)
-                for k, v in self.running_metrics.items()
-            }
-            self.running_metric_steps += 1
-
-        rows = []
-        for method in self.cfg.methods:
-            row = [
-                f"{self.running_metrics[f'{metric}_{method.key}']:.3f}"
-                for metric in ("psnr", "lpips", "ssim")
-            ]
-            rows.append((method.key, *row))
-
-        metrics_table(rows, ["Method", "PSNR (dB)", "LPIPS", "SSIM"])

optgs/evaluation/metrics.py

@@ -6,7 +6,7 @@ from jaxtyping import Float
 # from lpips import LPIPS
 # from skimage.metrics import structural_similarity
 from torch import Tensor
-from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics.image import StructuralSimilarityIndexMeasure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 from tqdm import tqdm
 

optgs/experimental/api/api.py

@@ -21,7 +21,7 @@ For non-inria codebases use :meth:`OptGS.initialize_from_ply` /
 
 External SfM scenes carry no optgs encoder features, so checkpoints trained
 with ``init_state_wo_features=False`` are coerced at construction (with a
-warning): the feature-conditioned ``update_proj`` weights are dropped and the
+warning): the feature-conditioned ``state_proj`` weights are dropped and the
 optimizer state is initialized standard-normal.
 """
 
@@ -105,7 +105,7 @@ class OptGS:
                 "(scene_trainer.scene_optimizer.init_state_wo_features=False). "
                 "External SfM/inria scenes carry no optgs encoder features; "
                 "proceeding with init_state_wo_features=True — the "
-                "feature-conditioned update_proj weights are dropped and the "
+                "feature-conditioned state_proj weights are dropped and the "
                 "initial optimizer state is set to a standard-normal random "
                 "vector (init_state_type='random', init_state_scale=1.0)."
             )
@@ -348,7 +348,6 @@ class OptGS:
             context=self._context,
             renderer=self.decoder,
             prev_output=self._init_output,
-            num_refine=self.num_refine,
             iter_batch_size=self.iter_batch_size,
             target=self._context,
         )
@@ -416,27 +415,36 @@ class OptGS:
             OptimizerPreviousOutput,
         )
 
+        # torch's grad mode is thread-local, and a generator consumer (e.g.
+        # gradio) may resume this generator on a *different* worker thread each
+        # step — so a single ``with torch.no_grad()`` spanning the ``yield`` does
+        # not reliably hold across steps. Enter no_grad *per step* and yield
+        # outside it: the Adam baseline's in-place parameter updates require
+        # no_grad ambient (calc_input_gradients re-enables grad internally just
+        # for the gradient computation), so a step that resumes on a grad-enabled
+        # thread would otherwise raise "a leaf Variable that requires grad is
+        # being used in an in-place operation" and leak the autograd graph.
+        inp = OptimizerInput(
+            context=self._context,
+            renderer=self.decoder,
+            prev_output=self._init_output,
+            iter_batch_size=self.iter_batch_size,
+            target=self._context,
+        )
         with torch.no_grad():
-            inp = OptimizerInput(
-                context=self._context,
-                renderer=self.decoder,
-                prev_output=self._init_output,
-                num_refine=self.num_refine,
-                iter_batch_size=self.iter_batch_size,
-                target=self._context,
-            )
             opt.validate_input(inp)
             opt.on_scene_start(inp)  # InitializerOutput -> OptimizerPreviousOutput
-            if not isinstance(inp.prev_output, OptimizerPreviousOutput):
-                raise OptGSError(
-                    "optimizer.on_scene_start did not produce an "
-                    f"OptimizerPreviousOutput (got {type(inp.prev_output)})."
-                )
-
-            out = OptimizerOutput.empty(t=0)
-            out.T = self.num_refine
-            try:
-                for step in range(self.num_refine):
+        if not isinstance(inp.prev_output, OptimizerPreviousOutput):
+            raise OptGSError(
+                "optimizer.on_scene_start did not produce an "
+                f"OptimizerPreviousOutput (got {type(inp.prev_output)})."
+            )
+
+        out = OptimizerOutput.empty(t=0)
+        out.T = self.num_refine
+        try:
+            for step in range(self.num_refine):
+                with torch.no_grad():
                     # Fresh view minibatch each step (the regime the optimizer
                     # was trained with); full_context/target stay the whole scene.
                     batch = self._view_minibatch(self._context)
@@ -447,12 +455,14 @@ class OptGS:
                         full_context=self._context, full_target=self._context,
                     )
                     out.t = (out.t or 0) + 1
-                    yield step, inp.prev_output.gaussians
-            finally:
+                    g = inp.prev_output.gaussians
+                yield step, g
+        finally:
+            with torch.no_grad():
                 if torch.cuda.is_available():
                     torch.cuda.synchronize()
                 opt.on_scene_end()
-                self._refined = inp.prev_output.gaussians
+            self._refined = inp.prev_output.gaussians
 
     def export_ply(self, path: str) -> None:
         """Write the most recently refined Gaussians to a 3DGS PLY."""

optgs/experimental/api/integration/config_bridge.py

@@ -190,22 +190,21 @@ def build_optimizer_cfg(cfg_path: Path) -> tuple["KnnBasedOptimizerCfg", int | N
 
     # Mirror SceneTrainerCfg (scene_trainer_cfg.py: scene_optimizer.update(
     # scene_initializer)): wire the checkpoint's initializer cfg into the
-    # optimizer cfg so the runtime-only fields init_gaussian_param_num /
-    # init_sh_d / sh_d — absent from every config file — are populated before
-    # the optimizer nn.Module is built.
+    # optimizer cfg so the runtime-only fields init_sh_d / sh_d — absent from
+    # every config file — are populated before the optimizer nn.Module is built.
     si = OmegaConf.select(cfg, "scene_trainer.scene_initializer")
     si_name = OmegaConf.select(cfg, "scene_trainer.scene_initializer.name")
     if si is None or si_name in (None, "none"):
         raise OptGSError(
             f"checkpoint config at {cfg_path} has no scene_initializer "
-            f"(name={si_name!r}); cannot derive init_gaussian_param_num "
-            f"required to build the optimizer."
+            f"(name={si_name!r}); cannot derive the optimizer's initializer "
+            f"settings required to build it."
         )
     init_cls = _initializer_cfg_class(str(si_name))
     if init_cls is None:
         raise OptGSError(
             f"unsupported scene_initializer.name={si_name!r} in {cfg_path}; "
-            f"cannot derive init_gaussian_param_num for the optimizer."
+            f"cannot derive the optimizer's initializer settings."
         )
     default_si = _compose_default_group("scene_initializer", str(si_name))
     if default_si is not None:
@@ -215,7 +214,7 @@ def build_optimizer_cfg(cfg_path: Path) -> tuple["KnnBasedOptimizerCfg", int | N
         merged_si = si
     try:
         init_cfg = load_typed_config(merged_si, init_cls)
-        opt_cfg.update(init_cfg)  # sets init_gaussian_param_num/init_sh_d/sh_d
+        opt_cfg.update(init_cfg)  # sets init_sh_d/sh_d
     except Exception as e:
         raise OptGSError(
             f"failed to wire scene_initializer ({si_name!r}) into the "
@@ -335,7 +334,7 @@ def build_adam_baseline(num_refine: int) -> "nn.Module":
         )
     OmegaConf.set_struct(composed, False)
     # gsplat decays the means LR over the full step budget.
-    composed.means_lr_max_steps = int(num_refine)
+    composed.lr_scheduler.max_steps = int(num_refine)
     # Disable densification — the baseline refines the same fixed Gaussian set
     # as the learned optimizer (a like-for-like comparison of the update rule).
     for flag in ("do_densify", "do_prune", "do_opacity_reset"):
@@ -357,14 +356,6 @@ def build_adam_baseline(num_refine: int) -> "nn.Module":
     return optimizer
 
 
-# Module-attribute renames applied when the legacy Resplat encoder was split
-# into separate initializer/optimizer modules (transcribed from
-# optgs/main.py:load_optimizer).
-_ORIG_OPTIMIZER_ATTR_RENAMES = {
-    "render_error_mv_attn": "update_error_attn",
-}
-
-
 def load_optimizer_state(
     optimizer: "nn.Module",
     ckpt_path: str,
@@ -379,6 +370,8 @@ def load_optimizer_state(
     """
     import torch
 
+    from optgs.misc.checkpointing import _rename_optimizer_attrs
+
     state = torch.load(ckpt_path, map_location="cpu")
     if isinstance(state, dict) and "state_dict" in state:
         state = state["state_dict"]
@@ -399,14 +392,10 @@ def load_optimizer_state(
             for k, v in state.items()
             if k.startswith("encoder.")
         }
-        renamed = {}
-        for k, v in osd.items():
-            for old, new in _ORIG_OPTIMIZER_ATTR_RENAMES.items():
-                if k == old or k.startswith(old + "."):
-                    k = new + k[len(old):]
-                    break
-            renamed[k] = v
-        osd = renamed
+
+    # Rename module attributes renamed in the update_/refine_ cleanup pass
+    # (and the resplat-era render_error_mv_attn).
+    osd = _rename_optimizer_attrs(osd)
 
     if not osd:
         raise OptGSError(
@@ -415,6 +404,6 @@ def load_optimizer_state(
         )
 
     if init_state_wo_features:
-        osd = {k: v for k, v in osd.items() if "update_proj" not in k}
+        osd = {k: v for k, v in osd.items() if "state_proj" not in k}
 
     optimizer.load_state_dict(osd, strict=strict)

optgs/global_cfg.py

@@ -1,19 +0,0 @@
-from typing import Optional
-
-from omegaconf import DictConfig
-
-cfg: Optional[DictConfig] = None
-
-
-def get_cfg() -> DictConfig:
-    global cfg
-    return cfg
-
-
-def set_cfg(new_cfg: DictConfig) -> None:
-    global cfg
-    cfg = new_cfg
-
-
-def get_seed() -> int:
-    return cfg.seed

optgs/loss/loss_deltas.py

@@ -24,6 +24,13 @@ class LossDeltasCfgWrapper:
 
 
 class LossDeltas(Loss[LossDeltasCfg, LossDeltasCfgWrapper]):
+    """L1 regularization on the optimizer's predicted per-Gaussian deltas, keeping the updates small.
+
+    With cfg.exclude_by_norm_grad the penalty is restricted to Gaussians with small gradients (and,
+    with exclude_by_norm_grad_opposite, those whose delta already agrees in sign with the gradient).
+    Returns 0 until global_step >= cfg.apply_after_step.
+    """
+
     def forward(
         self,
         prediction: DecoderOutput,
@@ -32,8 +39,6 @@ class LossDeltas(Loss[LossDeltasCfg, LossDeltasCfgWrapper]):
         gt_rgb: Tensor,
         pred_rgb: Tensor,
         valid_depth_mask: Tensor | None,
-        l1_loss: bool,
-        clamp_large_error: float,
         **kwargs,
     ) -> Float[Tensor, ""]:
 

optgs/loss/loss_lpips.py

@@ -20,6 +20,7 @@ class LossLpipsCfg:
     weight: float
     apply_after_step: int
     perceptual_loss: bool
+    half_res: bool  # downsample to half resolution before LPIPS to save memory
 
 
 @dataclass
@@ -28,6 +29,12 @@ class LossLpipsCfgWrapper:
 
 
 class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]):
+    """LPIPS perceptual distance between the rendered and ground-truth views.
+
+    Returns 0 until global_step >= cfg.apply_after_step. Uses VGG-LPIPS, or the custom
+    PerceptualLoss when cfg.perceptual_loss.
+    """
+
     lpips: LPIPS
 
     def __init__(self, cfg: LossLpipsCfgWrapper) -> None:
@@ -48,7 +55,6 @@ class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]):
         gt_rgb: Tensor,
         pred_rgb: Tensor,
         valid_depth_mask: Tensor | None,
-        half_res_lpips: bool = False,
         **kwargs,
     ) -> Float[Tensor, ""]:
 
@@ -64,7 +70,7 @@ class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]):
         pred = rearrange(pred_rgb, "b v c h w -> (b v) c h w")
         gt = rearrange(gt_rgb, "b v c h w -> (b v) c h w")
 
-        if half_res_lpips:
+        if self.cfg.half_res:
             pred = F.interpolate(pred, scale_factor=0.5, mode="bilinear", align_corners=True)
             gt = F.interpolate(gt, scale_factor=0.5, mode="bilinear", align_corners=True)
 

optgs/loss/loss_monodepth.py

@@ -1,11 +1,6 @@
 import torch
 
-
-import cv2
-import torch
-
-from optgs.model.encoder.depth_anything_v2.dpt import DepthAnythingV2
-
+from optgs.model.backbones.depth_anything_v2.dpt import DepthAnythingV2
 
 
 def get_monodepth_model():
@@ -16,7 +11,7 @@ def get_monodepth_model():
         'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
     }
 
-    encoder = 'vitl' # or 'vits', 'vitb', 'vitg'
+    encoder = 'vitl'  # or 'vits', 'vitb', 'vitg'
 
     model = DepthAnythingV2(**model_configs[encoder])
     model.load_state_dict(torch.load(f'pretrained/depth_anything_v2_{encoder}.pth', map_location='cpu'))
@@ -26,17 +21,3 @@ def get_monodepth_model():
         param.requires_grad = False
 
     return model
-
-
-
-def get_monodepth_pred(img, model):
-
-    with torch.no_grad():
-        pass
-
-
-def get_monodepth_loss(pred_depth, img):
-    pass
-
-
-

optgs/loss/loss_mse.py

@@ -12,6 +12,8 @@ from .loss import Loss
 @dataclass
 class LossMseCfg:
     weight: float
+    l1_loss: bool  # use L1 instead of L2 on the RGB error
+    clamp_large_error: float  # if > 0, drop per-pixel errors above this threshold
 
 
 @dataclass
@@ -20,6 +22,8 @@ class LossMseCfgWrapper:
 
 
 class LossMse(Loss[LossMseCfg, LossMseCfgWrapper]):
+    """RGB reconstruction loss between the rendered and ground-truth views (L2, or L1 if cfg.l1_loss)."""
+
     def forward(
             self,
             prediction: DecoderOutput,
@@ -28,8 +32,6 @@ class LossMse(Loss[LossMseCfg, LossMseCfgWrapper]):
             gt_rgb: Tensor,
             pred_rgb: Tensor,
             valid_depth_mask: Tensor | None,
-            l1_loss: bool,
-            clamp_large_error: float,
             **kwargs,
     ) -> Float[Tensor, ""]:
 
@@ -38,15 +40,15 @@ class LossMse(Loss[LossMseCfg, LossMseCfgWrapper]):
         if valid_depth_mask is not None and valid_depth_mask.max() > 0.5 and valid_depth_mask.min() < 0.5:
             error = error[~valid_depth_mask]
 
-        if l1_loss:
+        if self.cfg.l1_loss:
             # l1 loss
             error = error.abs()
         else:
             # l2 loss
             error = error ** 2
 
-        if clamp_large_error > 0:
-            valid_mask = error < clamp_large_error
+        if self.cfg.clamp_large_error > 0:
+            valid_mask = error < self.cfg.clamp_large_error
             error = error[valid_mask]
 
         error = error.mean()

optgs/loss/loss_sgd.py

@@ -3,7 +3,6 @@ from dataclasses import dataclass
 from jaxtyping import Float
 from torch import Tensor
 
-from optgs.dataset.data_types import BatchedExample
 from optgs.loss import Loss
 from optgs.model.decoder.decoder import DecoderOutput
 from optgs.model.types import Gaussians
@@ -12,7 +11,8 @@ from optgs.scene_trainer.gaussian_module import GaussiansModule
 
 @dataclass
 class LossSGDCfg:
-    pass
+    l1_loss: bool  # use L1 instead of L2 on the delta-vs-gradient error
+    clamp_large_error: float  # if > 0, clamp the loss to this maximum
 
 @dataclass
 class LossSGDCfgWrapper:
@@ -20,20 +20,21 @@ class LossSGDCfgWrapper:
 
 
 class LossSGD(Loss[LossSGDCfg, LossSGDCfgWrapper]):
+    """Debugging / sanity-check loss: supervises the optimizer's predicted per-Gaussian deltas to match
+    the true gradients, i.e. checks the learned optimizer can reproduce a plain gradient-descent step
+    (L2, or L1 if cfg.l1_loss)."""
+
     def forward(
         self,
         prediction: DecoderOutput,
-        batch: BatchedExample,
         gaussians: Gaussians | GaussiansModule | None,
         global_step: int,
-        l1_loss: bool,
-        clamp_large_error: float,
         valid_depth_mask: Tensor | None,
         **kwargs,
     ) -> Float[Tensor, ""]:
-        
+
         if gaussians is None:
-            raise ValueError("Gaussians must be provided for LossDeltas.")        
+            raise ValueError("Gaussians must be provided for LossSGD.")
 
         predicted_deltas = gaussians.deltas
         gt_gradients = gaussians.gradients
@@ -42,10 +43,10 @@ class LossSGD(Loss[LossSGDCfg, LossSGDCfgWrapper]):
         if predicted_deltas.dtype != gt_gradients.dtype:
             gt_gradients = gt_gradients.to(predicted_deltas.dtype)
 
-        if l1_loss:
+        if self.cfg.l1_loss:
             loss = (predicted_deltas - gt_gradients).abs().mean()
         else:
             loss = ((predicted_deltas - gt_gradients) ** 2).mean()
-        if clamp_large_error > 0:
-            loss = loss.clamp(max=clamp_large_error)
+        if self.cfg.clamp_large_error > 0:
+            loss = loss.clamp(max=self.cfg.clamp_large_error)
         return loss

optgs/loss/loss_sh0.py

@@ -18,6 +18,9 @@ class LossSh0CfgWrapper:
     sh0: LossSh0Cfg
 
 class LossSh0(Loss[LossSh0Cfg, LossSh0CfgWrapper]):
+    """Supervises each Gaussian's SH degree-0 (base) color against the ground-truth image colour
+    sampled at the Gaussian's projected 2D location, averaged over the views where it is visible."""
+
     def forward(
         self,
         prediction: DecoderOutput,