Merge pull request #518 from ionite34/choose-download-location

Added the ability to choose where model downloads are saved
9 months ago · b88784f069
7 changed files with 192 additions and 45 deletions
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 - Added new package: [StableSwarmUI](https://github.com/Stability-AI/StableSwarmUI) by Stability AI
 - Added new package: [Stable Diffusion WebUI Forge](https://github.com/lllyasviel/stable-diffusion-webui-forge) by lllyasviel
 - Added extension management for SD.Next and Stable Diffusion WebUI-UX
+- Added the ability to choose where CivitAI model downloads are saved

 ## v2.8.3
 ### Fixed
--- a/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs
+++ b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs
@ -644,6 +644,15 @@ namespace StabilityMatrix.Avalonia.Languages {
            }
        }
        
+        /// <summary>
+        ///   Looks up a localized string similar to Please select a download location..
+        /// </summary>
+        public static string Error_PleaseSelectDownloadLocation {
+            get {
+                return ResourceManager.GetString("Error_PleaseSelectDownloadLocation", resourceCulture);
+            }
+        }
+        
        /// <summary>
        ///   Looks up a localized string similar to About.
        /// </summary>
@ -1671,7 +1680,7 @@ namespace StabilityMatrix.Avalonia.Languages {
        }
        
        /// <summary>
-        ///   Looks up a localized string similar to .
+        ///   Looks up a localized string similar to None.
        /// </summary>
        public static string Label_NotificationOption_None {
            get {
@ -2039,6 +2048,15 @@ namespace StabilityMatrix.Avalonia.Languages {
            }
        }
        
+        /// <summary>
+        ///   Looks up a localized string similar to Select Download Location:.
+        /// </summary>
+        public static string Label_SelectDownloadLocation {
+            get {
+                return ResourceManager.GetString("Label_SelectDownloadLocation", resourceCulture);
+            }
+        }
+        
        /// <summary>
        ///   Looks up a localized string similar to Select new Data Directory.
        /// </summary>
--- a/StabilityMatrix.Avalonia/Languages/Resources.resx
+++ b/StabilityMatrix.Avalonia/Languages/Resources.resx
@ -946,7 +946,7 @@
    <value>Notifications</value>
  </data>
  <data name="Label_NotificationOption_None" xml:space="preserve">
-    <value />
+    <value>None</value>
  </data>
  <data name="Action_Download" xml:space="preserve">
    <value>Download</value>
@ -966,4 +966,10 @@
  <data name="Label_ComfyRequiredDetail" xml:space="preserve">
    <value>ComfyUI is required to install this package. Would you like to install it now?</value>
  </data>
+  <data name="Error_PleaseSelectDownloadLocation" xml:space="preserve">
+    <value>Please select a download location.</value>
+  </data>
+  <data name="Label_SelectDownloadLocation" xml:space="preserve">
+    <value>Select Download Location:</value>
+  </data>
 </root>
--- a/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj.DotSettings
+++ b/StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj.DotSettings
@ -1,3 +1,5 @@
 <wpf:ResourceDictionary xml:space="preserve" xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml" xmlns:s="clr-namespace:System;assembly=mscorlib" xmlns:wpf="http://schemas.microsoft.com/winfx/2006/xaml/presentation">
+	<s:String x:Key="/Default/CodeEditing/Localization/Localizable/@EntryValue">Yes</s:String>
+	<s:String x:Key="/Default/CodeEditing/Localization/LocalizableInspector/@EntryValue">Pessimistic</s:String>
 	<s:String x:Key="/Default/CodeInspection/Daemon/ConfigureAwaitAnalysisMode/@EntryValue">UI</s:String>
 	<s:Boolean x:Key="/Default/CodeInspection/NamespaceProvider/NamespaceFoldersToSkip/=controls_005Cpropertygrid/@EntryIndexedValue">True</s:Boolean></wpf:ResourceDictionary>
--- a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CheckpointBrowserCardViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CheckpointBrowserCardViewModel.cs
@ -190,13 +190,6 @@ public partial class CheckpointBrowserCardViewModel : Base.ProgressViewModel
        IsFavorite = settingsManager.Settings.FavoriteModels.Contains(CivitModel.Id);
    }

-    [RelayCommand]
-    private async Task Import(CivitModel model)
-    {
-        await DoImport(model);
-        CheckIfInstalled();
-    }
-
    [RelayCommand]
    private async Task ShowVersionDialog(CivitModel model)
    {
@ -220,7 +213,7 @@ public partial class CheckpointBrowserCardViewModel : Base.ProgressViewModel
            IsSecondaryButtonEnabled = false,
            IsFooterVisible = false,
            MaxDialogWidth = 750,
-            MaxDialogHeight = 850,
+            MaxDialogHeight = 950,
        };

        var prunedDescription = PruneDescription(model);
@ -229,6 +222,7 @@ public partial class CheckpointBrowserCardViewModel : Base.ProgressViewModel
        viewModel.Dialog = dialog;
        viewModel.Title = model.Name;
        viewModel.Description = prunedDescription;
+        viewModel.CivitModel = model;
        viewModel.Versions = versions
            .Select(
                version =>
@ -252,8 +246,21 @@ public partial class CheckpointBrowserCardViewModel : Base.ProgressViewModel
        var selectedVersion = viewModel?.SelectedVersionViewModel?.ModelVersion;
        var selectedFile = viewModel?.SelectedFile?.CivitFile;

+        DirectoryPath downloadPath;
+        if (viewModel?.IsCustomSelected is true)
+        {
+            downloadPath = viewModel.CustomInstallLocation;
+        }
+        else
+        {
+            var subFolder =
+                viewModel?.SelectedInstallLocation
+                ?? Path.Combine("Models", model.Type.ConvertTo<SharedFolderType>().GetStringValue());
+            downloadPath = Path.Combine(settingsManager.LibraryDir, subFolder);
+        }
+
        await Task.Delay(100);
-        await DoImport(model, selectedVersion, selectedFile);
+        await DoImport(model, downloadPath, selectedVersion, selectedFile);
    }

    private static string PruneDescription(CivitModel model)
@ -322,6 +329,7 @@ public partial class CheckpointBrowserCardViewModel : Base.ProgressViewModel

    private async Task DoImport(
        CivitModel model,
+        DirectoryPath downloadFolder,
        CivitModelVersion? selectedVersion = null,
        CivitFile? selectedFile = null
    )
@ -362,18 +370,13 @@ public partial class CheckpointBrowserCardViewModel : Base.ProgressViewModel
            return;
        }

-        var rootModelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory);
-
-        var downloadDirectory = rootModelsDirectory.JoinDir(
-            model.Type.ConvertTo<SharedFolderType>().GetStringValue()
-        );
        // Folders might be missing if user didn't install any packages yet
-        downloadDirectory.Create();
+        downloadFolder.Create();

-        var downloadPath = downloadDirectory.JoinFile(modelFile.Name);
+        var downloadPath = downloadFolder.JoinFile(modelFile.Name);

        // Download model info and preview first
-        var cmInfoPath = await SaveCmInfo(model, modelVersion, modelFile, downloadDirectory);
+        var cmInfoPath = await SaveCmInfo(model, modelVersion, modelFile, downloadFolder);
        var previewImagePath = await SavePreviewImage(modelVersion, downloadPath);

        // Create tracked download
--- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs
+++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/SelectModelVersionViewModel.cs
@ -6,6 +6,7 @@ using System.Linq;
 using System.Threading.Tasks;
 using Avalonia.Controls.Notifications;
 using Avalonia.Media.Imaging;
+using Avalonia.Platform.Storage;
 using Avalonia.Threading;
 using CommunityToolkit.Mvvm.ComponentModel;
 using FluentAvalonia.UI.Controls;
@ -15,7 +16,10 @@ using StabilityMatrix.Avalonia.Models;
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Base;
 using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.Api;
 using StabilityMatrix.Core.Models.FileInterfaces;
 using StabilityMatrix.Core.Services;

@ -23,17 +27,20 @@ namespace StabilityMatrix.Avalonia.ViewModels.Dialogs;

 [ManagedService]
 [Transient]
-public partial class SelectModelVersionViewModel : ContentDialogViewModelBase
+public partial class SelectModelVersionViewModel(
+    ISettingsManager settingsManager,
+    IDownloadService downloadService,
+    IModelIndexService modelIndexService,
+    INotificationService notificationService
+) : ContentDialogViewModelBase
 {
-    private readonly ISettingsManager settingsManager;
-    private readonly IDownloadService downloadService;
-    private readonly IModelIndexService modelIndexService;
-    private readonly INotificationService notificationService;
+    private readonly IDownloadService downloadService = downloadService;

    public required ContentDialog Dialog { get; set; }
    public required IReadOnlyList<ModelVersionViewModel> Versions { get; set; }
    public required string Description { get; set; }
    public required string Title { get; set; }
+    public required CivitModel CivitModel { get; set; }

    [ObservableProperty]
    private Bitmap? previewImage;
@ -63,25 +70,28 @@ public partial class SelectModelVersionViewModel : ContentDialogViewModelBase
    [ObservableProperty]
    private string importTooltip = string.Empty;

-    public int DisplayedPageNumber => SelectedImageIndex + 1;
+    [ObservableProperty]
+    [NotifyPropertyChangedFor(nameof(IsCustomSelected))]
+    [NotifyPropertyChangedFor(nameof(ShowEmptyPathWarning))]
+    private string selectedInstallLocation = string.Empty;

-    public SelectModelVersionViewModel(
-        ISettingsManager settingsManager,
-        IDownloadService downloadService,
-        IModelIndexService modelIndexService,
-        INotificationService notificationService
-    )
-    {
-        this.settingsManager = settingsManager;
-        this.downloadService = downloadService;
-        this.modelIndexService = modelIndexService;
-        this.notificationService = notificationService;
-    }
+    [ObservableProperty]
+    private ObservableCollection<string> availableInstallLocations = [];
+
+    [ObservableProperty]
+    [NotifyPropertyChangedFor(nameof(ShowEmptyPathWarning))]
+    private string customInstallLocation = string.Empty;
+
+    public bool IsCustomSelected => SelectedInstallLocation == "Custom...";
+    public bool ShowEmptyPathWarning => IsCustomSelected && string.IsNullOrWhiteSpace(CustomInstallLocation);
+
+    public int DisplayedPageNumber => SelectedImageIndex + 1;

    public override void OnLoaded()
    {
        SelectedVersionViewModel = Versions[0];
        CanGoToNextImage = true;
+        LoadInstallLocations();
    }

    partial void OnSelectedVersionViewModelChanged(ModelVersionViewModel? value)
@ -138,7 +148,26 @@ public partial class SelectModelVersionViewModel : ContentDialogViewModelBase
            ImportTooltip = "Please set the library directory in settings";
        }

-        IsImportEnabled = value?.CivitFile != null && canImport;
+        IsImportEnabled = value?.CivitFile != null && canImport && !ShowEmptyPathWarning;
+    }
+
+    partial void OnSelectedInstallLocationChanged(string value)
+    {
+        if (value.Equals("Custom...", StringComparison.OrdinalIgnoreCase))
+        {
+            Dispatcher.UIThread.InvokeAsync(SelectCustomFolder);
+        }
+        else
+        {
+            CustomInstallLocation = string.Empty;
+        }
+
+        IsImportEnabled = !ShowEmptyPathWarning;
+    }
+
+    partial void OnCustomInstallLocationChanged(string value)
+    {
+        IsImportEnabled = !ShowEmptyPathWarning;
    }

    public void Cancel()
@ -250,4 +279,48 @@ public partial class SelectModelVersionViewModel : ContentDialogViewModelBase
        CanGoToPreviousImage = SelectedImageIndex > 0;
        CanGoToNextImage = SelectedImageIndex < ImageUrls.Count - 1;
    }
+
+    public async Task SelectCustomFolder()
+    {
+        var files = await App.StorageProvider.OpenFolderPickerAsync(
+            new FolderPickerOpenOptions
+            {
+                Title = "Select Download Folder",
+                AllowMultiple = false,
+                SuggestedStartLocation = await App.StorageProvider.TryGetFolderFromPathAsync(
+                    Path.Combine(
+                        settingsManager.ModelsDirectory,
+                        CivitModel.Type.ConvertTo<SharedFolderType>().GetStringValue()
+                    )
+                )
+            }
+        );
+
+        if (files.FirstOrDefault()?.TryGetLocalPath() is { } path)
+        {
+            CustomInstallLocation = path;
+        }
+    }
+
+    private void LoadInstallLocations()
+    {
+        var rootModelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory);
+        var downloadDirectory = rootModelsDirectory.JoinDir(
+            CivitModel.Type.ConvertTo<SharedFolderType>().GetStringValue()
+        );
+        var installLocations = new ObservableCollection<string>
+        {
+            downloadDirectory.ToString().Replace(rootModelsDirectory, "Models")
+        };
+
+        foreach (var directory in downloadDirectory.EnumerateDirectories())
+        {
+            installLocations.Add(directory.ToString().Replace(rootModelsDirectory, "Models"));
+        }
+
+        installLocations.Add("Custom...");
+
+        AvailableInstallLocations = installLocations;
+        SelectedInstallLocation = installLocations.FirstOrDefault();
+    }
 }
--- a/StabilityMatrix.Avalonia/Views/Dialogs/SelectModelVersionDialog.axaml
+++ b/StabilityMatrix.Avalonia/Views/Dialogs/SelectModelVersionDialog.axaml
@ -21,7 +21,7 @@
        Margin="8"
        MinHeight="450"
        MinWidth="700"
-        RowDefinitions="Auto, *, Auto"
+        RowDefinitions="Auto, Auto, *, Auto"
        ColumnDefinitions="*,Auto,*">
        
        <Grid.Resources>
@ -86,8 +86,12 @@
            </StackPanel>
        </Grid>

+        <Border Grid.Row="0" Grid.Column="1"
+                MaxHeight="360"
+                Background="#33000000"
+                CornerRadius="4"/>
        <ListBox Grid.Row="0" Grid.Column="1"
-                  Margin="8"
+                  Margin="0, 8, 0, 8"
                  MaxHeight="360"
                  VerticalAlignment="Top"
                  ItemsSource="{Binding Versions}"
@ -120,15 +124,22 @@
            <ListBox.Template>
                <ControlTemplate>
                    <ScrollViewer>
-                        <ItemsPresenter />
+                        <ItemsPresenter Margin="8" />
                    </ScrollViewer>
                </ControlTemplate>
            </ListBox.Template>
        </ListBox>

+        <Border Grid.Row="0" Grid.Column="2"
+                MaxHeight="360"
+                Background="#33000000"
+                Margin="8,0"
+                CornerRadius="4"/>
        <ListBox Grid.Row="0" Grid.Column="2"
                 MaxHeight="360"
+                 Margin="8,8,8,0"
                 VerticalAlignment="Top"
+                 Background="#AA000000"
                 ItemsSource="{Binding SelectedVersionViewModel.CivitFileViewModels}"
                 SelectedItem="{Binding SelectedFile}">
            <ListBox.Template>
@ -179,20 +190,53 @@
            </ListBox.ItemTemplate>
        </ListBox>
        
+        <StackPanel Grid.Row="1"
+                    Grid.Column="0"
+                    Grid.ColumnSpan="3"
+                    Margin="0,8"
+                    Orientation="Vertical"
+                    HorizontalAlignment="Center"
+                    IsVisible="{Binding !SelectedFile.IsInstalled}">
+            <TextBlock Text="{x:Static lang:Resources.Label_SelectDownloadLocation}"
+                       TextAlignment="Center"
+                       HorizontalAlignment="Center"
+                       Margin="0,4"/>
+            <ComboBox ItemsSource="{Binding AvailableInstallLocations}"
+                      SelectedItem="{Binding SelectedInstallLocation}"
+                      HorizontalAlignment="Center"
+                      MinWidth="300"/>
+            <StackPanel Orientation="Horizontal"
+                        Margin="0,4"
+                        HorizontalAlignment="Center"
+                        IsVisible="{Binding IsCustomSelected}">
+                <TextBox Text="{Binding CustomInstallLocation}"
+                         MinWidth="260"/>
+                <Button Margin="4,0,0,0"
+                        VerticalAlignment="Stretch"
+                        Command="{Binding SelectCustomFolder}">
+                    <controls1:SymbolIcon Symbol="OpenFolder" />
+                </Button>
+            </StackPanel>
+            <TextBlock Text="{x:Static lang:Resources.Error_PleaseSelectDownloadLocation}"
+                       IsVisible="{Binding ShowEmptyPathWarning}"
+                       Foreground="{DynamicResource ThemeRedColor}"
+                       TextAlignment="Center"/>
+        </StackPanel>
+        
        <Expander Grid.Column="0"
                  Grid.ColumnSpan="3"
-                  Grid.Row="1"
+                  Grid.Row="2"
                  Header="{x:Static lang:Resources.Label_ModelDescription}"
                  ExpandDirection="Down"
-                  Margin="8, 16">
+                  Margin="8, 8">
            <ScrollViewer MaxHeight="200">
                <TextBlock Text="{Binding Description}"
                           TextWrapping="Wrap"
                           Margin="4"/>
            </ScrollViewer>
        </Expander>
-
-        <StackPanel Grid.Column="0" Grid.ColumnSpan="3" Grid.Row="2"
+        
+        <StackPanel Grid.Column="0" Grid.ColumnSpan="3" Grid.Row="3"
                    Orientation="Horizontal"
                    Margin="0,8,0,0"
                    HorizontalAlignment="Center">