using DynamicData.Tests;
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Models.FileInterfaces;
using JsonSerializer = System.Text.Json.JsonSerializer;
using Size = System.Drawing.Size;
namespace StabilityMatrix.Core.Models.Database;
///
/// Represents a locally indexed image file.
///
public record LocalImageFile
{
public required string AbsolutePath { get; init; }
///
/// Type of the model file.
///
public LocalImageFileType ImageType { get; init; }
///
/// Creation time of the file.
///
public DateTimeOffset CreatedAt { get; init; }
///
/// Last modified time of the file.
///
public DateTimeOffset LastModifiedAt { get; init; }
///
/// Generation parameters metadata of the file.
///
public GenerationParameters? GenerationParameters { get; init; }
///
/// Dimensions of the image
///
public Size? ImageSize { get; init; }
///
/// File name of the relative path.
///
public string FileName => Path.GetFileName(AbsolutePath);
///
/// File name of the relative path without extension.
///
public string FileNameWithoutExtension => Path.GetFileNameWithoutExtension(AbsolutePath);
public (
string? Parameters,
string? ParametersJson,
string? SMProject,
string? ComfyNodes
) ReadMetadata()
{
using var stream = new FileStream(
AbsolutePath,
FileMode.Open,
FileAccess.Read,
FileShare.Read
);
using var reader = new BinaryReader(stream);
var parameters = ImageMetadata.ReadTextChunk(reader, "parameters");
var parametersJson = ImageMetadata.ReadTextChunk(reader, "parameters-json");
var smProject = ImageMetadata.ReadTextChunk(reader, "smproj");
var comfyNodes = ImageMetadata.ReadTextChunk(reader, "prompt");
return (
string.IsNullOrEmpty(parameters) ? null : parameters,
string.IsNullOrEmpty(parametersJson) ? null : parametersJson,
string.IsNullOrEmpty(smProject) ? null : smProject,
string.IsNullOrEmpty(comfyNodes) ? null : comfyNodes
);
}
public static LocalImageFile FromPath(FilePath filePath)
{
// TODO: Support other types
const LocalImageFileType imageType =
LocalImageFileType.Inference | LocalImageFileType.TextToImage;
// Get metadata
using var stream = filePath.Info.OpenRead();
using var reader = new BinaryReader(stream);
var imageSize = ImageMetadata.GetImageSize(reader);
var metadata = ImageMetadata.ReadTextChunk(reader, "parameters-json");
GenerationParameters? genParams;
if (!string.IsNullOrWhiteSpace(metadata))
{
genParams = JsonSerializer.Deserialize(metadata);
}
else
{
metadata = ImageMetadata.ReadTextChunk(reader, "parameters");
GenerationParameters.TryParse(metadata, out genParams);
}
filePath.Info.Refresh();
return new LocalImageFile
{
AbsolutePath = filePath,
ImageType = imageType,
CreatedAt = filePath.Info.CreationTimeUtc,
LastModifiedAt = filePath.Info.LastWriteTimeUtc,
GenerationParameters = genParams,
ImageSize = imageSize
};
}
public static readonly HashSet SupportedImageExtensions =
new() { ".png", ".jpg", ".jpeg", ".webp" };
}