Aceradio/AceStepWorker.cpp

#include "AceStepWorker.h"
#include <QFile>
#include <QJsonDocument>
#include <QJsonObject>
#include <QProcess>
#include <QDir>
#include <QStandardPaths>
#include <QDebug>
#include <QCoreApplication>

AceStepWorker::AceStepWorker(QObject *parent)
    : QObject(parent),
      currentWorker(nullptr)
{
}

AceStepWorker::~AceStepWorker()
{
    cancelGeneration();
}

void AceStepWorker::generateSong(const QString &caption, const QString &lyrics, const QString &jsonTemplate,
                                const QString &aceStepPath, const QString &qwen3ModelPath,
                                const QString &textEncoderModelPath, const QString &ditModelPath,
                                const QString &vaeModelPath)
{
    // Cancel any ongoing generation
    cancelGeneration();
    
    // Create worker and start it
    currentWorker = new Worker(this, caption, lyrics, jsonTemplate, aceStepPath, qwen3ModelPath,
                               textEncoderModelPath, ditModelPath, vaeModelPath);
    QThreadPool::globalInstance()->start(currentWorker);
}

void AceStepWorker::cancelGeneration()
{
    // Note: In a real implementation, we would need to implement proper cancellation
    // For now, we just clear the reference
    currentWorker = nullptr;
}

void AceStepWorker::workerFinished()
{
    emit generationFinished();
}

// Worker implementation
void AceStepWorker::Worker::run()
{
    
    // Create temporary JSON file for the request
    QString tempDir = QStandardPaths::writableLocation(QStandardPaths::TempLocation);
    QString requestFile = tempDir + "/request_" + QString::number(QCoreApplication::applicationPid()) + ".json";
    
    // Parse and modify the template
    QJsonParseError parseError;
    QJsonDocument templateDoc = QJsonDocument::fromJson(jsonTemplate.toUtf8(), &parseError);
    if (!templateDoc.isObject()) {
        emit parent->generationError("Invalid JSON template: " + QString(parseError.errorString()));
        return;
    }
    
    QJsonObject requestObj = templateDoc.object();
    requestObj["caption"] = caption;
    
    if (!lyrics.isEmpty()) {
        requestObj["lyrics"] = lyrics;
    } else {
        // Remove lyrics field if empty to let the LLM generate them
        requestObj.remove("lyrics");
    }
    
    // Write the request file
    QFile requestFileHandle(requestFile);
    if (!requestFileHandle.open(QIODevice::WriteOnly | QIODevice::Text)) {
        emit parent->generationError("Failed to create request file: " + requestFileHandle.errorString());
        return;
    }
    
    requestFileHandle.write(QJsonDocument(requestObj).toJson(QJsonDocument::Indented));
    requestFileHandle.close();
    
    // Use provided paths for acestep.cpp binaries
    QString qwen3Binary = this->aceStepPath + "/ace-qwen3";
    QString ditVaeBinary = this->aceStepPath + "/dit-vae";
    
    // Check if binaries exist
    QFileInfo qwen3Info(qwen3Binary);
    QFileInfo ditVaeInfo(ditVaeBinary);
    
    if (!qwen3Info.exists() || !qwen3Info.isExecutable()) {
        emit parent->generationError("ace-qwen3 binary not found at: " + qwen3Binary);
        return;
    }
    
    if (!ditVaeInfo.exists() || !ditVaeInfo.isExecutable()) {
        emit parent->generationError("dit-vae binary not found at: " + ditVaeBinary);
        return;
    }
    
    // Use provided model paths
    QString qwen3Model = this->qwen3ModelPath;
    QString textEncoderModel = this->textEncoderModelPath;
    QString ditModel = this->ditModelPath;
    QString vaeModel = this->vaeModelPath;
    
    if (!QFileInfo::exists(qwen3Model)) {
        emit parent->generationError("Qwen3 model not found: " + qwen3Model);
        return;
    }
    
    if (!QFileInfo::exists(textEncoderModel)) {
        emit parent->generationError("Text encoder model not found: " + textEncoderModel);
        return;
    }
    
    if (!QFileInfo::exists(ditModel)) {
        emit parent->generationError("DiT model not found: " + ditModel);
        return;
    }
    
    if (!QFileInfo::exists(vaeModel)) {
        emit parent->generationError("VAE model not found: " + vaeModel);
        return;
    }
    
    // Step 1: Run ace-qwen3 to generate lyrics and audio codes
    QProcess qwen3Process;
    QStringList qwen3Args;
    qwen3Args << "--request" << requestFile;
    qwen3Args << "--model" << qwen3Model;
    
    emit parent->progressUpdate(20);
    
    qwen3Process.start(qwen3Binary, qwen3Args);
    if (!qwen3Process.waitForStarted()) {
        emit parent->generationError("Failed to start ace-qwen3: " + qwen3Process.errorString());
        return;
    }
    
    if (!qwen3Process.waitForFinished(30000)) { // 30 second timeout
        qwen3Process.terminate();
        qwen3Process.waitForFinished(5000);
        emit parent->generationError("ace-qwen3 timed out after 30 seconds");
        return;
    }
    
    int exitCode = qwen3Process.exitCode();
    if (exitCode != 0) {
        QString errorOutput = qwen3Process.readAllStandardError();
        emit parent->generationError("ace-qwen3 exited with code " + QString::number(exitCode) + ": " + errorOutput);
        return;
    }
    
    emit parent->progressUpdate(50);
    
    // Step 2: Run dit-vae to generate audio
    QProcess ditVaeProcess;
    QStringList ditVaeArgs;
    ditVaeArgs << "--request" << requestFile;
    ditVaeArgs << "--text-encoder" << textEncoderModel;
    ditVaeArgs << "--dit" << ditModel;
    ditVaeArgs << "--vae" << vaeModel;
    
    emit parent->progressUpdate(60);
    
    ditVaeProcess.start(ditVaeBinary, ditVaeArgs);
    if (!ditVaeProcess.waitForStarted()) {
        emit parent->generationError("Failed to start dit-vae: " + ditVaeProcess.errorString());
        return;
    }
    
    if (!ditVaeProcess.waitForFinished(120000)) { // 2 minute timeout
        ditVaeProcess.terminate();
        ditVaeProcess.waitForFinished(5000);
        emit parent->generationError("dit-vae timed out after 2 minutes");
        return;
    }
    
    exitCode = ditVaeProcess.exitCode();
    if (exitCode != 0) {
        QString errorOutput = ditVaeProcess.readAllStandardError();
        emit parent->generationError("dit-vae exited with code " + QString::number(exitCode) + ": " + errorOutput);
        return;
    }
    
    emit parent->progressUpdate(90);
    
    // Find the generated WAV file
    QDir requestDir(QFileInfo(requestFile).absolutePath());
    QStringList wavFiles = requestDir.entryList(QStringList("request*.wav"), QDir::Files, QDir::Name);
    
    if (wavFiles.isEmpty()) {
        emit parent->generationError("No WAV file generated");
        return;
    }
    
    QString wavFile = requestDir.absoluteFilePath(wavFiles.first());
    
    // Clean up temporary files
    QFile::remove(requestFile);
    
    emit parent->progressUpdate(100);
    emit parent->songGenerated(wavFile);
}