Skip to content

Checkpoints memory leak #1199

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Checkpoints memory leak #1199

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e4f5335
Merge branch 'batch-move'
RTHilton Mar 12, 2025
d3c6eee
Merge pull request #1 from RTHilton/batch-move
RTHilton Mar 19, 2025
4d992d5
Merge pull request #2 from RTHilton/tag-search
RTHilton Mar 19, 2025
2a5014e
Merge branch 'main' of https://github.com/RTHilton/StabilityMatrix
RTHilton Mar 19, 2025
3168636
Merge branch 'main' of https://github.com/RTHilton/StabilityMatrix
RTHilton Mar 24, 2025
9245a43
Add DisposeMany to ModelsCache disposable
RTHilton Mar 24, 2025
a6b04b5
Just in time loading of SafetensorMetadata when opening the metadata …
RTHilton Mar 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 23 additions & 35 deletions StabilityMatrix.Avalonia/ViewModels/CheckpointManager/CheckpointFileViewModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,50 +310,38 @@ private async Task RenameAsync()
[RelayCommand]
private async Task OpenSafetensorMetadataViewer()
{
if (!CheckpointFile.SafetensorMetadataParsed)
if (
!settingsManager.IsLibraryDirSet
|| new DirectoryPath(settingsManager.ModelsDirectory) is not { Exists: true } modelsDir
)
{
if (
!settingsManager.IsLibraryDirSet
|| new DirectoryPath(settingsManager.ModelsDirectory) is not { Exists: true } modelsDir
)
{
return;
}
return;
}

try
try
{
var safetensorPath = CheckpointFile.GetFullPath(modelsDir);
var metadata = await SafetensorMetadata.ParseAsync(safetensorPath);

var vm = vmFactory.Get<SafetensorMetadataViewModel>(vm =>
{
var safetensorPath = CheckpointFile.GetFullPath(modelsDir);
vm.ModelName = CheckpointFile.DisplayModelName;
vm.Metadata = metadata;
});

var metadata = await SafetensorMetadata.ParseAsync(safetensorPath);
var dialog = vm.GetDialog();
dialog.MinDialogHeight = 800;
dialog.MinDialogWidth = 700;
dialog.CloseButtonText = "Close";
dialog.DefaultButton = ContentDialogButton.Close;

CheckpointFile.SafetensorMetadataParsed = true;
CheckpointFile.SafetensorMetadata = metadata;
}
catch (Exception ex)
{
logger.LogWarning(ex, "Failed to parse safetensor metadata");
return;
}
await dialog.ShowAsync();
}

if (!CheckpointFile.SafetensorMetadataParsed)
catch (Exception ex)
{
logger.LogWarning(ex, "Failed to parse safetensor metadata");
return;
}

var vm = vmFactory.Get<SafetensorMetadataViewModel>(vm =>
{
vm.ModelName = CheckpointFile.DisplayModelName;
vm.Metadata = CheckpointFile.SafetensorMetadata;
});

var dialog = vm.GetDialog();
dialog.MinDialogHeight = 800;
dialog.MinDialogWidth = 700;
dialog.CloseButtonText = "Close";
dialog.DefaultButton = ContentDialogButton.Close;

await dialog.ShowAsync();
}

[RelayCommand]
Expand Down
34 changes: 27 additions & 7 deletions StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -358,6 +358,7 @@ or nameof(SortConnectedModelsFirst)
x
)
)
.DisposeMany()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This DisposeMany() call is great for ensuring that the CheckpointFileViewModel instances are properly disposed of when they are removed from the Models collection, preventing memory leaks. Good job!

.SortAndBind(Models, comparerObservable)
.WhenPropertyChanged(p => p.IsSelected)
.Throttle(TimeSpan.FromMilliseconds(50))
Expand Down Expand Up @@ -841,11 +842,16 @@ public async Task ImportFilesAsync(IEnumerable<string> files, DirectoryPath dest
.FirstOrDefault(x => x.Path == destinationFolder.FullPath);
}

public async Task MoveBetweenFolders(List<CheckpointFileViewModel>? sourceFiles, DirectoryPath destinationFolder)
public async Task MoveBetweenFolders(
List<CheckpointFileViewModel>? sourceFiles,
DirectoryPath destinationFolder
)
{
if (sourceFiles != null && sourceFiles.Count() > 0)
{
var sourceDirectory = Path.GetDirectoryName(sourceFiles[0].CheckpointFile.GetFullPath(settingsManager.ModelsDirectory));
var sourceDirectory = Path.GetDirectoryName(
sourceFiles[0].CheckpointFile.GetFullPath(settingsManager.ModelsDirectory)
);
foreach (CheckpointFileViewModel sourceFile in sourceFiles)
{
if (
Expand All @@ -863,13 +869,27 @@ public async Task MoveBetweenFolders(List<CheckpointFileViewModel>? sourceFiles,

try
{
var sourcePath = new FilePath(sourceFile.CheckpointFile.GetFullPath(settingsManager.ModelsDirectory));
var sourcePath = new FilePath(
sourceFile.CheckpointFile.GetFullPath(settingsManager.ModelsDirectory)
);
var fileNameWithoutExt = Path.GetFileNameWithoutExtension(sourcePath);
var sourceCmInfoPath = Path.Combine(sourcePath.Directory!, $"{fileNameWithoutExt}.cm-info.json");
var sourcePreviewPath = Path.Combine(sourcePath.Directory!, $"{fileNameWithoutExt}.preview.jpeg");
var sourceCmInfoPath = Path.Combine(
sourcePath.Directory!,
$"{fileNameWithoutExt}.cm-info.json"
);
var sourcePreviewPath = Path.Combine(
sourcePath.Directory!,
$"{fileNameWithoutExt}.preview.jpeg"
);
var destinationFilePath = Path.Combine(destinationFolder, sourcePath.Name);
var destinationCmInfoPath = Path.Combine(destinationFolder, $"{fileNameWithoutExt}.cm-info.json");
var destinationPreviewPath = Path.Combine(destinationFolder, $"{fileNameWithoutExt}.preview.jpeg");
var destinationCmInfoPath = Path.Combine(
destinationFolder,
$"{fileNameWithoutExt}.cm-info.json"
);
var destinationPreviewPath = Path.Combine(
destinationFolder,
$"{fileNameWithoutExt}.preview.jpeg"
);

// Move files
if (File.Exists(sourcePath))
Expand Down
6 changes: 0 additions & 6 deletions StabilityMatrix.Core/Models/Database/LocalModelFile.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,12 +154,6 @@ public override int GetHashCode()
[MemberNotNullWhen(true, nameof(ConnectedModelInfo))]
public bool HasCivitMetadata => HasConnectedModel && ConnectedModelInfo.ModelId != null;

[BsonIgnore]
public SafetensorMetadata? SafetensorMetadata { get; set; }

[BsonIgnore]
public bool SafetensorMetadataParsed { get; set; }

public string GetFullPath(string rootModelDirectory)
{
return Path.Combine(rootModelDirectory, RelativePath);
Expand Down
80 changes: 0 additions & 80 deletions StabilityMatrix.Core/Services/ModelIndexService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -544,86 +544,6 @@ await liteDbContext
);

EventManager.Instance.OnModelIndexChanged();

Task.Run(LoadSafetensorMetadataAsync)
.SafeFireAndForget(ex =>
{
logger.LogError(ex, "Error loading safetensor metadata");
});
}

private async Task LoadSafetensorMetadataAsync()
{
if (!settingsManager.IsLibraryDirSet)
{
logger.LogTrace("Safetensor metadata loading skipped, library directory not set");
return;
}

if (new DirectoryPath(settingsManager.ModelsDirectory) is not { Exists: true } modelsDir)
{
logger.LogTrace("Safetensor metadata loading skipped, model directory does not exist");
return;
}

await safetensorMetadataParseLock.WaitAsync().ConfigureAwait(false);
try
{
var stopwatch = Stopwatch.StartNew();
var readSuccess = 0;
var readFail = 0;
logger.LogInformation("Loading safetensor metadata...");

var models = ModelIndex
.Values.SelectMany(x => x)
.Where(m => !m.SafetensorMetadataParsed && m.RelativePath.EndsWith(".safetensors"));

await Parallel
.ForEachAsync(
models,
new ParallelOptions
{
MaxDegreeOfParallelism = Math.Max(1, Math.Min(Environment.ProcessorCount / 2, 6)),
TaskScheduler = TaskScheduler.Default,
},
async (model, token) =>
{
if (model.SafetensorMetadataParsed)
return;

if (!model.RelativePath.EndsWith(".safetensors"))
return;

try
{
var safetensorPath = model.GetFullPath(modelsDir);
var metadata = await SafetensorMetadata
.ParseAsync(safetensorPath)
.ConfigureAwait(false);
model.SafetensorMetadata = metadata;
model.SafetensorMetadataParsed = true;

Interlocked.Increment(ref readSuccess);
}
catch
{
Interlocked.Increment(ref readFail);
}
}
)
.ConfigureAwait(false);

logger.LogInformation(
"Loaded safetensor metadata for {Success} models, failed to load for {Fail} models in {Time:F2}ms",
readSuccess,
readFail,
stopwatch.Elapsed.TotalMilliseconds
);
}
finally
{
safetensorMetadataParseLock.Release();
}
}

/// <inheritdoc />
Expand Down
Loading