Adds support for split LTX 2.3 model

src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmModels.py

@@ -0,0 +1,30 @@
+import folder_paths
+import comfy.utils, comfy.sd
+from comfy_api.latest import io
+
+class SwarmLTXVAudioVAELoader(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="SwarmLTXVAudioVAELoader",
+            display_name="Swarm LTXV Audio VAE Loader",
+            category="SwarmUI/video",
+            description="Loads an LTX-2 audio VAE from the VAE models folder.",
+            inputs=[
+                io.Combo.Input("vae_name", options=folder_paths.get_filename_list("vae"), tooltip="Audio VAE file."),
+            ],
+            outputs=[io.Vae.Output(display_name="Audio VAE")],
+        )
+
+    @classmethod
+    def execute(cls, vae_name: str) -> io.NodeOutput:
+        vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
+        sd, metadata = comfy.utils.load_torch_file(vae_path, return_metadata=True)
+        sd = comfy.utils.state_dict_prefix_replace(sd, {"audio_vae.": "autoencoder."})
+        vae = comfy.sd.VAE(sd=sd, metadata=metadata)
+        vae.throw_exception_if_invalid()
+        return io.NodeOutput(vae)
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmLTXVAudioVAELoader": SwarmLTXVAudioVAELoader,
+}

src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmVideo.py

@@ -84,7 +84,6 @@ def _sample_linear(cls, source_frames: torch.Tensor, source_positions: torch.Ten
         upper_frames = source_frames[upper_idx]
         return ((1.0 - blend_weight) * lower_frames + blend_weight * upper_frames).contiguous()
 
-
 NODE_CLASS_MAPPINGS = {
     "SwarmVideoResampleFPS": SwarmVideoResampleFPS,
 }

src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/__init__.py

@@ -1,6 +1,6 @@
 import os, folder_paths, traceback
 
-from . import SwarmBlending, SwarmImages, SwarmInternalUtil, SwarmKSampler, SwarmLoadImageB64, SwarmLoraLoader, SwarmMasks, SwarmSaveImageWS, SwarmTiling, SwarmExtractLora, SwarmUnsampler, SwarmLatents, SwarmInputNodes, SwarmTextHandling, SwarmReference, SwarmMath, SwarmSam2, SwarmAudio, SwarmVideo
+from . import SwarmBlending, SwarmImages, SwarmInternalUtil, SwarmKSampler, SwarmLoadImageB64, SwarmLoraLoader, SwarmMasks, SwarmSaveImageWS, SwarmTiling, SwarmExtractLora, SwarmUnsampler, SwarmLatents, SwarmInputNodes, SwarmTextHandling, SwarmReference, SwarmMath, SwarmSam2, SwarmAudio, SwarmVideo, SwarmModels
 
 WEB_DIRECTORY = "./web"
 
@@ -24,6 +24,7 @@
     | SwarmSam2.NODE_CLASS_MAPPINGS
     | SwarmAudio.NODE_CLASS_MAPPINGS
     | SwarmVideo.NODE_CLASS_MAPPINGS
+    | SwarmModels.NODE_CLASS_MAPPINGS
 )
 
 try:

src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs

@@ -48,6 +48,9 @@ public bool IsSD3()
     /// <summary>Returns true if the current model is Lightricks LTX Video 2.</summary>
     public bool IsLTXV2() => IsModelCompatClass(T2IModelClassSorter.CompatLtxv2);
 
+    /// <summary>Returns true if the current model is Lightricks LTX Video 2.3.</summary>
+    public bool IsLTXV23() => CurrentModelClass()?.ID == "lightricks-ltx-video-2-3";
+
     /// <summary>Returns true if the current model is Black Forest Labs' Flux.1.</summary>
     public bool IsFlux() => IsModelCompatClass(T2IModelClassSorter.CompatFlux);
 
@@ -532,6 +535,22 @@ public void AudioVaeLoad(string ckpt)
             g.CurrentAudioVae = new WGNodeData([avaeLoader, 0], g, WGNodeData.DT_AUDIOVAE, g.CurrentCompat());
         }
 
+        public void LTXAudioVaeLoad(string knownName)
+        {
+            CommonModels.ModelInfo knownFile = CommonModels.Known[knownName];
+            string vaeFile = knownFile.FileName;
+            if (!Program.T2IModelSets["VAE"].Models.ContainsKey(vaeFile))
+            {
+                knownFile.DownloadNow().Wait();
+                Program.RefreshAllModelSets();
+            }
+            string avaeLoader = g.CreateNode("SwarmLTXVAudioVAELoader", new JObject()
+            {
+                ["vae_name"] = vaeFile.Replace('\\', '/').Replace("/", g.ModelFolderFormat ?? $"{Path.DirectorySeparatorChar}")
+            });
+            g.CurrentAudioVae = new WGNodeData([avaeLoader, 0], g, WGNodeData.DT_AUDIOVAE, g.CurrentCompat());
+        }
+
         public string RequireClipModel(string name, string url, string hash, T2IRegisteredParam<T2IModel> param)
         {
             if (param is not null && g.UserInput.TryGet(param, out T2IModel model))
@@ -699,6 +718,12 @@ public string GetLTX2EmbedClip()
             return RequireClipModel("ltx2/ltx2-embeddings-connector-distill.safetensors", "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/text_encoders/ltx-2-19b-embeddings_connector_distill_bf16.safetensors", "8990ec3fe88396ca33ac1795c89b1771d88190e51e24084b21f54b25399acbed", null);
         }
 
+        public string GetLTX23TextProjectionClip()
+        {
+            // TODO: Still cursed!
+            return RequireClipModel("LTX-2/ltx-2.3_text_projection_bf16.safetensors", "https://huggingface.co/Kijai/LTX2.3_comfy/resolve/main/text_encoders/ltx-2.3_text_projection_bf16.safetensors", "911d59bb4cb7708179c9a0045ea0fe41212ecfb77aed3a02702b7c0a8274911f", null);
+        }
+
         public void LoadClip(string type, string model)
         {
             string loaderType = "CLIPLoader";
@@ -1277,12 +1302,13 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC)
         {
             if (LoadingVAE is null)
             {
-                // Hypothetical approximation of what would probably be right if comfy wasn't just entirely broken on handling this
-                helpers.LoadClip2("ltxv", helpers.GetGemma3_12bModel(), helpers.GetLTX2EmbedClip());
-                helpers.DoVaeLoader(null, (string)null, "ltx2-audio-vae");
-                CurrentAudioVae = new WGNodeData([LoadingVAE, 0], this, WGNodeData.DT_AUDIOVAE, CurrentCompat());
-                helpers.DoVaeLoader(null, "lightricks-ltx-video-2", "ltx2-video-vae");
-                throw new SwarmUserErrorException("LTX2 requires the safetensors checkpoint format currently due to comfy limitations.");
+                if (!IsLTXV23())
+                {
+                    throw new SwarmUserErrorException("LTX2 requires the safetensors checkpoint format currently due to comfy limitations.");
+                }
+                helpers.LoadClip2("ltxv", helpers.GetGemma3_12bModel(), helpers.GetLTX23TextProjectionClip());
+                helpers.DoVaeLoader(null, "lightricks-ltx-video-2", "ltx2-3-video-vae");
+                helpers.LTXAudioVaeLoad("ltx2-3-audio-vae");
             }
             else
             {

src/Text2Image/CommonModels.cs

@@ -89,6 +89,8 @@ public static void RegisterCoreSet()
         Register(new("ltxv-vae", "LTX-V VAE", "The VAE for Lightricks LTX-Video.", "https://huggingface.co/wsbagnsv1/ltxv-13b-0.9.7-dev-GGUF/resolve/c4296d06bab7719ce08e68bfa7a35042898e538b/ltxv-13b-0.9.7-vae-BF16.safetensors", "ee5ddcebc0b92d81b8aed9ee43445b7a4e66df1acf180678c5aa40e82f898dc5", "VAE", "LTXV/ltxv_vae.safetensors"));
         Register(new("ltx2-video-vae", "LTX-2 Video VAE", "The video VAE for Lightricks LTX-2.", "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/VAE/LTX2_video_vae_bf16.safetensors", "8ea083ef4e7118c7fb5c78834fc6b76106bd31464b0c7912bb9dbc7557b0b91d", "VAE", "LTX2/LTX2_video_vae.safetensors"));
         Register(new("ltx2-audio-vae", "LTX-2 Audio VAE", "The audio VAE for Lightricks LTX-2.", "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/VAE/LTX2_audio_vae_bf16.safetensors", "f00fcace3ef1ebfa6dca9fcb4b925df47d57f826ae7df3bc8fb98f577f213204", "VAE", "LTX2/LTX2_audio_vae.safetensors"));
+        Register(new("ltx2-3-video-vae", "LTX-2.3 Video VAE", "The video VAE for Lightricks LTX-2.3.", "https://huggingface.co/Kijai/LTX2.3_comfy/resolve/main/vae/LTX23_video_vae_bf16.safetensors", "01ea62d09bc139f95c5dee7b5c062ad6a3e6cd8be910a1983ac02e7eb5b8ee3b", "VAE", "LTX-2/LTX23_video_vae_bf16.safetensors"));
+        Register(new("ltx2-3-audio-vae", "LTX-2.3 Audio VAE", "The audio VAE for Lightricks LTX-2.3.", "https://huggingface.co/Kijai/LTX2.3_comfy/resolve/main/vae/LTX23_audio_vae_bf16.safetensors", "5bc10fa4adecf99dda132d916e23048cbd56797702c5fa50eb5d2079048a38c3", "VAE", "LTX-2/LTX23_audio_vae_bf16.safetensors"));
         Register(new("qwen-image-vae", "Qwen Image VAE", "The VAE for Qwen Image", "https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors", "a70580f0213e67967ee9c95f05bb400e8fb08307e017a924bf3441223e023d1f", "VAE", "QwenImage/qwen_image_vae.safetensors"));
         Register(new("hunyuan-image-2_1-vae", "Hunyuan Image 2.1 VAE", "The VAE for Hunyuan Image 2.1 Base", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_2.1_vae_fp16.safetensors", "f2ae19863609206196b5e3a86bfd94f67bd3866f5042004e3994f07e3c93b2f9", "VAE", "HunyuanImage/hunyuan_image_2.1_vae_fp16.safetensors"));
         Register(new("hunyuan-image-2_1-refiner-vae", "Hunyuan Image 2.1 Refiner VAE", "The VAE for Hunyuan Image 2.1 Refiner", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_refiner_vae_fp16.safetensors", "e1b74e85d61b65e18cc05ca390e387d93cfadf161e737de229ebb800ea3db769", "VAE", "HunyuanImage/hunyuan_image_2.1_refiner_vae_fp16.safetensors"));

src/Text2Image/T2IModelClassSorter.cs

@@ -194,8 +194,10 @@ bool isKrea2Lora(JObject h) => (hasLoraKey(h, "blocks.0.attn.gate") && hasLoraKe
         bool isMochiVae(JObject h) => h.ContainsKey("encoder.layers.4.layers.1.attn_block.attn.qkv.weight") || h.ContainsKey("layers.4.layers.1.attn_block.attn.qkv.weight") || h.ContainsKey("blocks.2.blocks.3.stack.5.weight") || h.ContainsKey("decoder.blocks.2.blocks.3.stack.5.weight");
         bool isLtxv(JObject h) => hasKey(h, "adaln_single.emb.timestep_embedder.linear_1.bias");
         bool isLtxvVae(JObject h) => h.ContainsKey("decoder.conv_in.conv.bias") && h.ContainsKey("decoder.last_time_embedder.timestep_embedder.linear_1.bias");
+        bool isLtxv23Vae(JObject h) => (h.ContainsKey("decoder.conv_in.conv.bias") && !h.ContainsKey("decoder.last_time_embedder.timestep_embedder.linear_1.bias") && !h.ContainsKey("per_channel_statistics.channel"))
+            || (h.ContainsKey("audio_vae.decoder.conv_in.conv.bias") && h.ContainsKey("vocoder.mel_stft.mel_basis") && !isLtxv2(h));
         bool isLtxv2(JObject h) => hasKey(h, "transformer_blocks.1.audio_to_video_attn.k_norm.weight");
-        bool isLtxv23(JObject h) => hasKey(h, "text_embedding_projection.audio_aggregate_embed.weight");
+        bool isLtxv23(JObject h) => hasKey(h, "text_embedding_projection.audio_aggregate_embed.weight") || hasKey(h, "transformer_blocks.1.prompt_scale_shift_table");
         bool isLtxv2Lora(JObject h) => (hasLoraKey(h, "transformer_blocks.0.attn1.to_k") && hasLoraKey(h, "transformer_blocks.0.attn1.to_out.0") && hasLoraKey(h, "transformer_blocks.9.attn2.to_v"))
             || (hasLoraKey(h, "transformer_blocks.0.audio_attn1.to_k") && hasLoraKey(h, "transformer_blocks.0.audio_attn1.to_out.0") && hasLoraKey(h, "transformer_blocks.9.audio_attn1.to_v"));
         bool isSana(JObject h) => h.ContainsKey("attention_y_norm.weight") && h.ContainsKey("blocks.0.attn.proj.weight");
@@ -740,6 +742,10 @@ JToken GetEmbeddingKey(JObject h)
         {
             return isLtxv2(h) && isLtxv23(h);
         }});
+        Register(new() { ID = "lightricks-ltx-video-2-3/vae", CompatClass = CompatLtxv2, Name = "Lightricks LTX Video 2.3 VAE", StandardWidth = 960, StandardHeight = 960, IsThisModelOfClass = (m, h) =>
+        {
+            return isLtxv23Vae(h);
+        }});
         // ====================== Random Other Models ======================
         Register(new() { ID = "chroma", CompatClass = CompatChroma, Name = "Chroma", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) =>
         {