AudioLab.cs

using Hartsy.Extensions.AudioLab.AudioAPI;
using Hartsy.Extensions.AudioLab.AudioBackends;
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 static new readonly string Version = "4.0.0";

    /// <summary>Pre-initialization — registers providers and web assets before SwarmUI core is ready.</summary>
    public override void OnPreInit()
    {
        try
        {
            // Set extension directory for Python path resolution
            string projectRoot = Path.GetFullPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "..", ".."));
            AudioConfiguration.ExtensionDirectory = Path.GetFullPath(Path.Combine(projectRoot, "Extensions", "SwarmUI-AudioLab"));
            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)}");

            // 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");
            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.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");
        }
        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
        {
            // Register T2I parameters for audio workflows (TTS, STT, Music, Clone, FX, SFX)
            AudioLabParams.RegisterAll();
            Logs.Info("[AudioLab] Registered audio T2I parameters");

            // Register feature flags so SwarmUI knows these are extension-managed
            RegisterFeatureFlags();
            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();
            Logs.Info("[AudioLab] Registered wake-word 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()
    {
        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
        {
            WakeWordService.Stop();
        }
        catch (Exception ex)
        {
            Logs.Error($"[AudioLab][Wake] Error during shutdown: {ex.ReadableString()}");
        }
    }

    /// <summary>Registers all feature flags that should be disregarded for audio backends.
    /// Mirrors the pattern from SwarmUI-API-Backends RegisterFeatureFlags().</summary>
    private static void RegisterFeatureFlags()
    {
        // Category-level flags (one per AudioCategory)
        string[] categoryFlags = ["audiolab_tts", "audiolab_stt", "audiolab_audiogen", "audiolab_clone", "audiolab_audioproc"];

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

        // Image-only features incompatible with audio models
        string[] incompatibleFlags = [
            "sampling", "zero_negative", "refiners", "controlnet", "variation_seed",
            "video", "autowebui", "comfyui", "frameinterps", "ipadapter", "sdxl",
            "dynamic_thresholding", "cascade", "sd3", "flux-dev", "seamless",
            "freeu", "teacache", "text2video", "yolov8", "aitemplate", "sdcpp"
        ];

        foreach (string flag in categoryFlags) T2IEngine.DisregardedFeatureFlags.Add(flag);
        foreach (string flag in providerFlags) T2IEngine.DisregardedFeatureFlags.Add(flag);
        foreach (string flag in incompatibleFlags) T2IEngine.DisregardedFeatureFlags.Add(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;
    }
}