using DynamicData.Tests;
using MetadataExtractor.Formats.Exif;
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()
{
if (AbsolutePath.EndsWith("webp"))
{
var paramsJson = ImageMetadata.ReadTextChunkFromWebp(
AbsolutePath,
ExifDirectoryBase.TagImageDescription
);
var smProj = ImageMetadata.ReadTextChunkFromWebp(AbsolutePath, ExifDirectoryBase.TagSoftware);
return (null, paramsJson, smProj, null);
}
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;
if (filePath.Extension.Equals(".webp", StringComparison.OrdinalIgnoreCase))
{
var paramsJson = ImageMetadata.ReadTextChunkFromWebp(
filePath,
ExifDirectoryBase.TagImageDescription
);
var parameters = string.IsNullOrWhiteSpace(paramsJson)
? null
: JsonSerializer.Deserialize(paramsJson);
filePath.Info.Refresh();
return new LocalImageFile
{
AbsolutePath = filePath,
ImageType = imageType,
CreatedAt = filePath.Info.CreationTimeUtc,
LastModifiedAt = filePath.Info.LastWriteTimeUtc,
GenerationParameters = parameters,
ImageSize = new Size(parameters?.Width ?? 0, parameters?.Height ?? 0)
};
}
// 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 =
[
".png",
".jpg",
".jpeg",
".gif",
".webp"
];
}