comfy-diffusion 2.1.1__tar.gz → 2.1.2__tar.gz

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-diffusion
-Version: 2.1.1
+Version: 2.1.2
 Summary: ComfyUI inference engine as a standalone Python library (no server, no UI).
 Author: Alejandro Quintero
 License: GPL-3.0-only

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/comfy_diffusion/audio.py

@@ -120,6 +120,37 @@ def ltxv_audio_vae_decode(vae: _LtxvAudioVaeDecoder, latent: Any) -> dict[str, A
     if getattr(latent_tensor, "is_nested", False):
         latent_tensor = latent_tensor.unbind()[-1]
 
+    first_stage_model = getattr(vae, "first_stage_model", None)
+    if first_stage_model is not None:
+        if hasattr(vae, "patcher"):
+            from ._runtime import ensure_comfyui_on_path
+
+            ensure_comfyui_on_path()
+            import comfy.model_management
+
+            memory_required = 0
+            if hasattr(vae, "memory_used_decode") and hasattr(latent_tensor, "shape"):
+                memory_required = vae.memory_used_decode(latent_tensor.shape, vae.vae_dtype)
+            comfy.model_management.load_models_gpu(
+                [vae.patcher],
+                memory_required=memory_required,
+                force_full_load=getattr(vae, "disable_offload", False),
+            )
+        if hasattr(latent_tensor, "to") and hasattr(vae, "device") and hasattr(vae, "vae_dtype"):
+            latent_tensor = latent_tensor.to(device=vae.device, dtype=vae.vae_dtype)
+        with _inference_mode_context():
+            audio = first_stage_model.decode(latent_tensor)
+        if hasattr(audio, "to"):
+            output_device = getattr(vae, "output_device", getattr(latent_tensor, "device", None))
+            if hasattr(vae, "vae_output_dtype"):
+                audio = audio.to(device=output_device, dtype=vae.vae_output_dtype(), copy=True)
+            elif output_device is not None:
+                audio = audio.to(output_device)
+        return {
+            "waveform": audio,
+            "sample_rate": int(first_stage_model.output_sample_rate),
+        }
+
     if hasattr(vae, "to"):
         vae.to("cpu")
     if hasattr(latent_tensor, "cpu"):

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/comfy_diffusion/models.py

@@ -259,8 +259,8 @@ class ModelManager:
         ensure_comfyui_on_path()
 
         import folder_paths
+        from comfy import sd as comfy_sd
         from comfy import utils as comfy_utils
-        from comfy.ldm.lightricks.vae.audio_vae import AudioVAE
 
         p = Path(path)
         if p.is_absolute() and p.is_file():
@@ -274,7 +274,14 @@ class ModelManager:
         state_dict, metadata = comfy_utils.load_torch_file(
             checkpoint_path, return_metadata=True
         )
-        return AudioVAE(state_dict, metadata)
+        state_dict = comfy_utils.state_dict_prefix_replace(
+            state_dict,
+            {"audio_vae.": "autoencoder.", "vocoder.": "vocoder."},
+            filter_keys=True,
+        )
+        vae = comfy_sd.VAE(sd=state_dict, metadata=metadata)
+        vae.throw_exception_if_invalid()
+        return vae
 
     def load_ltxav_text_encoder(
         self,

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/comfy_diffusion.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-diffusion
-Version: 2.1.1
+Version: 2.1.2
 Summary: ComfyUI inference engine as a standalone Python library (no server, no UI).
 Author: Alejandro Quintero
 License: GPL-3.0-only

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "comfy-diffusion"
-version = "2.1.1"
+version = "2.1.2"
 description = "ComfyUI inference engine as a standalone Python library (no server, no UI)."
 readme = "README.md"
 license = { text = "GPL-3.0-only" }

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/tests/test_audio.py

@@ -146,6 +146,47 @@ def test_ltxv_audio_vae_decode_calls_vae_decode_and_returns_audio_dict() -> None
     assert vae2.calls == [fake_latent_tensor]
 
 
+def test_ltxv_audio_vae_decode_supports_comfy_vae_wrapper() -> None:
+    class FakeAudioTensor:
+        def __init__(self) -> None:
+            self.calls: list[tuple[str, Any]] = []
+
+        def movedim(self, source: int, destination: int) -> "FakeAudioTensor":
+            self.calls.append(("movedim", (source, destination)))
+            return self
+
+        def to(self, device: Any) -> "FakeAudioTensor":
+            self.calls.append(("to", device))
+            return self
+
+    class FakeFirstStageModel:
+        output_sample_rate = 48000
+
+        def __init__(self) -> None:
+            self.calls: list[Any] = []
+
+        def decode(self, latent: Any) -> FakeAudioTensor:
+            self.calls.append(latent)
+            return fake_audio
+
+    class FakeLatentTensor:
+        device = "cuda:0"
+
+    fake_audio = FakeAudioTensor()
+    fake_latent = FakeLatentTensor()
+
+    class FakeComfyVae:
+        def __init__(self) -> None:
+            self.first_stage_model = FakeFirstStageModel()
+
+    vae = FakeComfyVae()
+    result = ltxv_audio_vae_decode(vae, {"samples": fake_latent})
+
+    assert result == {"waveform": fake_audio, "sample_rate": 48000}
+    assert vae.first_stage_model.calls == [fake_latent]
+    assert fake_audio.calls == [("to", "cuda:0")]
+
+
 def test_ltxv_empty_latent_audio_signature_matches_contract() -> None:
     signature = inspect.signature(ltxv_empty_latent_audio)
     assert (

{comfy_diffusion-2.1.1 → comfy_diffusion-2.1.2}/tests/test_model_manager_ltxv_audio_vae_loading.py

@@ -6,6 +6,7 @@ import sys
 from pathlib import Path
 from types import ModuleType
 from typing import Any
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -24,7 +25,9 @@ def _install_fake_ltxv_audio_vae_loader_modules(
         "add_model_folder_path": [],
         "get_full_path_or_raise": [],
         "load_torch_file": [],
-        "AudioVAE": [],
+        "state_dict_prefix_replace": [],
+        "VAE": [],
+        "throw_exception_if_invalid": [],
     }
 
     folder_paths_module = ModuleType("folder_paths")
@@ -50,33 +53,40 @@ def _install_fake_ltxv_audio_vae_loader_modules(
         return state_dict, metadata
 
     setattr(comfy_utils_module, "load_torch_file", load_torch_file)
+    remapped_state_dict = {"autoencoder.encoder.weight": object(), "vocoder.weight": object()}
 
-    audio_vae_module = ModuleType("comfy.ldm.lightricks.vae.audio_vae")
+    def state_dict_prefix_replace(
+        sd: dict[str, Any],
+        replace_prefix: dict[str, str],
+        *,
+        filter_keys: bool,
+    ) -> dict[str, Any]:
+        calls["state_dict_prefix_replace"].append((sd, replace_prefix, filter_keys))
+        return remapped_state_dict
 
-    def AudioVAE(sd: dict[str, Any], model_metadata: dict[str, Any] | None) -> Any:
-        calls["AudioVAE"].append((sd, model_metadata))
+    setattr(comfy_utils_module, "state_dict_prefix_replace", state_dict_prefix_replace)
+
+    comfy_sd_module = ModuleType("comfy.sd")
+
+    def throw_exception_if_invalid() -> None:
+        calls["throw_exception_if_invalid"].append(())
+
+    def make_vae(*, sd: dict[str, Any], metadata: dict[str, Any] | None) -> Any:
+        calls["VAE"].append((sd, metadata))
+        audio_vae_object.throw_exception_if_invalid.side_effect = throw_exception_if_invalid
         return audio_vae_object
 
-    setattr(audio_vae_module, "AudioVAE", AudioVAE)
+    setattr(comfy_sd_module, "VAE", make_vae)
 
     comfy_module = ModuleType("comfy")
-    comfy_ldm_module = ModuleType("comfy.ldm")
-    comfy_lightricks_module = ModuleType("comfy.ldm.lightricks")
-    comfy_lightricks_vae_module = ModuleType("comfy.ldm.lightricks.vae")
 
     setattr(comfy_module, "utils", comfy_utils_module)
-    setattr(comfy_module, "ldm", comfy_ldm_module)
-    setattr(comfy_ldm_module, "lightricks", comfy_lightricks_module)
-    setattr(comfy_lightricks_module, "vae", comfy_lightricks_vae_module)
-    setattr(comfy_lightricks_vae_module, "audio_vae", audio_vae_module)
+    setattr(comfy_module, "sd", comfy_sd_module)
 
     monkeypatch.setitem(sys.modules, "folder_paths", folder_paths_module)
     monkeypatch.setitem(sys.modules, "comfy", comfy_module)
     monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils_module)
-    monkeypatch.setitem(sys.modules, "comfy.ldm", comfy_ldm_module)
-    monkeypatch.setitem(sys.modules, "comfy.ldm.lightricks", comfy_lightricks_module)
-    monkeypatch.setitem(sys.modules, "comfy.ldm.lightricks.vae", comfy_lightricks_vae_module)
-    monkeypatch.setitem(sys.modules, "comfy.ldm.lightricks.vae.audio_vae", audio_vae_module)
+    monkeypatch.setitem(sys.modules, "comfy.sd", comfy_sd_module)
 
     return calls
 
@@ -97,7 +107,7 @@ def test_load_ltxv_audio_vae_calls_loader_and_returns_raw_object(
 
     fake_state_dict = {"audio_vae.encoder.weight": object()}
     fake_metadata = {"format": "safetensors"}
-    fake_audio_vae = object()
+    fake_audio_vae = MagicMock()
 
     calls = _install_fake_ltxv_audio_vae_loader_modules(
         monkeypatch,
@@ -113,7 +123,18 @@ def test_load_ltxv_audio_vae_calls_loader_and_returns_raw_object(
     assert result is fake_audio_vae
     assert calls["get_full_path_or_raise"] == []
     assert calls["load_torch_file"] == [(str(audio_vae_file.resolve()), True)]
-    assert calls["AudioVAE"] == [(fake_state_dict, fake_metadata)]
+    assert calls["state_dict_prefix_replace"] == [
+        (
+            fake_state_dict,
+            {"audio_vae.": "autoencoder.", "vocoder.": "vocoder."},
+            True,
+        )
+    ]
+    assert len(calls["VAE"]) == 1
+    assert calls["VAE"][0][0] is not fake_state_dict
+    assert set(calls["VAE"][0][0]) == {"autoencoder.encoder.weight", "vocoder.weight"}
+    assert calls["VAE"][0][1] == fake_metadata
+    assert calls["throw_exception_if_invalid"] == [()]
     assert calls["add_model_folder_path"] == [
         ("checkpoints", str(checkpoints_dir), True),
         ("embeddings", str(embeddings_dir), True),
@@ -138,7 +159,7 @@ def test_load_ltxv_audio_vae_is_callable_from_models_import(
     (models_dir / "embeddings").mkdir(parents=True)
 
     expected_resolved = str((models_dir / "checkpoints" / "audio_vae.safetensors").resolve())
-    fake_audio_vae = object()
+    fake_audio_vae = MagicMock()
     calls = _install_fake_ltxv_audio_vae_loader_modules(
         monkeypatch,
         state_dict={"weights": object()},
@@ -167,7 +188,7 @@ def test_load_ltxv_audio_vae_raises_file_not_found_before_loader(
         monkeypatch,
         state_dict={},
         metadata=None,
-        audio_vae_object=object(),
+        audio_vae_object=MagicMock(),
         resolved_checkpoint_path="/unused/path.safetensors",
     )
 
@@ -180,4 +201,4 @@ def test_load_ltxv_audio_vae_raises_file_not_found_before_loader(
     assert expected_path in str(exc_info.value)
     assert calls["get_full_path_or_raise"] == []
     assert calls["load_torch_file"] == []
-    assert calls["AudioVAE"] == []
+    assert calls["VAE"] == []