feat: add Reconstruct button and optimize codec processing logic

app.py

@@ -43,6 +43,7 @@ def build_ui() -> gr.Blocks:
             type="filepath",
             label="Input audio",
         )
+        run_btn = gr.Button("Reconstruct", variant="primary")
 
         tab_components: dict[str, dict] = {}
 
@@ -87,75 +88,54 @@ def build_ui() -> gr.Blocks:
             c = tab_components[name]
             all_outputs.extend([c["audio_out"], c["stats_md"], c["spec_img"]])
 
-        def process_all(audio_path: str | None, current_tab: str, *dropdown_values):
-            """Generator that yields updates tab-by-tab, active tab first."""
+        def process_active(audio_path: str | None, current_tab: str, *dropdown_values):
+            """Reconstruct only the active tab's codec."""
             if audio_path is None:
-                return
+                return [gr.update()] * len(all_outputs)
 
             dd_map = dict(zip(ordered_names, dropdown_values))
-            order = [current_tab] + [n for n in ordered_names if n != current_tab]
-            results: dict[str, tuple] = {}
+            comps = tab_components[current_tab]
+            cfg_label = dd_map[current_tab]
+            cfg = next(c for c in comps["configs"] if c.name == cfg_label)
+
+            audio_out, sr, elapsed = _encode_decode_one(
+                Path(audio_path), current_tab, cfg
+            )
+            spec_path = make_spectrogram(audio_out, sr)
+
+            stats_text = (
+                f"**{elapsed:.2f}s**  |  "
+                f"{sr / 1000:.0f} kHz  |  "
+                f"{cfg_label}"
+            )
+
+            flat = []
             for name in ordered_names:
-                results[name] = (
-                    gr.update(value=None),
-                    gr.update(value="*Processing...*"),
-                    gr.update(value=None),
-                )
-
-            for name in order:
-                comps = tab_components[name]
-                cfg_label = dd_map[name]
-                cfg = next(c for c in comps["configs"] if c.name == cfg_label)
-
-                audio_out, sr, elapsed = _encode_decode_one(
-                    Path(audio_path), name, cfg
-                )
-                spec_path = make_spectrogram(audio_out, sr)
-
-                stats_text = (
-                    f"**{elapsed:.2f}s**  |  "
-                    f"{sr / 1000:.0f} kHz  |  "
-                    f"{cfg_label}"
-                )
-
-                results[name] = (
-                    gr.update(value=(sr, audio_out)),
-                    gr.update(value=stats_text),
-                    gr.update(value=str(spec_path)),
-                )
-
-                flat = []
-                for n in ordered_names:
-                    flat.extend(results[n])
-                yield flat
+                if name == current_tab:
+                    flat.extend(
+                        [
+                            gr.update(value=(sr, audio_out)),
+                            gr.update(value=stats_text),
+                            gr.update(value=str(spec_path)),
+                        ]
+                    )
+                else:
+                    flat.extend([gr.update(), gr.update(), gr.update()])
+            return flat
 
         all_dropdowns = [tab_components[n]["config_dd"] for n in ordered_names]
 
-        upload_event = audio_in.change(
-            fn=process_all,
+        btn_event = run_btn.click(
+            fn=process_active,
             inputs=[audio_in, active_tab] + all_dropdowns,
             outputs=all_outputs,
         )
 
-        dropdown_events = []
-        for name in ordered_names:
-            evt = tab_components[name]["config_dd"].change(
-                fn=process_all,
-                inputs=[audio_in, active_tab] + all_dropdowns,
-                outputs=all_outputs,
-            )
-            dropdown_events.append(evt)
-
         for codec_name, comps in tab_components.items():
             comps["tab"].select(
                 fn=lambda name=codec_name: name,
                 inputs=[],
                 outputs=[active_tab],
-            ).then(
-                fn=process_all,
-                inputs=[audio_in, active_tab] + all_dropdowns,
-                outputs=all_outputs,
-                cancels=[upload_event] + dropdown_events,
             )
 
     return demo
@@ -164,4 +144,4 @@ def build_ui() -> gr.Blocks:
 demo = build_ui()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

compare_codec/dac.py

@@ -9,6 +9,14 @@ import torch
 
 from compare_codec import CodecConfig, register
 
+_device = torch.device(
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
 
 class DACCodec:
     """DAC codec with lazy model loading."""
@@ -50,7 +58,7 @@ class DACCodec:
 
             model_path = _dac.utils.download(model_type=model_type)
             model = _dac.DAC.load(model_path)
-            model.eval()
+            model.eval().to(_device)
             self._models[model_type] = model
         return self._models[model_type]
 

compare_codec/encodec_codec.py

@@ -41,7 +41,9 @@ class EnCodecCodec:
         if self._model is None:
             from transformers import AutoProcessor, EncodecModel
 
-            self._model = EncodecModel.from_pretrained("facebook/encodec_24khz")
+            self._model = EncodecModel.from_pretrained(
+                "facebook/encodec_24khz", device_map="auto"
+            )
             self._model.eval()
             self._processor = AutoProcessor.from_pretrained("facebook/encodec_24khz")
 
@@ -62,6 +64,8 @@ class EnCodecCodec:
             sampling_rate=target_sr,
             return_tensors="pt",
         )
+        device = self._model.device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         enc = self._model.encode(
             inputs["input_values"],
             inputs["padding_mask"],
@@ -76,4 +80,4 @@ class EnCodecCodec:
         return audio_out.squeeze(0).squeeze(0).cpu().numpy()
 
 
-register(EnCodecCodec())
+register(EnCodecCodec())

compare_codec/mimi_codec.py

@@ -38,7 +38,7 @@ class MimiCodec:
         if self._model is None:
             from transformers import AutoFeatureExtractor, MimiModel
 
-            self._model = MimiModel.from_pretrained("kyutai/mimi")
+            self._model = MimiModel.from_pretrained("kyutai/mimi", device_map="auto")
             self._model.eval()
             self._fe = AutoFeatureExtractor.from_pretrained("kyutai/mimi")
 
@@ -59,14 +59,14 @@ class MimiCodec:
             sampling_rate=target_sr,
             return_tensors="pt",
         )
+        device = self._model.device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         enc = self._model.encode(inputs["input_values"], inputs["padding_mask"])
-        audio_out = self._model.decode(
-            enc.audio_codes, inputs["padding_mask"]
-        )[0]
+        audio_out = self._model.decode(enc.audio_codes, inputs["padding_mask"])[0]
 
         # Trim to original length (Mimi may pad).
         audio_out = audio_out.squeeze(0).squeeze(0).cpu().numpy()[:original_len]
         return audio_out
 
 
-register(MimiCodec())
+register(MimiCodec())

compare_codec/snac_codec.py

@@ -10,6 +10,14 @@ import torchaudio
 
 from compare_codec import CodecConfig, register
 
+_device = torch.device(
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
 _MODELS = [
     ("hubertsiuzdak/snac_24khz", 24_000),
     ("hubertsiuzdak/snac_32khz", 32_000),
@@ -47,7 +55,7 @@ class SNACCodec:
         if model_id not in self._models:
             from snac import SNAC
 
-            model = SNAC.from_pretrained(model_id)
+            model = SNAC.from_pretrained(model_id).to(_device)
             self._models[model_id] = model
         return self._models[model_id]
 
@@ -66,7 +74,7 @@ class SNACCodec:
         if sr != target_sr:
             wav = torchaudio.functional.resample(wav, sr, target_sr)
         # SNAC expects (B, 1, T).
-        wav = wav.unsqueeze(0)
+        wav = wav.unsqueeze(0).to(_device)
 
         audio_hat, _ = model(wav)
 

pyproject.toml

@@ -6,6 +6,7 @@ readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.12"
 dependencies = [
+    "accelerate>=1.13.0",
     "descript-audio-codec>=1.0.0",
     "encodec>=0.1.1",
     "gradio",

uv.lock

@@ -25,6 +25,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl", hash = "sha256:88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d", size = 135750, upload-time = "2026-01-28T10:17:04.19Z" },
 ]
 
+[[package]]
+name = "accelerate"
+version = "1.13.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "huggingface-hub" },
+    { name = "numpy" },
+    { name = "packaging" },
+    { name = "psutil" },
+    { name = "pyyaml" },
+    { name = "safetensors" },
+    { name = "torch" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl", hash = "sha256:cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0", size = 383744, upload-time = "2026-03-04T19:34:10.313Z" },
+]
+
 [[package]]
 name = "annotated-doc"
 version = "0.0.4"
@@ -347,6 +365,7 @@ name = "compare-codec"
 version = "0.1.0"
 source = { virtual = "." }
 dependencies = [
+    { name = "accelerate" },
     { name = "descript-audio-codec" },
     { name = "encodec" },
     { name = "gradio" },
@@ -367,6 +386,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
+    { name = "accelerate", specifier = ">=1.13.0" },
     { name = "descript-audio-codec", specifier = ">=1.0.0" },
     { name = "encodec", specifier = ">=0.1.1" },
     { name = "gradio" },
@@ -1918,6 +1938,34 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/32/27/1141a8232723dcb10a595cc0ce4321dcbbd5215300bf4acfc142343205bf/protobuf-3.19.6-py2.py3-none-any.whl", hash = "sha256:14082457dc02be946f60b15aad35e9f5c69e738f80ebbc0900a19bc83734a5a4", size = 162648, upload-time = "2022-09-29T22:07:20.303Z" },
 ]
 
+[[package]]
+name = "psutil"
+version = "7.2.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/51/08/510cbdb69c25a96f4ae523f733cdc963ae654904e8db864c07585ef99875/psutil-7.2.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b", size = 130595, upload-time = "2026-01-28T18:14:57.293Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/f5/97baea3fe7a5a9af7436301f85490905379b1c6f2dd51fe3ecf24b4c5fbf/psutil-7.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea", size = 131082, upload-time = "2026-01-28T18:14:59.732Z" },
+    { url = "https://files.pythonhosted.org/packages/37/d6/246513fbf9fa174af531f28412297dd05241d97a75911ac8febefa1a53c6/psutil-7.2.2-cp313-cp313t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63", size = 181476, upload-time = "2026-01-28T18:15:01.884Z" },
+    { url = "https://files.pythonhosted.org/packages/b8/b5/9182c9af3836cca61696dabe4fd1304e17bc56cb62f17439e1154f225dd3/psutil-7.2.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312", size = 184062, upload-time = "2026-01-28T18:15:04.436Z" },
+    { url = "https://files.pythonhosted.org/packages/16/ba/0756dca669f5a9300d0cbcbfae9a4c30e446dfc7440ffe43ded5724bfd93/psutil-7.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b", size = 139893, upload-time = "2026-01-28T18:15:06.378Z" },
+    { url = "https://files.pythonhosted.org/packages/1c/61/8fa0e26f33623b49949346de05ec1ddaad02ed8ba64af45f40a147dbfa97/psutil-7.2.2-cp313-cp313t-win_arm64.whl", hash = "sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9", size = 135589, upload-time = "2026-01-28T18:15:08.03Z" },
+    { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" },
+    { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" },
+    { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" },
+    { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" },
+    { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" },
+    { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" },
+    { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" },
+    { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" },
+    { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" },
+]
+
 [[package]]
 name = "ptyprocess"
 version = "0.7.0"