From 8433349af5311afb000af08f3717ef354c2be407 Mon Sep 17 00:00:00 2001
From: fspecii <4722521+fspecii@users.noreply.github.com>
Date: Tue, 10 Feb 2026 13:24:36 +0200
Subject: [PATCH] Add AI Enhance toggle for better genre accuracy

When enabled, uses the LLM to enrich genre/style tags into detailed
music descriptions and generate proper BPM, key, and time signature
metadata (CoT features). Fixes genre tags like "pop, rock" producing
ballad-like output by matching Gradio UI default behavior.

- Add enhance toggle in Style of Music card header
- Gate CoT metas/caption/language by enhance OR thinking flag
- Remove unsupported --lm-backend/--lm-model from Python fallback
- Add i18n translations (en/zh/ja/ko) with tooltip
---
 components/CreatePanel.tsx     | 30 +++++++++++++++++++++---------
 i18n/translations.ts           | 12 ++++++++++++
 server/src/services/acestep.ts | 14 +++++++++-----
 types.ts                       |  1 +
 4 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/components/CreatePanel.tsx b/components/CreatePanel.tsx
index 79be27e..6ecd102 100644
--- a/components/CreatePanel.tsx
+++ b/components/CreatePanel.tsx
@@ -168,6 +168,7 @@ export const CreatePanel: React.FC<CreatePanelProps> = ({
   const [randomSeed, setRandomSeed] = useState(true);
   const [seed, setSeed] = useState(-1);
   const [thinking, setThinking] = useState(false); // Default false for GPU compatibility
+  const [enhance, setEnhance] = useState(false); // AI Enhance: uses LLM to enrich caption & generate metadata
   const [audioFormat, setAudioFormat] = useState<'mp3' | 'flac'>('mp3');
   const [inferenceSteps, setInferenceSteps] = useState(12);
   const [inferMethod, setInferMethod] = useState<'ode' | 'sde'>('ode');
@@ -1003,6 +1004,7 @@ export const CreatePanel: React.FC<CreatePanelProps> = ({
         randomSeed: randomSeed || i > 0, // Force random for subsequent bulk jobs
         seed: jobSeed,
         thinking,
+        enhance,
         audioFormat,
         inferMethod,
         lmBackend,
@@ -1601,7 +1603,17 @@ export const CreatePanel: React.FC<CreatePanelProps> = ({
             <div className="bg-white dark:bg-suno-card rounded-xl border border-zinc-200 dark:border-white/5 overflow-hidden transition-colors group focus-within:border-zinc-400 dark:focus-within:border-white/20">
               <div className="flex items-center justify-between px-3 py-2.5 bg-zinc-50 dark:bg-white/5 border-b border-zinc-100 dark:border-white/5">
                 <div>
-                  <span className="text-xs font-bold text-zinc-500 dark:text-zinc-400 uppercase tracking-wide">{t('styleOfMusic')}</span>
+                  <div className="flex items-center gap-2">
+                    <span className="text-xs font-bold text-zinc-500 dark:text-zinc-400 uppercase tracking-wide">{t('styleOfMusic')}</span>
+                    <button
+                      onClick={() => setEnhance(!enhance)}
+                      className={`flex items-center gap-1 px-1.5 py-0.5 rounded text-[10px] font-medium transition-all cursor-pointer ${enhance ? 'bg-violet-100 dark:bg-violet-500/20 text-violet-600 dark:text-violet-400' : 'text-zinc-400 dark:text-zinc-500 hover:text-zinc-600 dark:hover:text-zinc-300'}`}
+                      title={t('enhanceTooltip')}
+                    >
+                      <Sparkles size={9} />
+                      <span>{enhance ? 'ON' : 'OFF'}</span>
+                    </button>
+                  </div>
                   <p className="text-[11px] text-zinc-400 dark:text-zinc-500 mt-0.5">{t('genreMoodInstruments')}</p>
                 </div>
                 <div className="flex items-center gap-1">
@@ -1618,15 +1630,15 @@ export const CreatePanel: React.FC<CreatePanelProps> = ({
                   >
                     <Trash2 size={14} />
                   </button>
+                  <button
+                    className={`p-1.5 hover:bg-zinc-200 dark:hover:bg-white/10 rounded transition-colors ${isFormattingStyle ? 'text-pink-500' : 'text-zinc-500 hover:text-black dark:hover:text-white'}`}
+                    title="AI Format - Enhance style & auto-fill parameters"
+                    onClick={() => handleFormat('style')}
+                    disabled={isFormattingStyle || !style.trim()}
+                  >
+                    {isFormattingStyle ? <Loader2 size={14} className="animate-spin" /> : <Sparkles size={14} />}
+                  </button>
                 </div>
-                <button
-                  className={`p-1.5 hover:bg-zinc-200 dark:hover:bg-white/10 rounded transition-colors ${isFormattingStyle ? 'text-pink-500' : 'text-zinc-500 hover:text-black dark:hover:text-white'}`}
-                  title="AI Format - Enhance style & auto-fill parameters"
-                  onClick={() => handleFormat('style')}
-                  disabled={isFormattingStyle || !style.trim()}
-                >
-                  {isFormattingStyle ? <Loader2 size={14} className="animate-spin" /> : <Sparkles size={14} />}
-                </button>
               </div>
               <textarea
                 value={style}
diff --git a/i18n/translations.ts b/i18n/translations.ts
index 2e0e782..53c4a42 100644
--- a/i18n/translations.ts
+++ b/i18n/translations.ts
@@ -265,6 +265,9 @@ export const translations = {
     randomSeedRecommended: 'Randomized every run (recommended)',
     fixedSeedReproducible: 'Fixed seed for reproducible results',
     enterFixedSeed: 'Enter fixed seed',
+    enhance: 'AI Enhance',
+    enhanceHint: 'better genre accuracy, slightly slower',
+    enhanceTooltip: 'Uses the AI language model to enrich your genre/style tags into a detailed music description and generate proper BPM, key, and time signature. Improves genre accuracy but adds 10-20s to generation time. No extra VRAM needed.',
     thinkingCot: 'Thinking (CoT)',
     mp3Smaller: 'MP3 (smaller)',
     flacLossless: 'FLAC (lossless)',
@@ -861,6 +864,9 @@ export const translations = {
     randomSeedRecommended: '每次运行随机（推荐）',
     fixedSeedReproducible: '固定种子以获得可重现结果',
     enterFixedSeed: '输入固定种子',
+    enhance: 'AI 增强',
+    enhanceHint: '更准确的风格，稍慢',
+    enhanceTooltip: '使用AI语言模型将风格标签丰富为详细的音乐描述，并生成准确的BPM、调性和拍号。提高风格准确度，但增加10-20秒生成时间。无需额外显存。',
     thinkingCot: '思考模式（CoT）',
     mp3Smaller: 'MP3（较小）',
     flacLossless: 'FLAC（无损）',
@@ -1457,6 +1463,9 @@ export const translations = {
     randomSeedRecommended: '毎回ランダム化（推奨）',
     fixedSeedReproducible: '再現可能な結果のための固定シード',
     enterFixedSeed: '固定シードを入力',
+    enhance: 'AI エンハンス',
+    enhanceHint: 'ジャンル精度向上、やや遅い',
+    enhanceTooltip: 'AIがジャンル/スタイルタグを詳細な音楽説明に変換し、適切なBPM、キー、拍子を生成します。ジャンルの精度が向上しますが、生成に10-20秒追加されます。追加VRAMは不要です。',
     thinkingCot: '思考（CoT）',
     mp3Smaller: 'MP3（小さい）',
     flacLossless: 'FLAC（ロスレス）',
@@ -2053,6 +2062,9 @@ export const translations = {
     randomSeedRecommended: '매 실행마다 무작위 (권장)',
     fixedSeedReproducible: '재현 가능한 결과를 위한 고정 시드',
     enterFixedSeed: '고정 시드 입력',
+    enhance: 'AI 향상',
+    enhanceHint: '더 정확한 장르, 약간 느림',
+    enhanceTooltip: 'AI 언어 모델을 사용하여 장르/스타일 태그를 상세한 음악 설명으로 변환하고 적절한 BPM, 키, 박자를 생성합니다. 장르 정확도가 향상되지만 생성 시간이 10-20초 추가됩니다. 추가 VRAM이 필요하지 않습니다.',
     thinkingCot: '생각 (CoT)',
     mp3Smaller: 'MP3 (작음)',
     flacLossless: 'FLAC (무손실)',
diff --git a/server/src/services/acestep.ts b/server/src/services/acestep.ts
index cb76c54..2acd633 100644
--- a/server/src/services/acestep.ts
+++ b/server/src/services/acestep.ts
@@ -135,11 +135,15 @@ async function buildGradioArgs(params: GenerationParams): Promise<unknown[]> {
   const prompt = params.customMode ? caption : (params.songDescription || caption);
   const lyrics = params.instrumental ? '' : (params.lyrics || '');
   const isThinking = params.thinking ?? false;
+  const isEnhance = params.enhance ?? false;
 
   // Prepare audio files (async — reads from disk)
   const referenceAudio = await prepareAudioFile(params.referenceAudioUrl);
   const sourceAudio = await prepareAudioFile(params.sourceAudioUrl);
 
+  // CoT features are gated by enhance OR thinking (either enables LLM enrichment)
+  const useCot = isEnhance || isThinking;
+
   return [
     prompt,                                                       //  0: Music Caption
     lyrics,                                                       //  1: Lyrics
@@ -174,9 +178,9 @@ async function buildGradioArgs(params: GenerationParams): Promise<unknown[]> {
     params.lmTopK ?? 0,                                           // 30: LM Top-K
     params.lmTopP ?? 0.9,                                         // 31: LM Top-P
     params.lmNegativePrompt || 'NO USER INPUT',                   // 32: LM Negative Prompt
-    isThinking ? (params.useCotMetas ?? true) : false,            // 33: CoT Metas
-    isThinking ? (params.useCotCaption ?? true) : false,          // 34: CaptionRewrite
-    isThinking ? (params.useCotLanguage ?? true) : false,         // 35: CoT Language
+    useCot ? (params.useCotMetas ?? true) : false,                // 33: CoT Metas
+    useCot ? (params.useCotCaption ?? true) : false,              // 34: CaptionRewrite
+    useCot ? (params.useCotLanguage ?? true) : false,             // 35: CoT Language
     params.isFormatCaption ?? false,                              // 36: Is Format Caption State
     params.constrainedDecodingDebug ?? false,                     // 37: Constrained Decoding Debug
     params.allowLmBatch ?? true,                                  // 38: ParallelThinking
@@ -264,6 +268,7 @@ export interface GenerationParams {
   randomSeed?: boolean;
   seed?: number;
   thinking?: boolean;
+  enhance?: boolean;
   audioFormat?: 'mp3' | 'flac';
   inferMethod?: 'ode' | 'sde';
   shift?: number;
@@ -654,8 +659,7 @@ async function processGenerationViaPython(
     if (params.lmTopK !== undefined && params.lmTopK > 0) args.push('--lm-top-k', String(params.lmTopK));
     if (params.lmTopP !== undefined) args.push('--lm-top-p', String(params.lmTopP));
     if (params.lmNegativePrompt) args.push('--lm-negative-prompt', params.lmNegativePrompt);
-    if (params.lmBackend) args.push('--lm-backend', params.lmBackend);
-    if (params.lmModel) args.push('--lm-model', params.lmModel);
+    // Note: --lm-backend and --lm-model are not supported by simple_generate.py
     if (params.useCotMetas === false) args.push('--no-cot-metas');
     if (params.useCotCaption === false) args.push('--no-cot-caption');
     if (params.useCotLanguage === false) args.push('--no-cot-language');
diff --git a/types.ts b/types.ts
index 8f827ec..df9b58d 100644
--- a/types.ts
+++ b/types.ts
@@ -78,6 +78,7 @@ export interface GenerationParams {
   randomSeed: boolean;
   seed: number;
   thinking: boolean;
+  enhance?: boolean;
   audioFormat: 'mp3' | 'flac';
   inferMethod: 'ode' | 'sde';
   shift: number;