h-school/tests/HSchool.Server.Tests/SwarmUiLoraFormatterTests.cs

using System.Text.Json;
using HSchool.Server.Game;

namespace HSchool.Server.Tests;

public class SwarmUiLoraFormatterTests
{
    [Fact]
    public void FormatForApi_SplitsNamesAndWeights()
    {
        var (names, weights) = SwarmUiLoraFormatter.FormatForApi(
        [
            new SwarmUiLoraEntry { Name = "style.safetensors", Weight = 0.8 },
            new SwarmUiLoraEntry { Name = "face.safetensors", Weight = 1 },
        ]);

        Assert.Equal("style.safetensors,face.safetensors", names);
        Assert.Equal("0.8,1", weights);
    }

    [Fact]
    public void Concat_StacksModelPresetAndKind()
    {
        var stacked = SwarmUiLoraFormatter.Concat(
            [new SwarmUiLoraEntry { Name = "model-lora", Weight = 1 }],
            [new SwarmUiLoraEntry { Name = "preset-lora", Weight = 0.5 }],
            [new SwarmUiLoraEntry { Name = "kind-lora", Weight = 0.2 }]);

        Assert.Equal(["model-lora", "preset-lora", "kind-lora"], stacked.Select(entry => entry.Name));
    }

    [Fact]
    public void AppendEmbedTags_AddsSwarmEmbedSyntax()
    {
        var prompt = SwarmUiLoraFormatter.AppendEmbedTags(
            "a student",
            [
                new SwarmUiLoraEntry { Name = "pos.safetensors", Weight = 1 },
                new SwarmUiLoraEntry { Name = "soft.safetensors", Weight = 0.6 },
            ]);

        Assert.Contains("<embed:pos.safetensors>", prompt, StringComparison.Ordinal);
        Assert.Contains("<embed:soft.safetensors:0.6>", prompt, StringComparison.Ordinal);
        Assert.DoesNotContain("<lora:", prompt, StringComparison.Ordinal);
    }

    [Fact]
    public void DiscoveryParser_ReadsEmbeddings()
    {
        using var document = JsonDocument.Parse(
            """
            {
              "models": {
                "Stable-Diffusion": ["base.safetensors"],
                "LoRA": ["style.safetensors"],
                "Embedding": ["pos.safetensors", "neg.safetensors"]
              },
              "list": []
            }
            """);

        var discovery = SwarmUiDiscoveryParser.Parse(document.RootElement);

        Assert.Equal(["pos.safetensors", "neg.safetensors"], discovery.Embeddings);
        Assert.Equal(["style.safetensors"], discovery.Loras);
    }

    [Fact]
    public void ToProfile_StacksLorasAndEmbeddingsAcrossLayers()
    {
        var model = new SwarmUiModelDefinition
        {
            Id = "m.safetensors",
            PositiveLoras = [new SwarmUiLoraEntry { Name = "model-lora", Weight = 1 }],
            PositiveEmbeddings = [new SwarmUiLoraEntry { Name = "model-embed", Weight = 1 }],
        };
        var preset = new SwarmUiPresetDefinition
        {
            Id = "default",
            Label = "Default",
            Model = "m.safetensors",
            PositiveLoras = [new SwarmUiLoraEntry { Name = "preset-lora", Weight = 0.5 }],
            Avatar = new SwarmUiKindPreset
            {
                Width = 512,
                Height = 512,
                PositiveEmbeddings = [new SwarmUiLoraEntry { Name = "avatar-embed", Weight = 0.8 }],
            },
        };

        var profile = preset.ToProfile(PortraitKind.Avatar, model);

        Assert.Equal(["model-lora", "preset-lora"], profile.PositiveLoras.Select(entry => entry.Name));
        Assert.Equal(["model-embed", "avatar-embed"], profile.PositiveEmbeddings.Select(entry => entry.Name));
    }

    [Fact]
    public void Build_PutsEmbedsInPromptAndOmitsLoraTags()
    {
        var model = new SwarmUiModelDefinition
        {
            Id = "m.safetensors",
            PositiveLoras = [new SwarmUiLoraEntry { Name = "style-lora", Weight = 1 }],
            NegativeEmbeddings = [new SwarmUiLoraEntry { Name = "neg-embed", Weight = 1 }],
        };
        var preset = SwarmUiPresetDefinition.CreateDefault();
        preset.PositiveEmbeddings = [new SwarmUiLoraEntry { Name = "pos-embed", Weight = 1 }];
        var card = new HSchool.Server.Api.PersonCardResponse(
            "f0.c0",
            "Maria Ivanova",
            "Ivanova",
            "Maria",
            "",
            true,
            16,
            new DateTime(2000, 3, 14, 0, 0, 0, DateTimeKind.Utc),
            ["student"],
            5,
            "A",
            "class-1",
            null,
            null,
            [],
            [],
            [],
            [],
            null,
            null,
            new HSchool.Server.Api.PersonFamilyResponse([], [], [], []),
            [],
            [],
            0f,
            0f);

        var (positive, negative) = PortraitPromptBuilder.Build(
            card,
            preset.ToProfile(PortraitKind.Avatar, model),
            PortraitKind.Avatar);

        Assert.Contains("<embed:pos-embed>", positive, StringComparison.Ordinal);
        Assert.Contains("<embed:neg-embed>", negative, StringComparison.Ordinal);
        Assert.DoesNotContain("<lora:", positive, StringComparison.Ordinal);
        Assert.DoesNotContain("<lora:", negative, StringComparison.Ordinal);
    }
}