Add download dialog for controlnet model

StabilityMatrix.Avalonia/Controls/ControlNetCard.axaml

@@ -8,6 +8,8 @@
     xmlns:sg="clr-namespace:SpacedGridControl.Avalonia;assembly=SpacedGridControl.Avalonia"
     xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia"
     xmlns:vmInference="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Inference"
+    xmlns:input="clr-namespace:FluentAvalonia.UI.Input;assembly=FluentAvalonia"
+    xmlns:fluentIcons="clr-namespace:FluentIcons.FluentAvalonia;assembly=FluentIcons.FluentAvalonia"
     x:DataType="vmInference:ControlNetCardViewModel">
     <Design.PreviewWith>
         <StackPanel Width="350" Height="400">
@@ -59,12 +61,48 @@
                             Text="{x:Static lang:Resources.Label_Model}"
                             TextAlignment="Left" />
                         <ui:FAComboBox
+                            x:Name="PART_ModelComboBox"
                             Grid.Row="3"
                             Grid.Column="1"
                             HorizontalAlignment="Stretch"
                             ItemsSource="{Binding ClientManager.ControlNetModels}"
                             SelectedItem="{Binding SelectedModel}"
-                            Theme="{StaticResource FAComboBoxHybridModelTheme}" />
+                            ItemContainerTheme="{StaticResource FAComboBoxItemHybridModelTheme}">
+                            
+                            <ui:FAComboBox.Resources>
+                                    <input:StandardUICommand 
+                                        x:Key="RemoteDownloadCommand" 
+                                        Command="{Binding RemoteDownloadCommand}"/>
+                                </ui:FAComboBox.Resources>
+                                
+                                <ui:FAComboBox.DataTemplates>
+                                    <controls:HybridModelTemplateSelector>
+                                        <DataTemplate DataType="models:HybridModelFile" x:Key="{x:Static models:HybridModelType.Downloadable}">
+                                            <Grid ColumnDefinitions="*,Auto">
+                                                <TextBlock
+                                                    Text="{Binding ShortDisplayName}"
+                                                    Foreground="{DynamicResource ThemeGreyColor}"/>
+                                                <Button 
+                                                    Grid.Column="1" 
+                                                    Classes="transparent-full"
+                                                    Margin="8,0,0,0"
+                                                    Padding="0">
+                                                    <fluentIcons:SymbolIcon
+                                                        VerticalAlignment="Center"
+                                                        FontSize="18"
+                                                        Foreground="{DynamicResource ThemeGreyColor}"
+                                                        IsFilled="True"
+                                                        Symbol="CloudArrowDown" />
+                                                </Button>
+                                            </Grid>
+                                        </DataTemplate>
+                                        <DataTemplate DataType="models:HybridModelFile" x:Key="{x:Static models:HybridModelType.None}">
+                                            <TextBlock Text="{Binding ShortDisplayName}"/>
+                                        </DataTemplate>
+                                    </controls:HybridModelTemplateSelector>
+                                </ui:FAComboBox.DataTemplates>
+                            
+                        </ui:FAComboBox>
 
                         <sg:SpacedGrid
                             RowDefinitions="Auto,Auto,Auto,Auto"

StabilityMatrix.Avalonia/Controls/ControlNetCard.axaml.cs

@@ -1,7 +1,52 @@
-using Avalonia.Controls.Primitives;
+using AsyncAwaitBestPractices;
+using Avalonia.Controls;
+using Avalonia.Controls.Primitives;
+using FluentAvalonia.UI.Controls;
+using StabilityMatrix.Avalonia.ViewModels.Inference;
 using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Models;
 
 namespace StabilityMatrix.Avalonia.Controls;
 
 [Transient]
-public class ControlNetCard : TemplatedControl { }
+public class ControlNetCard : TemplatedControl
+{
+    /// <inheritdoc />
+    protected override void OnApplyTemplate(TemplateAppliedEventArgs e)
+    {
+        base.OnApplyTemplate(e);
+
+        var upscalerComboBox = e.NameScope.Find("PART_ModelComboBox") as FAComboBox;
+        upscalerComboBox!.SelectionChanged += UpscalerComboBox_OnSelectionChanged;
+    }
+
+    private void UpscalerComboBox_OnSelectionChanged(object? sender, SelectionChangedEventArgs e)
+    {
+        if (e.AddedItems.Count == 0)
+            return;
+
+        var item = e.AddedItems[0];
+        if (item is HybridModelFile { IsDownloadable: true })
+        {
+            // Reset the selection
+            e.Handled = true;
+
+            if (
+                e.RemovedItems.Count > 0
+                && e.RemovedItems[0] is HybridModelFile { IsDownloadable: false } removedItem
+            )
+            {
+                (sender as FAComboBox)!.SelectedItem = removedItem;
+            }
+            else
+            {
+                (sender as FAComboBox)!.SelectedItem = null;
+            }
+
+            // Show dialog to download the model
+            (DataContext as ControlNetCardViewModel)!.RemoteDownloadCommand
+                .ExecuteAsync(item)
+                .SafeFireAndForget();
+        }
+    }
+}

StabilityMatrix.Avalonia/Controls/HybridModelTemplateSelector.cs

@@ -24,15 +24,15 @@ public class HybridModelTemplateSelector : IDataTemplate
     // Build the DataTemplate here
     public Control Build(object? data)
     {
-        if (data is not HybridModelTemplateSelector card)
+        if (data is not HybridModelFile modelFile)
             throw new ArgumentException(null, nameof(data));
 
-        if (Templates.TryGetValue(card.Type, out var type))
+        if (Templates.TryGetValue(modelFile.Type, out var type))
         {
-            return type.Build(card)!;
+            return type.Build(modelFile)!;
         }
 
         // Fallback to Local
-        return Templates[HybridModelType.Local].Build(card)!;
+        return Templates[HybridModelType.None].Build(modelFile)!;
     }
 }

StabilityMatrix.Avalonia/Services/InferenceClientManager.cs

@@ -69,6 +69,9 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
     private readonly SourceCache<HybridModelFile, string> controlNetModelsSource =
         new(p => p.GetId());
 
+    private readonly SourceCache<HybridModelFile, string> downloadableControlNetModelsSource =
+        new(p => p.GetId());
+
     public IObservableCollection<HybridModelFile> ControlNetModels { get; } =
         new ObservableCollectionExtended<HybridModelFile>();
 
@@ -119,10 +122,11 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
 
         controlNetModelsSource
             .Connect()
-            .SortBy(
+            .Or(downloadableControlNetModelsSource.Connect())
-                f => f.ShortDisplayName,
+            .Sort(
-                SortDirection.Ascending,
+                SortExpressionComparer<HybridModelFile>
-                SortOptimisations.ComparesImmutableValuesOnly
+                    .Ascending(f => f.Type)
+                    .ThenByAscending(f => f.ShortDisplayName)
             )
             .DeferUntilLoaded()
             .Bind(ControlNetModels)
@@ -274,6 +278,15 @@ public partial class InferenceClientManager : ObservableObject, IInferenceClient
             HybridModelFile.Comparer
         );
 
+        // Downloadable ControlNet models
+        var downloadableControlNets = RemoteModels.ControlNetModels.Where(
+            u => !modelUpscalersSource.Lookup(u.GetId()).HasValue
+        );
+        downloadableControlNetModelsSource.EditDiff(
+            downloadableControlNets,
+            HybridModelFile.Comparer
+        );
+
         // Load local VAE models
         vaeModelsSource.EditDiff(
             modelIndexService

StabilityMatrix.Avalonia/ViewModels/Inference/ControlNetCardViewModel.cs

@@ -1,10 +1,19 @@
-using System.ComponentModel.DataAnnotations;
+using System;
+using System.ComponentModel.DataAnnotations;
+using System.Threading.Tasks;
 using CommunityToolkit.Mvvm.ComponentModel;
+using CommunityToolkit.Mvvm.Input;
+using FluentAvalonia.UI.Controls;
 using StabilityMatrix.Avalonia.Controls;
 using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Avalonia.ViewModels.Dialogs;
 using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Extensions;
+using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Models;
+using StabilityMatrix.Core.Models.FileInterfaces;
+using StabilityMatrix.Core.Services;
 
 namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 
@@ -13,6 +22,10 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 [Transient]
 public partial class ControlNetCardViewModel : LoadableViewModelBase
 {
+    private readonly ITrackedDownloadService trackedDownloadService;
+    private readonly ISettingsManager settingsManager;
+    private readonly ServiceManager<ViewModelBase> vmFactory;
+
     [ObservableProperty]
     [Required]
     private HybridModelFile? selectedModel;
@@ -41,11 +54,50 @@ public partial class ControlNetCardViewModel : LoadableViewModelBase
     public IInferenceClientManager ClientManager { get; }
 
     public ControlNetCardViewModel(
+        ITrackedDownloadService trackedDownloadService,
+        ISettingsManager settingsManager,
         IInferenceClientManager clientManager,
         ServiceManager<ViewModelBase> vmFactory
     )
     {
+        this.trackedDownloadService = trackedDownloadService;
+        this.settingsManager = settingsManager;
+        this.vmFactory = vmFactory;
+
         ClientManager = clientManager;
         SelectImageCardViewModel = vmFactory.Get<SelectImageCardViewModel>();
     }
+
+    [RelayCommand]
+    private async Task RemoteDownload(HybridModelFile? modelFile)
+    {
+        if (modelFile?.DownloadableResource is not { } resource)
+            return;
+
+        var sharedFolderType =
+            resource.ContextType as SharedFolderType?
+            ?? throw new InvalidOperationException("ContextType is not SharedFolderType");
+
+        var confirmDialog = vmFactory.Get<DownloadResourceViewModel>();
+        confirmDialog.Resource = resource;
+        confirmDialog.FileName = modelFile.FileName;
+
+        if (await confirmDialog.GetDialog().ShowAsync() != ContentDialogResult.Primary)
+        {
+            return;
+        }
+
+        var modelsDir = new DirectoryPath(settingsManager.ModelsDirectory).JoinDir(
+            sharedFolderType.GetStringValue()
+        );
+
+        var download = trackedDownloadService.NewDownload(
+            resource.Url,
+            modelsDir.JoinFile(modelFile.FileName)
+        );
+        download.ContextAction = new ModelPostDownloadContextAction();
+        download.Start();
+
+        EventManager.Instance.OnToggleProgressFlyout();
+    }
 }