Adds YuE2 support

src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs

@@ -1014,6 +1014,11 @@ public string CreateKSampler(JArray model, JArray pos, JArray neg, JArray latent
             defsampler ??= "res_multistep";
             defscheduler ??= "simple";
         }
+        else if (IsYue2())
+        {
+            defsampler ??= "dpm_2";
+            defscheduler ??= "sgm_uniform";
+        }
         else if (IsAnima())
         {
             defsampler ??= "er_sde";
@@ -2560,6 +2565,36 @@ public JArray CreateConditioningDirect(string prompt, JArray clip, T2IModel mode
                 ["min_p"] = 0
             }, id);
         }
+        else if (IsYue2())
+        {
+            if (!isPositive)
+            {
+                return FinalPrompt;
+            }
+            string abc = CreateNode("YuE2GenerateABC", new JObject()
+            {
+                ["clip"] = clip,
+                ["style"] = UserInput.Get(T2IParamTypes.Text2AudioStyle, ""),
+                ["lyrics"] = prompt,
+                ["seed"] = UserInput.Get(T2IParamTypes.Seed, 0) + 10,
+                ["mode"] = "full",
+                ["max_abc_tokens"] = 1024
+            });
+            node = CreateNode("YuE2GenerateMusic", new JObject()
+            {
+                ["clip"] = clip,
+                ["style"] = UserInput.Get(T2IParamTypes.Text2AudioStyle, ""),
+                ["lyrics"] = prompt,
+                ["abc"] = NodePath(abc, 0),
+                ["seed"] = UserInput.Get(T2IParamTypes.Seed, 0) + 20,
+                ["mode"] = "full",
+                ["max_duration"] = Math.Clamp(UserInput.Get(T2IParamTypes.Text2AudioDuration, 120), 0.04, 900),
+                ["temperature"] = 1,
+                ["top_p"] = 0.95,
+                ["top_k"] = 100,
+                ["repetition_penalty"] = 1.2
+            }, id);
+        }
         else if (IsMiniMaxMusic3())
         {
             node = CreateNode("MiniMaxMusic3TextEncode", new JObject()

src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs

@@ -278,6 +278,12 @@ public bool IsMiniMaxMusic3()
         return IsModelCompatClass(T2IModelClassSorter.CompatMiniMaxMusic3);
     }
 
+    /// <summary>Returns true if the current model is YuE2.</summary>
+    public bool IsYue2()
+    {
+        return IsModelCompatClass(T2IModelClassSorter.CompatYue2);
+    }
+
     /// <summary>Returns true if the current model primarily operates on audio.</summary>
     public bool IsAudioModel()
     {
@@ -419,9 +425,9 @@ public WGNodeData EmptyImage(int width, int height, int batchSize, string id = n
                 ["seconds"] = UserInput.Get(T2IParamTypes.Text2AudioDuration, 120)
             }, id));
         }
-        else if (IsMiniMaxMusic3())
+        else if (IsMiniMaxMusic3() || IsYue2())
         {
-            JProperty encodedNode = NodesOfClass("MiniMaxMusic3TextEncode").FirstOrDefault();
+            JProperty encodedNode = NodesOfClass(IsYue2() ? "YuE2GenerateMusic" : "MiniMaxMusic3TextEncode").FirstOrDefault();
             JToken targetSeconds;
             if (encodedNode is not null)
             {
@@ -431,7 +437,7 @@ public WGNodeData EmptyImage(int width, int height, int batchSize, string id = n
             {
                 targetSeconds = NodePath("6", 1); // TODO: This is a very wrong hack hardcoding the prompt path. This will break in many practical edge cases. Need to figure out the special routing that applies here.
             }
-            return resultAudio(CreateNode("EmptyMiniMaxMusic3LatentAudio", new JObject()
+            return resultAudio(CreateNode(IsYue2() ? "EmptyYuE2LatentAudio" : "EmptyMiniMaxMusic3LatentAudio", new JObject()
             {
                 ["batch_size"] = batchSize,
                 ["seconds"] = targetSeconds
@@ -1529,6 +1535,10 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC)
             helpers.DoVaeLoader(null, T2IModelClassSorter.CompatMiniMaxMusic3, "minimax-music-3-vae");
             CurrentAudioVae = new WGNodeData([LoadingVAE, 0], this, WGNodeData.DT_AUDIOVAE, CurrentCompat());
         }
+        else if (IsYue2())
+        {
+            CurrentAudioVae = new WGNodeData([LoadingVAE, 0], this, WGNodeData.DT_AUDIOVAE, CurrentCompat());
+        }
         else if (IsAceStep15())
         {
             // TODO: WTF? these twin qwen tencs are wacky.

src/Text2Image/T2IModelClassSorter.cs

@@ -102,6 +102,7 @@ public static T2IModelCompatClass
         // Audio models
         CompatAceStep15 = RegisterCompat(new() { ID = "ace-step-1_5", ShortCode = "Ace15", IsAudioModel = true }),
         CompatMiniMaxMusic3 = RegisterCompat(new() { ID = "minimax-music-3", ShortCode = "MMM3", IsAudioModel = true }),
+        CompatYue2 = RegisterCompat(new() { ID = "yue-2", ShortCode = "YuE2", IsAudioModel = true }),
         // Obscure old random ones
         CompatAuraFlow = RegisterCompat(new() { ID = "auraflow-v1", ShortCode = "Aura", VaeFamily = VaeSdxl }),
         CompatHiDreamI1 = RegisterCompat(new() { ID = "hidream-i1", ShortCode = "HiDrm", LorasTargetTextEnc = false, VaeFamily = VaeFlux1 }),
@@ -286,6 +287,7 @@ bool isAnimaLora(JObject h) => (hasLoraKey(h, "llm_adapter.blocks.5.self_attn.v_
         // Audio models
         bool isAceStep15(JObject h) => hasKey(h, "encoder.lyric_encoder.layers.0.post_attention_layernorm.weight");
         bool isMiniMaxMusic3(JObject h) => hasKey(h, "cond_layer_logits") && hasKey(h, "latent_conditioners.0.weight") && hasKey(h, "diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight");
+        bool isYue2(JObject h) => hasKey(h, "vae2llm.weight") && hasKey(h, "llm2vae.weight") && hasKey(h, "latent_pos_embed.pe") && hasKey(h, "model.layers.0.self_attn.qkv_proj.weight") && hasKey(h, "time_embedder.mlp.0.weight");
         // ====================== Stable Diffusion v1 ======================
         Register(new() { ID = "stable-diffusion-v1", CompatClass = CompatSdv1, Name = "Stable Diffusion v1", StandardWidth = 512, StandardHeight = 512, IsThisModelOfClass = (m, h) =>
         {
@@ -956,6 +958,10 @@ JToken GetEmbeddingKey(JObject h)
         {
             return isMiniMaxMusic3(h);
         }});
+        Register(new() { ID = "yue-2", CompatClass = CompatYue2, Name = "YuE2", IsThisModelOfClass = (m, h) =>
+        {
+            return isYue2(h);
+        }});
         // ====================== Everything below this point does not autodetect, it must match through ModelSpec or be manually set ======================
         // General Stable Diffusion variants
         Register(new() { ID = "stable-diffusion-v1/vae", CompatClass = CompatSdv1, Name = "Stable Diffusion v1 VAE", StandardWidth = 512, StandardHeight = 512, IsThisModelOfClass = (m, h) => { return false; } });

src/wwwroot/js/genpage/main.js

@@ -209,7 +209,7 @@ function reviseBackendFeatureSet() {
     doAnyCompatFeature(['stable-diffusion-v1', 'stable-diffusion-v2', 'stable-diffusion-xl-v1'], 'supports_reference_only');
     doAnyCompatFeature(['stable-diffusion-v1', 'stable-diffusion-v2', 'stable-diffusion-xl-v1'], 'supports_hypertile');
     doAnyCompatFeature(['genmo-mochi-1', 'lightricks-ltx-video', 'hunyuan-video', 'nvidia-cosmos-1', `wan-21`, `wan-22`, 'kandinsky5-vidlite', 'kandinsky5-vidpro', 'minimax-h3'], 'text2video');
-    doAnyCompatFeature(['ace-step-1_5', 'minimax-music-3'], 'text2audio');
+    doAnyCompatFeature(['ace-step-1_5', 'minimax-music-3', 'yue-2'], 'text2audio');
     for (let changer of featureSetChangers) {
         let [add, remove] = changer();
         addMe.push(...add);