h-school/src/HSchool.Server/Game/SwarmUiClient.cs

using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Options;

namespace HSchool.Server.Game;

internal sealed class SwarmUiClient
{
    private static readonly JsonSerializerOptions Json = new()
    {
        PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
        PropertyNameCaseInsensitive = true,
    };

    private readonly HttpClient _http;
    private readonly SwarmUiOptions _options;
    private readonly ILogger<SwarmUiClient> _logger;
    private string? _sessionId;

    public SwarmUiClient(HttpClient http, IOptions<SwarmUiOptions> options, ILogger<SwarmUiClient> logger)
    {
        _http = http;
        _options = options.Value;
        _logger = logger;
    }

    public bool IsConfigured => !string.IsNullOrWhiteSpace(_options.BaseUrl);

    /// <summary>Opens a session to verify SwarmUI responds; does not generate an image.</summary>
    public async Task<bool> ProbeAsync(TimeSpan timeout, CancellationToken cancellationToken)
    {
        if (!IsConfigured)
        {
            return false;
        }

        using var timeoutSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeoutSource.CancelAfter(timeout);

        try
        {
            await RefreshSessionAsync(timeoutSource.Token);
            return true;
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            return false;
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "SwarmUI probe failed.");
            return false;
        }
    }

    public async Task<SwarmUiDiscovery> FetchDiscoveryAsync(CancellationToken cancellationToken)
    {
        if (!IsConfigured)
        {
            return SwarmUiDiscovery.Offline;
        }

        try
        {
            return await RunWithSessionAsync(async () =>
            {
                using var content = new StringContent(
                    JsonSerializer.Serialize(new { session_id = _sessionId }),
                    Encoding.UTF8,
                    "application/json");
                using var response = await _http.PostAsync("/API/ListT2IParams", content, cancellationToken);
                response.EnsureSuccessStatusCode();
                var json = await response.Content.ReadAsStringAsync(cancellationToken);
                using var document = JsonDocument.Parse(json);
                return SwarmUiDiscoveryParser.Parse(document.RootElement);
            }, cancellationToken);
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "SwarmUI discovery failed.");
            return SwarmUiDiscovery.Offline;
        }
    }

    public async Task<byte[]> GenerateAsync(
        string prompt,
        string negativePrompt,
        SwarmUiResolvedProfile profile,
        CancellationToken cancellationToken)
    {
        if (!IsConfigured)
        {
            throw new InvalidOperationException("SwarmUI is not configured.");
        }

        return await RunWithSessionAsync(async () =>
        {
            var kind = profile.KindPreset;
            var body = new Dictionary<string, object?>
            {
                ["session_id"] = _sessionId,
                ["images"] = 1,
                ["donotsave"] = true,
                ["prompt"] = prompt,
                ["negativeprompt"] = negativePrompt,
                ["model"] = profile.Model,
                ["steps"] = profile.Steps,
                ["cfgscale"] = profile.CfgScale,
                ["width"] = kind.Width,
                ["height"] = kind.Height,
                ["seed"] = profile.Seed,
            };

            if (!string.IsNullOrWhiteSpace(profile.Sampler))
            {
                body["sampler"] = profile.Sampler;
            }

            if (!string.IsNullOrWhiteSpace(profile.Scheduler))
            {
                body["scheduler"] = profile.Scheduler;
            }

            if (profile.ClipSkip > 0)
            {
                body["clipstopatlayer"] = -profile.ClipSkip;
            }

            var (loraNames, loraWeights) = SwarmUiLoraFormatter.FormatForApi(
                SwarmUiLoraFormatter.Concat(profile.PositiveLoras, profile.NegativeLoras));
            if (loraNames is not null)
            {
                body["loras"] = loraNames;
                body["loraweights"] = loraWeights;
            }

            using var content = new StringContent(JsonSerializer.Serialize(body), Encoding.UTF8, "application/json");
            using var response = await _http.PostAsync("/API/GenerateText2Image", content, cancellationToken);
            response.EnsureSuccessStatusCode();

            var json = await response.Content.ReadAsStringAsync(cancellationToken);
            using var document = JsonDocument.Parse(json);
            var root = document.RootElement;

            if (root.TryGetProperty("error_id", out var errorId)
                && errorId.GetString() == "invalid_session_id")
            {
                throw new SwarmUiSessionInvalidException();
            }

            if (root.TryGetProperty("error", out var error))
            {
                throw new InvalidOperationException($"SwarmUI error: {error.GetString()}");
            }

            if (!root.TryGetProperty("images", out var images) || images.GetArrayLength() == 0)
            {
                throw new InvalidOperationException("SwarmUI returned no images.");
            }

            var first = images[0].GetString() ?? throw new InvalidOperationException("SwarmUI image entry was empty.");
            return await ReadImageAsync(first, cancellationToken);
        }, cancellationToken);
    }

    private async Task<byte[]> ReadImageAsync(string reference, CancellationToken cancellationToken)
    {
        if (reference.StartsWith("data:", StringComparison.OrdinalIgnoreCase))
        {
            var comma = reference.IndexOf(',', StringComparison.Ordinal);
            if (comma < 0)
            {
                throw new InvalidOperationException("Malformed data URL from SwarmUI.");
            }

            return Convert.FromBase64String(reference[(comma + 1)..]);
        }

        var path = reference.StartsWith('/') ? reference : $"/{reference}";
        using var response = await _http.GetAsync(path, cancellationToken);
        response.EnsureSuccessStatusCode();
        return await response.Content.ReadAsByteArrayAsync(cancellationToken);
    }

    private async Task<T> RunWithSessionAsync<T>(Func<Task<T>> call, CancellationToken cancellationToken)
    {
        if (string.IsNullOrWhiteSpace(_sessionId))
        {
            await RefreshSessionAsync(cancellationToken);
        }

        try
        {
            return await call();
        }
        catch (SwarmUiSessionInvalidException)
        {
            await RefreshSessionAsync(cancellationToken);
            return await call();
        }
    }

    private async Task RefreshSessionAsync(CancellationToken cancellationToken)
    {
        using var content = new StringContent("{}", Encoding.UTF8, "application/json");
        using var response = await _http.PostAsync("/API/GetNewSession", content, cancellationToken);
        response.EnsureSuccessStatusCode();

        var json = await response.Content.ReadAsStringAsync(cancellationToken);
        using var document = JsonDocument.Parse(json);
        _sessionId = document.RootElement.GetProperty("session_id").GetString()
            ?? throw new InvalidOperationException("SwarmUI did not return a session_id.");
        _logger.LogDebug("SwarmUI session refreshed.");
    }

    internal sealed class SwarmUiSessionInvalidException : Exception;
}

internal static class SwarmUiClientRegistration
{
    public static IServiceCollection AddSwarmUi(this IServiceCollection services, IConfiguration configuration)
    {
        services
            .AddOptions<SwarmUiOptions>()
            .Bind(configuration.GetSection(SwarmUiOptions.SectionName))
            .Validate(options => options.TimeoutSeconds is > 0 and <= 3600, "SwarmUi:TimeoutSeconds must be between 1 and 3600.")
            .ValidateOnStart();

        var http = services.AddHttpClient<SwarmUiClient>((sp, client) =>
        {
            var options = sp.GetRequiredService<IOptions<SwarmUiOptions>>().Value;
            if (!string.IsNullOrWhiteSpace(options.BaseUrl))
            {
                client.BaseAddress = new Uri(options.BaseUrl.TrimEnd('/') + "/");
                client.Timeout = TimeSpan.FromSeconds(options.TimeoutSeconds);
            }

            if (!string.IsNullOrWhiteSpace(options.Authorization))
            {
                client.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", options.Authorization);
            }
        });

        // AddServiceDefaults puts a 10s AttemptTimeout on every HttpClient. Image generation
        // is a long POST and not safe to retry; HttpClient.Timeout (SwarmUi:TimeoutSeconds) is
        // the only deadline that should apply.
#pragma warning disable EXTEXP0001
        http.RemoveAllResilienceHandlers();
#pragma warning restore EXTEXP0001

        return services;
    }
}