From faad7bf69b459b4dd28e10ee77d95e7ec8bcfa5e Mon Sep 17 00:00:00 2001
From: Sricharan Reddy Varra <sricharan.varra@biohub.org>
Date: Fri, 6 Nov 2026 15:04:07 -0800
Subject: [PATCH] fix(cli): run apply-inv-tf on the config's device

apply_inverse_transfer_function_single_position never passed settings.device to the phase and fluorescence models, so apply-inv-tf and reconstruct ran them on the CPU whatever the config said, Richardson-Lucy included. Pass it through; the birefringence models take no device and stay on the CPU.
---
 tests/cli_tests/test_reconstruct.py           | 50 +++++++++++++++++++
 .../cli/apply_inverse_transfer_function.py    |  6 ++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/tests/cli_tests/test_reconstruct.py b/tests/cli_tests/test_reconstruct.py
index 040441d6..66baac03 100644
--- a/tests/cli_tests/test_reconstruct.py
+++ b/tests/cli_tests/test_reconstruct.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import pytest
+import torch
 from iohub.ngff import open_ome_zarr
 from iohub.ngff.models import TransformationMeta
 from typer.testing import CliRunner
@@ -703,6 +704,55 @@ def test_apply_inv_tf_process_pool_matches_serial(tmp_path):
     np.testing.assert_array_equal(pooled, serial)
 
 
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device required")
+@pytest.mark.parametrize(
+    "contrast",
+    [
+        {"phase": settings.PhaseSettings()},
+        {
+            "fluorescence": settings.FluorescenceSettings(
+                apply_inverse={"reconstruction_algorithm": "RL", "rl": {"iterations": 5}}
+            )
+        },
+    ],
+    ids=["phase", "fluorescence"],
+)
+def test_apply_inv_tf_runs_on_the_config_device(tmp_path, contrast):
+    """`device: cuda` reconstructs on the GPU, to the CPU's result."""
+    input_path = tmp_path / "input.zarr"
+    data = np.random.default_rng(0).uniform(1, 100, size=(2, 1, 6, 16, 16)).astype(np.float32)
+    with open_ome_zarr(input_path, layout="hcs", mode="w", channel_names=["C"]) as dataset:
+        dataset.create_position("0", "0", "0").create_image(
+            "0", data, transform=[TransformationMeta(type="scale", scale=[1, 1, 0.25, 0.1, 0.1])]
+        )
+    position_path = input_path / "0" / "0" / "0"
+    tf_path = tmp_path / "tf.zarr"
+    volumes, peak_gpu_bytes = {}, {}
+    for device in ("cpu", "cuda"):
+        config_path = tmp_path / f"{device}.yml"
+        utils.model_to_yaml(
+            settings.ReconstructionSettings(
+                input_channel_names=["C"], reconstruction_dimension=3, device=device, **contrast
+            ),
+            config_path,
+        )
+        if not tf_path.exists():
+            CliRunner().invoke(
+                app,
+                ["compute-tf", "-i", str(position_path), "-c", str(config_path), "-o", str(tf_path)],
+                catch_exceptions=False,
+            )
+        torch.cuda.reset_peak_memory_stats()
+        volumes[device] = _reconstruct(position_path, tf_path, config_path, tmp_path / f"{device}.zarr")
+        peak_gpu_bytes[device] = torch.cuda.max_memory_allocated()
+
+    assert peak_gpu_bytes["cpu"] == 0
+    assert peak_gpu_bytes["cuda"] > 0
+    assert np.any(volumes["cpu"] != 0)
+    np.testing.assert_allclose(volumes["cuda"], volumes["cpu"], rtol=1e-4, atol=1e-5)
+
+
 def test_apply_inverse_czyx_labels_the_volume_for_the_model():
     import xarray as xr
 
diff --git a/waveorder/cli/apply_inverse_transfer_function.py b/waveorder/cli/apply_inverse_transfer_function.py
index 41e023fe..3851f8d1 100644
--- a/waveorder/cli/apply_inverse_transfer_function.py
+++ b/waveorder/cli/apply_inverse_transfer_function.py
@@ -335,7 +335,9 @@ def apply_inverse_transfer_function_single_position(
                 settings.input_channel_names,
             )
 
-    # Prepare the apply_inverse_model_function and its arguments
+    # Prepare the apply_inverse_model_function and its arguments. The phase and
+    # fluorescence models run on the config's `device`; the birefringence models
+    # run on the CPU only.
 
     # [biref only]
     if recon_biref and (not recon_phase):
@@ -360,6 +362,7 @@ def apply_inverse_transfer_function_single_position(
             "transfer_function": transfer_function,
             "recon_dim": recon_dim,
             "settings": settings.phase,
+            "device": settings.device,
         }
 
     # [biref and phase]
@@ -388,6 +391,7 @@ def apply_inverse_transfer_function_single_position(
             "recon_dim": recon_dim,
             "settings": settings.fluorescence,
             "fluor_channel_name": settings.input_channel_names[0],
+            "device": settings.device,
         }
 
     # Label each CZYX volume with the input's own coordinates, taken once here