diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index a78662b8d..286f46154 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -769,9 +769,9 @@ public List LoadPromptImagesForMainRef(List images) public (JArray, JArray, JArray, JArray) BuildInputImageHandling(List images, JArray pos, JArray neg, JArray latent) { JArray imgNeg = null; - if (IsKontext() || IsOmniGen() || IsQwenImage() || IsAnyFlux2() || IsBoogu()) + if (IsKontext() || IsOmniGen() || IsQwenImage() || IsAnyFlux2() || IsBoogu() || IsMageFlow()) { - if (IsOmniGen() || IsQwenImageEditPlus() || IsBoogu()) + if (IsOmniGen() || IsQwenImageEditPlus() || IsBoogu() || IsMageFlow()) { imgNeg = neg; } @@ -803,7 +803,7 @@ void makeRefLatent(JArray image) } if (img is not null) { - if (IsQwenImageEditPlus() || IsBoogu()) + if (IsQwenImageEditPlus() || IsBoogu() || IsMageFlow()) { neg = imgNeg; } @@ -975,7 +975,7 @@ public string CreateKSampler(JArray model, JArray pos, JArray neg, JArray latent defscheduler ??= "simple"; } // TODO: Registry of model default preferences instead of this - else if (IsFlux() || IsWanVideo() || IsWanVideo22() || IsOmniGen() || IsQwenImage() || IsZImage() || IsZetaChroma() || IsErnie() || IsHiDreamO1() || IsLens() || IsPixelDiT() || IsKrea2()) + else if (IsFlux() || IsWanVideo() || IsWanVideo22() || IsOmniGen() || IsQwenImage() || IsZImage() || IsZetaChroma() || IsErnie() || IsHiDreamO1() || IsLens() || IsPixelDiT() || IsKrea2() || IsMageFlow()) { defscheduler ??= "simple"; } @@ -1305,7 +1305,7 @@ public JArray GetPromptImage(bool fixSize, bool promptSize = false, int index = doesFit = false; } } - else if ((IsBoogu() || IsQwenImageEditPlus()) && promptSize) + else if ((IsBoogu() || IsQwenImageEditPlus() || IsMageFlow()) && promptSize) { target = 384; doesFit = false; @@ -2353,7 +2353,7 @@ public JArray CreateConditioningDirect(string prompt, JArray clip, T2IModel mode ["text"] = prompt }, id); } - else if (IsIdeogram4() || IsKrea2() || (IsBoogu() && isPositive)) + else if (IsIdeogram4() || IsKrea2() || (IsBoogu() && isPositive) || IsMageFlow()) { JArray imageNode = GetPromptImage(true, true, 0); for (int i = 1; i < 10; i++) diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs index 83ffbb538..54be505c4 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs @@ -218,6 +218,9 @@ public bool IsNvidiaCosmos2() /// Returns true if the current model is any Kandinsky 5 variant. public bool IsKandinsky5() => IsKandinsky5ImgLite() || IsKandinsky5VidLite() || IsKandinsky5VidPro(); + /// Returns true if the current model is Mage-Flow. + public bool IsMageFlow() => IsModelCompatClass(T2IModelClassSorter.CompatMageFlow); + /// Returns true if the current model is any Wan-2.1 variant. public bool IsWanVideo() { @@ -1169,6 +1172,11 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) helpers.LoadClip("boogu", helpers.GetQwen3vl_8bModel()); helpers.DoVaeLoader(UserInput.SourceSession?.User?.Settings?.VAEs?.DefaultFluxVAE, "flux-1", "flux-ae"); } + else if (IsMageFlow()) + { + helpers.LoadClip("mage", helpers.GetQwen3vl_4bModel()); + helpers.DoVaeLoader(null, "mage-flow", "mage-flow-vae"); + } else if (IsFlux() && (LoadingClip is null || LoadingVAE is null || UserInput.Get(T2IParamTypes.T5XXLModel) is not null || UserInput.Get(T2IParamTypes.ClipLModel) is not null)) { helpers.LoadClip2("flux", helpers.GetT5XXLModel(), helpers.GetClipLModel()); diff --git a/src/Text2Image/CommonModels.cs b/src/Text2Image/CommonModels.cs index 772145796..1c5c8ce79 100644 --- a/src/Text2Image/CommonModels.cs +++ b/src/Text2Image/CommonModels.cs @@ -92,6 +92,7 @@ public static void RegisterCoreSet() 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("mage-flow-vae", "Mage-VAE", "The VAE for Mage-Flow", "https://huggingface.co/microsoft/Mage-Flow/resolve/main/vae/diffusion_pytorch_model.safetensors", "34e076dc1e8a15321e1e07be5111d59cf16dd10b804b7c7e20b4de29013427e0", "VAE", "MageFlow/mage_flow_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")); Register(new("hunyuan-video-1_5-vae", "Hunyuan Video 1.5 VAE", "The VAE for Hunyuan Video 1.5", "https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/resolve/main/split_files/vae/hunyuanvideo15_vae_fp16.safetensors", "e7c3091949c27e2d55ae6d5df917b99dadfebbf308e5a50d0ade0d16c90297ae", "VAE", "HunyuanVideo/hunyuanvideo15_vae_fp16.safetensors")); diff --git a/src/Text2Image/T2IModelClassSorter.cs b/src/Text2Image/T2IModelClassSorter.cs index 904098c39..18d841b8a 100644 --- a/src/Text2Image/T2IModelClassSorter.cs +++ b/src/Text2Image/T2IModelClassSorter.cs @@ -95,6 +95,7 @@ public static T2IModelCompatClass CompatIdeogram4 = RegisterCompat(new() { ID = "ideogram-4", ShortCode = "Ideo4", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatKrea2 = RegisterCompat(new() { ID = "krea-2", ShortCode = "Krea2", VaeFamily = VaeQwenImage }), CompatBoogu = RegisterCompat(new() { ID = "boogu", ShortCode = "Boogu", LorasTargetTextEnc = false, VaeFamily = VaeFlux1 }), + CompatMageFlow = RegisterCompat(new() { ID = "mage-flow", ShortCode = "Mage", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), // Audio models CompatAceStep15 = RegisterCompat(new() { ID = "ace-step-1_5", ShortCode = "Ace15", IsAudioModel = true }), // Obscure old random ones @@ -245,6 +246,7 @@ bool isZImageLora(JObject h) => (hasLoraKey(h, "layers.0.adaLN_modulation.0") && bool isQwenImage(JObject h) => (h.ContainsKey("time_text_embed.timestep_embedder.linear_1.bias") && h.ContainsKey("img_in.bias") && (h.ContainsKey("transformer_blocks.0.attn.add_k_proj.bias") || h.ContainsKey("transformer_blocks.0.attn.add_qkv_proj.bias"))) || (h.ContainsKey("model.diffusion_model.time_text_embed.timestep_embedder.linear_1.bias") && h.ContainsKey("model.diffusion_model.img_in.bias") && (h.ContainsKey("model.diffusion_model.transformer_blocks.0.attn.add_k_proj.bias") || h.ContainsKey("model.diffusion_model.transformer_blocks.0.attn.add_qkv_proj.bias"))); bool isQwenImageEdit2511(JObject h) => h.ContainsKey("__index_timestep_zero__"); + bool isMageFlow(JObject h) => tryGetKey(h, "txt_norm.weight", out JToken tok) && tok["shape"].ToArray()[0].Value() == 2560; bool isQwenImageLora(JObject h) => (hasLoraKey(h, "transformer_blocks.0.attn.add_k_proj") && hasLoraKey(h, "transformer_blocks.0.img_mlp.net.0.proj")) || (hasLoraKey(h, "transformer_blocks.0.attn.add_k_proj") && hasLoraKey(h, "transformer_blocks.0.img_mlp.net.2")) || (hasLoraKey(h, "transformer_blocks.59.attn.to_k") && hasLoraKey(h, "transformer_blocks.59.attn.to_out.0")) @@ -666,7 +668,7 @@ JToken GetEmbeddingKey(JObject h) // ====================== Qwen Image ====================== Register(new() { ID = "qwen-image", CompatClass = CompatQwenImage, Name = "Qwen Image", StandardWidth = 1328, StandardHeight = 1328, IsThisModelOfClass = (m, h) => { - return isQwenImage(h) && !isControlnetX(h) && !isSD3Controlnet(h) && !isQwenImageEdit2511(h); + return isQwenImage(h) && !isControlnetX(h) && !isSD3Controlnet(h) && !isQwenImageEdit2511(h) && !isMageFlow(h); }}); Register(new() { ID = "qwen-image-edit", CompatClass = CompatQwenImage, Name = "Qwen Image Edit", StandardWidth = 1328, StandardHeight = 1328, IsThisModelOfClass = (m, h) => { @@ -697,6 +699,11 @@ JToken GetEmbeddingKey(JObject h) { return isErnie(h); }}); + // ====================== Mage Flow ====================== + Register(new() { ID = "mage-flow", CompatClass = CompatMageFlow, Name = "Mage Flow", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => + { + return isMageFlow(h) && isQwenImage(h) && !isControlnetX(h) && !isSD3Controlnet(h) && !isQwenImageEdit2511(h); + }}); // ====================== Kandinsky5 ====================== Register(new() { ID = "kandinsky5-image-lite", CompatClass = CompatKandinsky5ImgLite, Name = "Kandinsky5 Image Lite", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => {