StabilityMatrix/StabilityMatrix.Avalonia/Helpers/ImageSearcher.cs at e944f5a8c2dcb3b45c91569e885b19b4459db8f8 · LykosAI/StabilityMatrix · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
using FuzzySharp;
using FuzzySharp.PreProcess;
using StabilityMatrix.Core.Models.Database;

namespace StabilityMatrix.Avalonia.Helpers;

public class ImageSearcher
{
    public int MinimumFuzzScore { get; init; } = 80;

    public ImageSearchOptions SearchOptions { get; init; } = ImageSearchOptions.All;

    public Func<LocalImageFile, bool> GetPredicate(string? searchQuery)
    {
        if (string.IsNullOrEmpty(searchQuery))
        {
            return _ => true;
        }

        return file =>
        {
            if (file.FileName.Contains(searchQuery, StringComparison.OrdinalIgnoreCase))
            {
                return true;
            }

            if (
                SearchOptions.HasFlag(ImageSearchOptions.FileName)
                && Fuzz.WeightedRatio(searchQuery, file.FileName, PreprocessMode.Full) > MinimumFuzzScore
            )
            {
                return true;
            }

            // Generation params
            if (file.GenerationParameters is { } parameters)
            {
                if (
                    SearchOptions.HasFlag(ImageSearchOptions.PositivePrompt)
                        && (
                            parameters.PositivePrompt?.Contains(
                                searchQuery,
                                StringComparison.OrdinalIgnoreCase
                            ) ?? false
                        )
                    || SearchOptions.HasFlag(ImageSearchOptions.NegativePrompt)
                        && (
                            parameters.NegativePrompt?.Contains(
                                searchQuery,
                                StringComparison.OrdinalIgnoreCase
                            ) ?? false
                        )
                    || SearchOptions.HasFlag(ImageSearchOptions.Seed)
                        && parameters
                            .Seed.ToString()
                            .StartsWith(searchQuery, StringComparison.OrdinalIgnoreCase)
                    || SearchOptions.HasFlag(ImageSearchOptions.Sampler)
                        && (
                            parameters.Sampler?.StartsWith(searchQuery, StringComparison.OrdinalIgnoreCase)
                            ?? false
                        )
                    || SearchOptions.HasFlag(ImageSearchOptions.ModelName)
                        && (
                            parameters.ModelName?.StartsWith(searchQuery, StringComparison.OrdinalIgnoreCase)
                            ?? false
                        )
                )
                {
                    return true;
                }
            }

            return false;
        };
    }

    [Flags]
    public enum ImageSearchOptions
    {
        None = 0,
        FileName = 1 << 0,
        PositivePrompt = 1 << 1,
        NegativePrompt = 1 << 2,
        Seed = 1 << 3,
        Sampler = 1 << 4,
        ModelName = 1 << 5,
        All = int.MaxValue,
    }
}