Attempting a fix for the audio switch

README.md

@@ -9,6 +9,10 @@ ComfyUI custom nodes for returning the first available value from six optional i
   - Returns the first non-empty input in order
   - Preserves original behavior for empty model/clip context dictionaries
   - Uses wildcard typing so it can pass through most ComfyUI value types
+- `SimpleLatentSwitch`
+  - Accepts six optional `LATENT` inputs
+  - Returns the first non-empty latent in order
+  - Keeps LTX latents on a typed `LATENT` path instead of a wildcard path
 - `SimpleAudioLatentSwitch`
   - Accepts six optional `LATENT` inputs
   - Returns the first latent compatible with LTX audio decode
@@ -36,7 +40,7 @@ Restart ComfyUI after installing.
 
 ## Usage
 
-1. Add `SimpleSwitch` for general pass-through selection, or `SimpleAudioLatentSwitch` for LTX audio pipelines.
+1. Add `SimpleSwitch` for general pass-through selection, `SimpleLatentSwitch` for LTX latent selection, or `SimpleAudioLatentSwitch` right before `LTXV Audio VAE Decode`.
 2. Connect values to `input01` ... `input06` in your preferred priority order.
 3. Use `output` downstream.
 
@@ -44,11 +48,12 @@ Restart ComfyUI after installing.
 
 LTX video latents and LTX audio latents both use the ComfyUI `LATENT` socket type, but
 the audio decoder only accepts audio-shaped latents or nested AV latents. If you route
-an LTX audio path through a generic wildcard switch, it is easy to accidentally forward
-an incompatible latent and only discover it at decode time.
+an LTX latent through the generic wildcard `SimpleSwitch`, it can lose the typed latent
+path and you may only discover the mismatch at decode time.
 
-Use `SimpleAudioLatentSwitch` before `LTXV Audio VAE Decode` when selecting between LTX
-audio latent branches.
+Use `SimpleLatentSwitch` anywhere you are switching LTX latents in general, and use
+`SimpleAudioLatentSwitch` immediately before `LTXV Audio VAE Decode` when selecting
+between LTX audio latent branches.
 
 ## License
 

__init__.py

@@ -1,12 +1,14 @@
-from .simple_switch import SimpleAudioLatentSwitch, SimpleSwitch
+from .simple_switch import SimpleAudioLatentSwitch, SimpleLatentSwitch, SimpleSwitch
 
 NODE_CLASS_MAPPINGS = {
     "SimpleSwitch": SimpleSwitch,
+    "SimpleLatentSwitch": SimpleLatentSwitch,
     "SimpleAudioLatentSwitch": SimpleAudioLatentSwitch,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "SimpleSwitch": "Simple Switch (6 Inputs)",
+    "SimpleLatentSwitch": "Simple Latent Switch (6 Inputs)",
     "SimpleAudioLatentSwitch": "Simple Audio Latent Switch (6 Inputs)",
 }
 

simple_switch.py

@@ -38,6 +38,14 @@ def _is_nested_tensor(value):
     return bool(getattr(value, "is_nested", False))
 
 
+def _get_samples_ndim(samples):
+    return getattr(samples, "ndim", None)
+
+
+def _is_latent(value):
+    return _get_latent_samples(value) is not None
+
+
 def _is_audio_latent(value):
     if not isinstance(value, dict):
         return False
@@ -46,15 +54,19 @@ def _is_audio_latent(value):
     if samples is None:
         return False
 
-    if value.get("type") == "audio":
-        return True
-
     # LTX AV latents are valid for audio decode because the decoder unwraps the
     # nested container and selects the audio branch internally.
     if _is_nested_tensor(samples):
         return True
 
-    return "sample_rate" in value and getattr(samples, "ndim", None) == 4
+    ndim = _get_samples_ndim(samples)
+    if ndim != 4:
+        return False
+
+    if value.get("type") == "audio":
+        return True
+
+    return "sample_rate" in value
 
 
 def _describe_audio_latent(value, index):
@@ -112,6 +124,44 @@ class SimpleSwitch:
         return (None,)
 
 
+class SimpleLatentSwitch:
+    """Return the first non-empty latent from input01..input06."""
+
+    CATEGORY = "Simple Switch"
+    FUNCTION = "switch"
+    RETURN_TYPES = (LATENT_TYPE,)
+    RETURN_NAMES = ("output",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "optional": {
+                "input01": (LATENT_TYPE,),
+                "input02": (LATENT_TYPE,),
+                "input03": (LATENT_TYPE,),
+                "input04": (LATENT_TYPE,),
+                "input05": (LATENT_TYPE,),
+                "input06": (LATENT_TYPE,),
+            },
+        }
+
+    def switch(
+        self,
+        input01=None,
+        input02=None,
+        input03=None,
+        input04=None,
+        input05=None,
+        input06=None,
+    ):
+        for candidate in (input01, input02, input03, input04, input05, input06):
+            if not _is_none(candidate) and _is_latent(candidate):
+                return (candidate,)
+
+        return (None,)
+
+
 class SimpleAudioLatentSwitch:
     """Return the first non-empty audio-compatible latent from input01..input06."""