JT
12 months ago
17 changed files with 769 additions and 114 deletions
@ -0,0 +1,36 @@
|
||||
using System; |
||||
using System.Threading.Tasks; |
||||
using StabilityMatrix.Core.Models; |
||||
using StabilityMatrix.Core.Models.FileInterfaces; |
||||
using StabilityMatrix.Core.Models.Progress; |
||||
using StabilityMatrix.Core.Services; |
||||
|
||||
namespace StabilityMatrix.Avalonia.DesignData; |
||||
|
||||
public class MockMetadataImportService : IMetadataImportService |
||||
{ |
||||
public Task ScanDirectoryForMissingInfo( |
||||
DirectoryPath directory, |
||||
IProgress<ProgressReport>? progress = null |
||||
) |
||||
{ |
||||
return Task.CompletedTask; |
||||
} |
||||
|
||||
public Task<ConnectedModelInfo?> GetMetadataForFile( |
||||
FilePath filePath, |
||||
IProgress<ProgressReport>? progress = null, |
||||
bool forceReimport = false |
||||
) |
||||
{ |
||||
return null; |
||||
} |
||||
|
||||
public Task UpdateExistingMetadata( |
||||
DirectoryPath directory, |
||||
IProgress<ProgressReport>? progress = null |
||||
) |
||||
{ |
||||
return Task.CompletedTask; |
||||
} |
||||
} |
||||
@ -0,0 +1,24 @@
|
||||
using StabilityMatrix.Core.Models; |
||||
using StabilityMatrix.Core.Models.FileInterfaces; |
||||
using StabilityMatrix.Core.Models.Progress; |
||||
|
||||
namespace StabilityMatrix.Core.Services; |
||||
|
||||
public interface IMetadataImportService |
||||
{ |
||||
Task ScanDirectoryForMissingInfo( |
||||
DirectoryPath directory, |
||||
IProgress<ProgressReport>? progress = null |
||||
); |
||||
|
||||
Task<ConnectedModelInfo?> GetMetadataForFile( |
||||
FilePath filePath, |
||||
IProgress<ProgressReport>? progress = null, |
||||
bool forceReimport = false |
||||
); |
||||
|
||||
Task UpdateExistingMetadata( |
||||
DirectoryPath directory, |
||||
IProgress<ProgressReport>? progress = null |
||||
); |
||||
} |
||||
@ -0,0 +1,308 @@
|
||||
using System.Diagnostics; |
||||
using System.Text.Json; |
||||
using Microsoft.Extensions.Logging; |
||||
using StabilityMatrix.Core.Attributes; |
||||
using StabilityMatrix.Core.Helper; |
||||
using StabilityMatrix.Core.Models; |
||||
using StabilityMatrix.Core.Models.Api; |
||||
using StabilityMatrix.Core.Models.Database; |
||||
using StabilityMatrix.Core.Models.FileInterfaces; |
||||
using StabilityMatrix.Core.Models.Progress; |
||||
|
||||
namespace StabilityMatrix.Core.Services; |
||||
|
||||
[Transient(typeof(IMetadataImportService))] |
||||
public class MetadataImportService( |
||||
ILogger<MetadataImportService> logger, |
||||
IDownloadService downloadService, |
||||
ModelFinder modelFinder |
||||
) : IMetadataImportService |
||||
{ |
||||
public async Task ScanDirectoryForMissingInfo( |
||||
DirectoryPath directory, |
||||
IProgress<ProgressReport>? progress = null |
||||
) |
||||
{ |
||||
progress?.Report(new ProgressReport(-1f, "Scanning directory...", isIndeterminate: true)); |
||||
|
||||
var checkpointsWithoutMetadata = directory |
||||
.EnumerateFiles(searchOption: SearchOption.AllDirectories) |
||||
.Where(FileHasNoCmInfo) |
||||
.ToList(); |
||||
|
||||
var scanned = 0; |
||||
var success = 0; |
||||
|
||||
foreach (var checkpointFilePath in checkpointsWithoutMetadata) |
||||
{ |
||||
if (scanned == 0) |
||||
{ |
||||
progress?.Report( |
||||
new ProgressReport( |
||||
current: scanned, |
||||
total: checkpointsWithoutMetadata.Count, |
||||
$"Scanning directory..." |
||||
) |
||||
); |
||||
} |
||||
else |
||||
{ |
||||
progress?.Report( |
||||
new ProgressReport( |
||||
current: scanned, |
||||
total: checkpointsWithoutMetadata.Count, |
||||
$"{success} files imported successfully" |
||||
) |
||||
); |
||||
} |
||||
|
||||
var fileNameWithoutExtension = checkpointFilePath.NameWithoutExtension; |
||||
var cmInfoPath = checkpointFilePath.Directory?.JoinFile( |
||||
$"{fileNameWithoutExtension}.cm-info.json" |
||||
); |
||||
var cmInfoExists = File.Exists(cmInfoPath); |
||||
if (cmInfoExists) |
||||
continue; |
||||
|
||||
var hashProgress = new Progress<ProgressReport>(report => |
||||
{ |
||||
progress?.Report( |
||||
new ProgressReport( |
||||
current: report.Current ?? 0, |
||||
total: report.Total ?? 0, |
||||
$"Scanning file {scanned}/{checkpointsWithoutMetadata.Count} ... {report.Percentage}%" |
||||
) |
||||
); |
||||
}); |
||||
|
||||
var blake3 = await GetBlake3Hash(cmInfoPath, checkpointFilePath, hashProgress) |
||||
.ConfigureAwait(false); |
||||
if (string.IsNullOrWhiteSpace(blake3)) |
||||
{ |
||||
logger.LogWarning($"Blake3 hash was null for {checkpointFilePath}"); |
||||
scanned++; |
||||
continue; |
||||
} |
||||
|
||||
var modelInfo = await modelFinder.RemoteFindModel(blake3).ConfigureAwait(false); |
||||
if (modelInfo == null) |
||||
{ |
||||
logger.LogWarning($"Could not find model for {blake3}"); |
||||
scanned++; |
||||
continue; |
||||
} |
||||
|
||||
var (model, modelVersion, modelFile) = modelInfo.Value; |
||||
|
||||
var updatedCmInfo = new ConnectedModelInfo( |
||||
model, |
||||
modelVersion, |
||||
modelFile, |
||||
DateTimeOffset.UtcNow |
||||
); |
||||
await updatedCmInfo |
||||
.SaveJsonToDirectory(checkpointFilePath.Directory, fileNameWithoutExtension) |
||||
.ConfigureAwait(false); |
||||
|
||||
var image = modelVersion.Images?.FirstOrDefault( |
||||
img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url)) |
||||
); |
||||
if (image == null) |
||||
{ |
||||
scanned++; |
||||
success++; |
||||
continue; |
||||
} |
||||
|
||||
await DownloadImage(image, checkpointFilePath, progress).ConfigureAwait(false); |
||||
|
||||
scanned++; |
||||
success++; |
||||
} |
||||
|
||||
progress?.Report( |
||||
new ProgressReport( |
||||
current: scanned, |
||||
total: checkpointsWithoutMetadata.Count, |
||||
$"Metadata found for {success}/{checkpointsWithoutMetadata.Count} files" |
||||
) |
||||
); |
||||
} |
||||
|
||||
private static bool FileHasNoCmInfo(FilePath file) |
||||
{ |
||||
return LocalModelFile.SupportedCheckpointExtensions.Contains(file.Extension) |
||||
&& !File.Exists(file.Directory?.JoinFile($"{file.NameWithoutExtension}.cm-info.json")); |
||||
} |
||||
|
||||
public async Task UpdateExistingMetadata( |
||||
DirectoryPath directory, |
||||
IProgress<ProgressReport>? progress = null |
||||
) |
||||
{ |
||||
progress?.Report(new ProgressReport(-1f, "Scanning directory...", isIndeterminate: true)); |
||||
|
||||
var cmInfoList = new Dictionary<FilePath, ConnectedModelInfo>(); |
||||
foreach ( |
||||
var cmInfoPath in directory.EnumerateFiles( |
||||
"*.cm-info.json", |
||||
SearchOption.AllDirectories |
||||
) |
||||
) |
||||
{ |
||||
var cmInfo = JsonSerializer.Deserialize<ConnectedModelInfo>( |
||||
await cmInfoPath.ReadAllTextAsync().ConfigureAwait(false) |
||||
); |
||||
if (cmInfo == null) |
||||
continue; |
||||
|
||||
cmInfoList.Add(cmInfoPath, cmInfo); |
||||
} |
||||
|
||||
var success = 1; |
||||
foreach (var (filePath, cmInfoValue) in cmInfoList) |
||||
{ |
||||
progress?.Report( |
||||
new ProgressReport( |
||||
current: success, |
||||
total: cmInfoList.Count, |
||||
$"Updating metadata {success}/{cmInfoList.Count}" |
||||
) |
||||
); |
||||
|
||||
var hash = cmInfoValue.Hashes.BLAKE3; |
||||
if (string.IsNullOrWhiteSpace(hash)) |
||||
continue; |
||||
|
||||
var modelInfo = await modelFinder.RemoteFindModel(hash).ConfigureAwait(false); |
||||
if (modelInfo == null) |
||||
{ |
||||
logger.LogWarning($"Could not find model for {hash}"); |
||||
continue; |
||||
} |
||||
|
||||
var (model, modelVersion, modelFile) = modelInfo.Value; |
||||
|
||||
var updatedCmInfo = new ConnectedModelInfo( |
||||
model, |
||||
modelVersion, |
||||
modelFile, |
||||
DateTimeOffset.UtcNow |
||||
); |
||||
|
||||
var nameWithoutCmInfo = filePath.NameWithoutExtension.Replace(".cm-info", string.Empty); |
||||
await updatedCmInfo |
||||
.SaveJsonToDirectory(filePath.Directory, nameWithoutCmInfo) |
||||
.ConfigureAwait(false); |
||||
|
||||
var image = modelVersion.Images?.FirstOrDefault( |
||||
img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url)) |
||||
); |
||||
if (image == null) |
||||
continue; |
||||
|
||||
await DownloadImage(image, filePath, progress).ConfigureAwait(false); |
||||
|
||||
success++; |
||||
} |
||||
} |
||||
|
||||
public async Task<ConnectedModelInfo?> GetMetadataForFile( |
||||
FilePath filePath, |
||||
IProgress<ProgressReport>? progress = null, |
||||
bool forceReimport = false |
||||
) |
||||
{ |
||||
progress?.Report(new ProgressReport(-1f, "Getting metadata...", isIndeterminate: true)); |
||||
|
||||
var fileNameWithoutExtension = filePath.NameWithoutExtension; |
||||
var cmInfoPath = filePath.Directory?.JoinFile($"{fileNameWithoutExtension}.cm-info.json"); |
||||
var cmInfoExists = File.Exists(cmInfoPath); |
||||
if (cmInfoExists && !forceReimport) |
||||
return null; |
||||
|
||||
var hashProgress = new Progress<ProgressReport>(report => |
||||
{ |
||||
progress?.Report( |
||||
new ProgressReport( |
||||
current: report.Current ?? 0, |
||||
total: report.Total ?? 0, |
||||
$"Getting metadata for {filePath} ... {report.Percentage}%" |
||||
) |
||||
); |
||||
}); |
||||
var blake3 = await GetBlake3Hash(cmInfoPath, filePath, hashProgress).ConfigureAwait(false); |
||||
|
||||
if (string.IsNullOrWhiteSpace(blake3)) |
||||
{ |
||||
logger.LogWarning($"Blake3 hash was null for {filePath}"); |
||||
return null; |
||||
} |
||||
|
||||
var modelInfo = await modelFinder.RemoteFindModel(blake3).ConfigureAwait(false); |
||||
if (modelInfo == null) |
||||
{ |
||||
logger.LogWarning($"Could not find model for {blake3}"); |
||||
return null; |
||||
} |
||||
|
||||
var (model, modelVersion, modelFile) = modelInfo.Value; |
||||
|
||||
var updatedCmInfo = new ConnectedModelInfo( |
||||
model, |
||||
modelVersion, |
||||
modelFile, |
||||
DateTimeOffset.UtcNow |
||||
); |
||||
await updatedCmInfo |
||||
.SaveJsonToDirectory(filePath.Directory, fileNameWithoutExtension) |
||||
.ConfigureAwait(false); |
||||
|
||||
var image = modelVersion.Images?.FirstOrDefault( |
||||
img => LocalModelFile.SupportedImageExtensions.Contains(Path.GetExtension(img.Url)) |
||||
); |
||||
|
||||
if (image == null) |
||||
return updatedCmInfo; |
||||
|
||||
await DownloadImage(image, filePath, progress).ConfigureAwait(false); |
||||
|
||||
return updatedCmInfo; |
||||
} |
||||
|
||||
private static async Task<string?> GetBlake3Hash( |
||||
FilePath? cmInfoPath, |
||||
FilePath checkpointFilePath, |
||||
IProgress<ProgressReport> hashProgress |
||||
) |
||||
{ |
||||
if (string.IsNullOrWhiteSpace(cmInfoPath?.ToString()) || !File.Exists(cmInfoPath)) |
||||
{ |
||||
return await FileHash |
||||
.GetBlake3Async(checkpointFilePath, hashProgress) |
||||
.ConfigureAwait(false); |
||||
} |
||||
|
||||
var cmInfo = JsonSerializer.Deserialize<ConnectedModelInfo>( |
||||
await cmInfoPath.ReadAllTextAsync().ConfigureAwait(false) |
||||
); |
||||
return cmInfo?.Hashes.BLAKE3; |
||||
} |
||||
|
||||
private Task DownloadImage( |
||||
CivitImage image, |
||||
FilePath modelFilePath, |
||||
IProgress<ProgressReport>? progress |
||||
) |
||||
{ |
||||
var imageExt = Path.GetExtension(image.Url).TrimStart('.'); |
||||
var nameWithoutCmInfo = modelFilePath.NameWithoutExtension.Replace( |
||||
".cm-info", |
||||
string.Empty |
||||
); |
||||
var imageDownloadPath = Path.GetFullPath( |
||||
Path.Combine(modelFilePath.Directory, $"{nameWithoutCmInfo}.preview.{imageExt}") |
||||
); |
||||
return downloadService.DownloadToFileAsync(image.Url, imageDownloadPath, progress); |
||||
} |
||||
} |
||||
Loading…
Reference in new issue