Add ACE-Step caption rewrite action

Author: 0xShug0

include/engine/models/ace_step/types.h

@@ -85,6 +85,7 @@ struct AceStepRequest {
     std::optional<float> repainting_end_seconds = std::nullopt;
     std::optional<runtime::AudioBuffer> source_audio = std::nullopt;
     std::optional<AceStepReferenceCondition> reference = std::nullopt;
+    bool rewrite_caption = false;
     AceStepGenerationOptions generation;
 };
 

src/models/ace_step/request_parser.cpp

@@ -126,6 +126,10 @@ AceStepRequest ace_step_parse_request(const runtime::TaskRequest &request) {
     if (const auto lyrics = runtime::find_option(request.options, {"lyrics"}); lyrics.has_value()) {
         out.lyrics = *lyrics;
     }
+    if (const auto rewrite_caption = runtime::find_option(request.options, {"rewrite_caption"});
+        rewrite_caption.has_value()) {
+        out.rewrite_caption = runtime::parse_bool_option(*rewrite_caption, "rewrite_caption");
+    }
     if (const auto negative = runtime::find_option(request.options, {"negative_prompt"}); negative.has_value()) {
         out.negative_prompt = *negative;
     }

src/models/ace_step/session.cpp

@@ -2,16 +2,19 @@
 
 #include "engine/framework/assets/tensor_source.h"
 #include "engine/framework/debug/profiler.h"
+#include "engine/framework/io/json.h"
 #include "engine/framework/runtime/options.h"
 #include "engine/models/ace_step/prompt_builder.h"
 #include "engine/models/ace_step/repaint.h"
 #include "engine/models/ace_step/request_parser.h"
 #include "engine/models/ace_step/task_route.h"
 
 #include <chrono>
+#include <sstream>
 #include <stdexcept>
 #include <string>
 #include <string_view>
+#include <unordered_map>
 #include <utility>
 
 namespace engine::models::ace_step {
@@ -143,6 +146,39 @@ bool mem_saver_from_options(const runtime::SessionOptions & options) {
     return false;
 }
 
+bool rewrite_caption_requested(const std::unordered_map<std::string, std::string> & options) {
+    const auto value = runtime::find_option(options, {"rewrite_caption"});
+    return value.has_value() && runtime::parse_bool_option(*value, "rewrite_caption");
+}
+
+std::string plan_json(const AceStepPlan & plan, const AceStepRequest & request) {
+    std::ostringstream out;
+    out << "{"
+        << "\"caption\":" << engine::io::json::stringify_string(plan.caption)
+        << ",\"cot_caption\":" << engine::io::json::stringify_string(plan.cot_caption)
+        << ",\"lyrics\":" << engine::io::json::stringify_string(request.lyrics);
+    if (plan.metadata.bpm.has_value()) {
+        out << ",\"bpm\":" << *plan.metadata.bpm;
+    }
+    if (plan.metadata.duration.has_value()) {
+        out << ",\"duration_seconds\":" << *plan.metadata.duration;
+    }
+    if (plan.metadata.keyscale.has_value()) {
+        out << ",\"keyscale\":" << engine::io::json::stringify_string(*plan.metadata.keyscale);
+    }
+    if (plan.metadata.timesignature.has_value()) {
+        out << ",\"timesignature\":" << engine::io::json::stringify_string(*plan.metadata.timesignature);
+    }
+    if (plan.metadata.language.has_value()) {
+        out << ",\"language\":" << engine::io::json::stringify_string(*plan.metadata.language);
+    }
+    if (plan.metadata.genres.has_value()) {
+        out << ",\"genres\":" << engine::io::json::stringify_string(*plan.metadata.genres);
+    }
+    out << "}";
+    return out.str();
+}
+
 }  // namespace
 
 AceStepSession::AceStepSession(runtime::TaskSpec task, runtime::SessionOptions options,
@@ -178,8 +214,11 @@ runtime::RunMode AceStepSession::run_mode() const {
 }
 
 void AceStepSession::prepare(const runtime::SessionPreparationRequest &request) {
-    (void)request;
     ensure_planner();
+    if (rewrite_caption_requested(request.options)) {
+        mark_prepared();
+        return;
+    }
     ensure_vae_decoder();
     ensure_pre_dit();
     pre_dit_->prepare_runtime();
@@ -196,9 +235,46 @@ runtime::TaskResult AceStepSession::run(const runtime::TaskRequest &request) {
     const AceStepTaskRoute &route = ace_step_task_route(ace_request);
     validate_task_route_request(ace_request, route);
     engine::debug::timing_log_scalar("ace_step.session.parse_request_ms", engine::debug::elapsed_ms(parse_start, Clock::now()));
+    const bool flow_edit_morph = ace_step_request_uses_flow_edit_morph(ace_request);
+    if (ace_request.rewrite_caption) {
+        if (flow_edit_morph || !route.uses_planner) {
+            throw std::runtime_error("ACE-Step rewrite_caption requires a planner route");
+        }
+        const auto planner_ensure_start = Clock::now();
+        ensure_planner();
+        engine::debug::timing_log_scalar("ace_step.session.ensure_planner_ms",
+                                         engine::debug::elapsed_ms(planner_ensure_start, Clock::now()));
+        const auto planner_start = Clock::now();
+        AceStepPlan plan = planner_->generate(ace_request, false);
+        engine::debug::timing_log_scalar("ace_step.session.planner_generate_ms",
+                                         engine::debug::elapsed_ms(planner_start, Clock::now()));
+        const auto planner_release_start = Clock::now();
+        planner_->release_graph_workspace();
+        if (execution_context().backend_type() == core::BackendType::Metal) {
+            planner_.reset();
+        }
+        engine::debug::timing_log_scalar("ace_step.session.planner_release.graph.workspace_ms",
+                                         engine::debug::elapsed_ms(planner_release_start, Clock::now()));
+
+        runtime::TaskResult result;
+        result.text_output = runtime::Transcript{
+            plan.caption,
+            plan.metadata.language.value_or(ace_request.vocal_language),
+        };
+        result.output_artifacts.push_back(runtime::make_text_artifact(
+            runtime::ArtifactKind::Custom,
+            "ace_step_caption_plan",
+            plan_json(plan, ace_request),
+            {
+                {"format", "json"},
+                {"extension", "json"},
+                {"mime", "application/json"},
+            }));
+        engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(total_start, Clock::now()));
+        return result;
+    }
 
     AceStepPlan plan;
-    const bool flow_edit_morph = ace_step_request_uses_flow_edit_morph(ace_request);
     const bool has_request_audio_codes = !flow_edit_morph && !ace_request.audio_code_ids.empty();
     const bool use_planner =
         !flow_edit_morph &&

webui/native/dist/index.html

@@ -161,6 +161,9 @@ const english: Record<string, string> = {
   'request.randomSeed': '-1 = random',
   'request.maxTokens': 'Maximum tokens',
   'request.duration': 'Duration seconds',
+  'request.autoDuration': '-1 = auto',
+  'request.rewriteCaption': 'Rewrite caption',
+  'request.rewritingCaption': 'Rewriting caption...',
   'request.minimaxFrames': '{frames} aligned output frames',
   'request.sourceAudio': 'Source audio',
   'request.stopRecording': 'Stop recording',

webui/native/src/lib/i18n.ts

@@ -61,6 +61,7 @@
   let installed: boolean | null = null;
   let loadingModel = false;
   let running = false;
+  let rewritingCaption = false;
   let status = 'Ready';
   let warningStatus = '';
   let errorStatus = '';
@@ -230,7 +231,8 @@
   }
 
   function setDuration(value: number) {
-    duration = Math.max(1, Number.isFinite(value) ? value : 1);
+    const minimum = selected?.family === 'ace_step' ? -1 : 1;
+    duration = Math.max(minimum, Number.isFinite(value) ? value : minimum);
     if (selected?.family === 'minimax_h3') {
       advancedValues = { ...advancedValues, num_frames: miniMaxFramesForDuration(duration) };
     }
@@ -386,6 +388,7 @@
     modelMatchesSelectedPackage(model, selected));
   $: isFireRedAudioEdit = selected?.id === 'firered-audio-semantic-edit' ||
     selected?.id === 'firered-audio-acoustic-edit';
+  $: allowsAutoDuration = selected?.family === 'ace_step';
   $: usesDurationSecOption =
     selected?.family === 'controlfoley' ||
     selected?.family === 'midashenglm_gen';
@@ -1220,6 +1223,74 @@
     return { ...defaults, ...advancedValues, ...raw };
   }
 
+  function base64Text(value: string): string {
+    const binary = atob(value);
+    const bytes = new Uint8Array(binary.length);
+    for (let index = 0; index < binary.length; index += 1) {
+      bytes[index] = binary.charCodeAt(index);
+    }
+    return new TextDecoder().decode(bytes);
+  }
+
+  function acePlanFromResult(result: Record<string, unknown>): Record<string, unknown> {
+    if (Array.isArray(result.artifacts)) {
+      const artifact = result.artifacts.find((entry): entry is { id: string; payload: string } =>
+        typeof entry === 'object' && entry !== null &&
+          (entry as { id?: unknown }).id === 'ace_step_caption_plan' &&
+          typeof (entry as { payload?: unknown }).payload === 'string');
+      if (artifact) return JSON.parse(base64Text(artifact.payload));
+    }
+    if (typeof result.text === 'string') return { caption: result.text };
+    return {};
+  }
+
+  async function rewriteAceCaption() {
+    if (selected?.family !== 'ace_step' || running || rewritingCaption) return;
+    if (!text.trim() && !lyrics.trim()) {
+      status = 'Enter a caption or lyrics to rewrite.';
+      warningStatus = status;
+      errorStatus = '';
+      return;
+    }
+    rewritingCaption = true;
+    warningStatus = '';
+    errorStatus = '';
+    status = tr('request.rewritingCaption');
+    try {
+      await ensureLoaded();
+      const options = { ...requestOptions(), rewrite_caption: true };
+      const request: Record<string, unknown> = {
+        text,
+        seed: resolveRequestSeed(seed),
+        duration_seconds: duration,
+        options
+      };
+      if (language.trim()) request.language = language;
+      if (lyrics.trim()) request.lyrics = lyrics;
+      const result = await runTask({ model: selected.id, request });
+      const plan = acePlanFromResult(result);
+      if (typeof plan.caption === 'string' && plan.caption.trim()) text = plan.caption;
+      if (typeof plan.language === 'string' && plan.language.trim()) language = plan.language;
+      if (typeof plan.duration_seconds === 'number' && Number.isFinite(plan.duration_seconds) && plan.duration_seconds > 0) {
+        duration = plan.duration_seconds;
+      }
+      const nextAdvanced = { ...advancedValues };
+      if (typeof plan.bpm === 'number' && Number.isFinite(plan.bpm)) nextAdvanced.bpm = plan.bpm;
+      if (typeof plan.keyscale === 'string') nextAdvanced.keyscale = plan.keyscale;
+      if (typeof plan.timesignature === 'string') nextAdvanced.timesignature = plan.timesignature;
+      advancedValues = nextAdvanced;
+      outputText = typeof result.text === 'string' ? result.text : '';
+      outputJson = JSON.stringify(plan, null, 2);
+      status = 'Caption rewritten.';
+    } catch (error) {
+      status = error instanceof Error ? error.message : String(error);
+      errorStatus = status;
+      log(`Caption rewrite failed: ${status}`);
+    } finally {
+      rewritingCaption = false;
+    }
+  }
+
   function clearOutput() {
     for (const output of outputAudio) URL.revokeObjectURL(output.url);
     outputAudio = [];
@@ -2112,6 +2183,14 @@
           <label for="lyrics">{tr('request.lyrics')} <span>{lyricsRequired ? tr('voice.required') : tr('request.optional')}</span></label>
           <textarea id="lyrics" rows="3" bind:value={lyrics} required={lyricsRequired}
             aria-required={lyricsRequired} placeholder="[Verse]…"></textarea>
+          {#if selected.family === 'ace_step'}
+            <div class="media-actions">
+              <button type="button" disabled={running || rewritingCaption || (!text.trim() && !lyrics.trim())}
+                on:click={rewriteAceCaption}>
+                {rewritingCaption ? tr('request.rewritingCaption') : tr('request.rewriteCaption')}
+              </button>
+            </div>
+          {/if}
         {/if}
 
         {#if selected.task === 'asr'}
@@ -2147,8 +2226,11 @@
           {#if selected.task === 'gen'}
             <div>
               <label for="duration">{tr('request.duration')}</label>
-              <input id="duration" type="number" min="1" step="0.1" value={duration}
+              <input id="duration" type="number" min={allowsAutoDuration ? -1 : 1} step="0.1" value={duration}
                 on:input={(event) => setDuration(event.currentTarget.valueAsNumber)} />
+              {#if allowsAutoDuration}
+                <small>{tr('request.autoDuration')}</small>
+              {/if}
               {#if selected.family === 'minimax_h3'}
                 <small>{tr('request.minimaxFrames', { frames: Number(advancedValues.num_frames || 0) })}</small>
               {/if}