AudioLab.cs

using Hartsy.Extensions.AudioLab.AudioAPI;
using Hartsy.Extensions.AudioLab.AudioBackends;
using Hartsy.Extensions.AudioLab.AudioModels;
using Hartsy.Extensions.AudioLab.AudioProviders;
using Hartsy.Extensions.AudioLab.AudioProviderTypes;
using Hartsy.Extensions.AudioLab.AudioServices;
using Newtonsoft.Json.Linq;
using SwarmUI.Core;
using SwarmUI.Text2Image;
using SwarmUI.Utils;
using System.IO;

namespace Hartsy.Extensions.AudioLab;

/// <summary>SwarmUI AudioLab Extension - Main Entry Point.
/// Provides modular audio processing (TTS, STT, music gen, voice cloning, etc.)
/// through a provider-based architecture integrated into SwarmUI's Generate tab.</summary>
public class AudioLab : Extension
{
    /// <summary>Current extension version.</summary>
    public const string ExtensionVersion = "4.0.0";

    /// <summary>Fills in the fields core shows on the Extensions tab. Called once, after OnFirstInit.</summary>
    public override void PopulateMetadata()
    {
        Version = ExtensionVersion;
        ExtensionAuthor = "Hartsy AI";
        Description = "Audio generation, transcription, voice conversion and a multi-track DAW, running on the "
            + "in-process HartsyInference engine.";
        License = "MIT";
        ReadmeURL = "https://github.com/HartsyAI/SwarmUI-AudioLab";
        Tags = ["audio", "tts", "stt", "music"];
    }

    /// <summary>Pre-initialization — registers providers and web assets before SwarmUI core is ready.</summary>
    public override void OnPreInit()
    {
        try
        {
            // Core sets FilePath while loading the extension, before this hook. Deriving it from the assembly
            // location instead used to hardcode the checkout folder name, so renaming the directory silently
            // broke every provider preview image.
            AudioConfiguration.ExtensionDirectory = Path.GetFullPath(FilePath);
            Logs.Info($"[AudioLab] Extension directory: {AudioConfiguration.ExtensionDirectory}");

            // Settings load well before extension pre-init, so the server's model root is known here.
            AudioConfiguration.SyncModelRootFromServer();
            Logs.Info($"[AudioLab] Audio model root: {Path.GetFullPath(AudioConfiguration.ModelRoot)}");

            // The engine's model cache root is fixed on first read; align it before any weight-presence check can
            // touch it, so downloads land under the configured Swarm model root instead of ~/.cache.
            AudioEngineBridge.AlignModelsRoot();

            // Register all built-in audio providers
            AudioProviderDefinitions.RegisterAll();
            Logs.Info($"[AudioLab] Registered {AudioProviderDefinitions.All.Count} audio providers");

            // Must run before the user-settings page first renders — its key-entry table is generated
            // from UserUpstreamApiKeys.KeysByType, so unregistered types have no input field at all.
            AudioApiKeys.RegisterAll();

            // Register web assets — libraries first, then DAW modules, then integration
            ScriptFiles.Add("Assets/lib/wavesurfer.min.js");
            ScriptFiles.Add("Assets/lib/wavesurfer-record.min.js");
            ScriptFiles.Add("Assets/lib/wavesurfer-regions.min.js");
            ScriptFiles.Add("Assets/lib/wavesurfer-timeline.min.js");
            ScriptFiles.Add("Assets/lib/wavesurfer-minimap.min.js");
            ScriptFiles.Add("Assets/lib/crunker.min.js");
            // abcjs (MIT) engraves the ABC scores YuE2 plans. Vendored rather than CDN-loaded so the Score tab
            // works on an offline install, like every other lib here.
            ScriptFiles.Add("Assets/lib/abcjs-basic-min.min.js");
            ScriptFiles.Add("Assets/audio-player.js");
            ScriptFiles.Add("Assets/audio-api.js");
            ScriptFiles.Add("Assets/audio-core.js");
            ScriptFiles.Add("Assets/audio-daw-timeline.js");
            ScriptFiles.Add("Assets/audio-daw-track.js");
            ScriptFiles.Add("Assets/audio-daw-mixer.js");
            ScriptFiles.Add("Assets/audio-daw-fx.js");
            ScriptFiles.Add("Assets/audio-daw-store.js");
            ScriptFiles.Add("Assets/audio-daw-input.js");
            ScriptFiles.Add("Assets/audio-daw-instruments.js");
            ScriptFiles.Add("Assets/audio-daw-score.js");
            ScriptFiles.Add("Assets/audio-voice.js");
            ScriptFiles.Add("Assets/audio-daw.js");
            ScriptFiles.Add("Assets/audio-editor.js");
            ScriptFiles.Add("Assets/audio-integration.js");
            ScriptFiles.Add("Assets/audio-wakeword.js");
            StyleSheetFiles.Add("Assets/audio-lab.css");
            // Loaded via audioContext.audioWorklet.addModule(...), not as a page <script> tag -- OtherAssets
            // just needs the file servable, the same treatment the soundfont notes below get.
            OtherAssets.Add("Assets/audio-voice-worklet.js");
            OtherAssets.Add("Assets/audio-daw-recorder-worklet.js");
            // Registered whether or not the files exist yet: the getters read on request, so a soundfont
            // fetched later in the session is served without a restart.
            foreach (string note in ScoreSoundfont.AssetPaths)
            {
                OtherAssets.Add(note);
            }
        }
        catch (Exception ex)
        {
            Logs.Error($"[AudioLab] Critical error during pre-initialization: {ex.Message}");
        }
    }

    /// <summary>Main initialization — registers backend, T2I params, feature flags, and API endpoints.</summary>
    // Not async: nothing here awaits, and `async void` turned any startup failure into an unobservable
    // crash on the synchronization context instead of a logged error.
    public override void OnInit()
    {
        try
        {
            // Before core's first RefreshAllModelSets, so Models/audio is in the very first scan.
            AudioModelTypeRegistration.Register();

            // Before anything reads a sidecar, so artifacts stamped under our old class ids still classify.
            AudioModelFactory.RegisterLegacyClassRemaps();

            // Register T2I parameters for audio workflows (TTS, STT, Music, Clone, FX, SFX)
            HashSet<string> foreignParamIds = [.. T2IParamTypes.Types.Keys];
            AudioLabParams.RegisterAll();
            T2IParamType[] ownParams = [.. T2IParamTypes.Types.Values.Where(t => !foreignParamIds.Contains(t.ID))];
            Logs.Info("[AudioLab] Registered audio T2I parameters");

            // Register feature flags so SwarmUI knows these are extension-managed
            RegisterFeatureFlags(ownParams);
            Logs.Info("[AudioLab] Registered feature flags");

            // Register ONE unified backend
            Program.Backends.RegisterBackendType<DynamicAudioBackend>(
                "audio-backend", "Audio Backend",
                "Dynamic audio backend supporting TTS, STT, music generation, and more.", true);

            // Register API endpoints
            AudioLabAPI.Register();
            VideoAudioEndpoints.Register();
            WakeWordEndpoints.Register();
            VoiceSessionEndpoints.Register();
            Logs.Info("[AudioLab] Registered wake-word and voice-agent endpoints");
        }
        catch (Exception ex)
        {
            Logs.Error($"[AudioLab] Critical error during initialization: {ex.Message}");
        }
    }

    /// <summary>Starts the wake-word listener if it is enabled in settings.
    ///
    /// <para>Deliberately in OnPreLaunch rather than OnInit: the listener accepts network connections and pushes
    /// events to the web UI, so it should not be live before the webserver is. It is off unless enabled, so an
    /// install with no voice satellite never binds a port or holds a detection thread.</para></summary>
    public override void OnPreLaunch()
    {
        // Core has finished mapping its own routes and has not started listening yet, which is the one window
        // an extension can add a raw HTTP endpoint in. Registered unconditionally: it costs nothing until
        // something posts to it, and a device that wants streamed speech has no other way to ask.
        try
        {
            SpeakStreamRoute.Register();
            VoiceTurnOrchestrator.Start();
        }
        catch (Exception ex)
        {
            Logs.Error($"[AudioLab] Could not register the streaming speech route: {ex.ReadableString()}");
        }
        try
        {
            if (!WakeWordService.GetSettings().Enabled)
            {
                Logs.Debug("[AudioLab][Wake] Listener is disabled in settings; not starting.");
                return;
            }
            string error = WakeWordService.Start();
            if (error is not null)
            {
                Logs.Error($"[AudioLab][Wake] Listener failed to start: {error}");
            }
        }
        catch (Exception ex)
        {
            // A failed listener must not take the rest of the extension down with it.
            Logs.Error($"[AudioLab][Wake] Unexpected error starting the listener: {ex.ReadableString()}");
        }
    }

    /// <summary>Releases the wake listener's port and joins its worker.
    ///
    /// <para>Core has already cancelled GlobalProgramCancel several steps before extensions are shut down, so
    /// this joins work that is already unwinding rather than initiating the stop.</para></summary>
    public override void OnShutdown()
    {
        try
        {
            AudioModelTypeRegistration.Unregister();
            VoiceTurnOrchestrator.Stop();
            WakeWordService.Stop();
        }
        catch (Exception ex)
        {
            Logs.Error($"[AudioLab][Wake] Error during shutdown: {ex.ReadableString()}");
        }
    }

    /// <summary>Marks AudioLab's own feature flags as UI-visibility-only, so they never gate backend selection.
    ///
    /// <para>Only flags this extension owns belong here. <see cref="T2IEngine.DisregardedFeatureFlags"/> is
    /// process-global and core consults it to decide whether a backend may be filtered out for lacking a feature
    /// a request needs — so adding a CORE flag here disables that check for every generation in the process,
    /// image ones included. This used to also add the image-only list ("controlnet", "refiners", "freeu",
    /// "ipadapter", ...), which meant a ControlNet request would happily route to a backend that cannot do
    /// ControlNet. Hiding image params on an audio model is the JS layer's job (see audio-integration.js), not
    /// this set's.</para></summary>
    private static void RegisterFeatureFlags(IEnumerable<T2IParamType> ownParams)
    {
        // Category-level flags (one per AudioCategory), plus the output-format flag the backend advertises.
        string[] categoryFlags = [.. DynamicAudioBackend.CategoryFlags.Values, DynamicAudioBackend.OutputFlag, DynamicAudioBackend.DurationFlag];

        // Per-provider flags from each provider's FeatureFlags list
        string[] providerFlags = AudioProviderRegistry.All
            .SelectMany(p => p.FeatureFlags).Distinct().ToArray();

        foreach (string flag in categoryFlags) T2IEngine.DisregardedFeatureFlags.Add(flag);
        foreach (string flag in providerFlags) T2IEngine.DisregardedFeatureFlags.Add(flag);
        WarnOnUndeclaredFeatureFlags(ownParams, [.. categoryFlags, .. providerFlags]);
    }

    /// <summary>Complains at startup about any param of ours carrying a flag nothing grants.
    ///
    /// <para>The failure it catches is silent by construction: <see cref="T2IEngine"/> drops a backend whose
    /// features don't cover a job's required flags and names neither the param nor the flag, so a param with a
    /// typo'd or never-registered flag just makes every generation touching it refuse. Flags in
    /// <c>DisregardedFeatureFlags</c> never gate a backend, so they are fine; core's own flags are core's to
    /// grant. Only params this extension registered are checked: other extensions use the same "_params" naming
    /// (API-Backends' "openai_sora_params" is granted by its own backends) and guessing ownership from the name
    /// reported their flags as broken. Mirrors SwarmUIHartsyInference.WarnOnUndeclaredFeatureFlags.</para></summary>
    private static void WarnOnUndeclaredFeatureFlags(IEnumerable<T2IParamType> ownParams, HashSet<string> declared)
    {
        foreach (T2IParamType type in ownParams)
        {
            if (string.IsNullOrEmpty(type.FeatureFlag))
            {
                continue;
            }
            foreach (string flag in type.FeatureFlag.Split(','))
            {
                bool isOurs = flag.StartsWith("audiolab_", StringComparison.Ordinal) || flag.EndsWith("_params", StringComparison.Ordinal);
                if (isOurs && !declared.Contains(flag) && !T2IEngine.DisregardedFeatureFlags.Contains(flag))
                {
                    Logs.Error($"[AudioLab] Param '{type.Name}' requires feature flag '{flag}', which nothing "
                        + "registers; every generation using that param will be refused with no explanation. "
                        + "Add it to a provider's FeatureFlags, or drop the flag.");
                }
            }
        }
    }

    /// <summary>Creates a standardized error response for API endpoints.</summary>
    public static JObject CreateErrorResponse(string message, string errorCode = null, Exception exception = null)
    {
        JObject response = new()
        {
            ["success"] = false,
            ["error"] = message,
            ["timestamp"] = DateTime.UtcNow.ToString("O")
        };
        if (!string.IsNullOrEmpty(errorCode))
        {
            response["error_code"] = errorCode;
        }
        if (exception != null)
        {
            Logs.Error($"[AudioLab] Exception details: {exception}");
            response["error_type"] = exception.GetType().Name;
        }
        return response;
    }

    /// <summary>Creates a standardized success response for API endpoints.</summary>
    public static JObject CreateSuccessResponse(object data = null, string message = null)
    {
        JObject response = new()
        {
            ["success"] = true,
            ["timestamp"] = DateTime.UtcNow.ToString("O")
        };
        if (!string.IsNullOrEmpty(message))
        {
            response["message"] = message;
        }
        if (data != null)
        {
            if (data is JObject jObject)
            {
                foreach (JProperty property in jObject.Properties())
                {
                    response[property.Name] = property.Value;
                }
            }
            else
            {
                response["data"] = JToken.FromObject(data);
            }
        }
        return response;
    }
}