diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmModels.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmModels.py new file mode 100644 index 000000000..74ab55638 --- /dev/null +++ b/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, +} diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmVideo.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmVideo.py index 2e5ec8908..8df773268 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmVideo.py +++ b/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, } diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/__init__.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/__init__.py index 366209ab3..9cc7e15da 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/__init__.py +++ b/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: diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs index 97ba2aa02..83ffbb538 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs @@ -48,6 +48,9 @@ public bool IsSD3() /// Returns true if the current model is Lightricks LTX Video 2. public bool IsLTXV2() => IsModelCompatClass(T2IModelClassSorter.CompatLtxv2); + /// Returns true if the current model is Lightricks LTX Video 2.3. + public bool IsLTXV23() => CurrentModelClass()?.ID == "lightricks-ltx-video-2-3"; + /// Returns true if the current model is Black Forest Labs' Flux.1. 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 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 { diff --git a/src/Text2Image/CommonModels.cs b/src/Text2Image/CommonModels.cs index 7f0fd349d..4c4442791 100644 --- a/src/Text2Image/CommonModels.cs +++ b/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")); diff --git a/src/Text2Image/T2IModelClassSorter.cs b/src/Text2Image/T2IModelClassSorter.cs index fdf807f59..514dc1917 100644 --- a/src/Text2Image/T2IModelClassSorter.cs +++ b/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) => {