Add upscale option and fix GAN upscales

StabilityMatrix.Avalonia/Controls/ImageGalleryCard.axaml

@@ -65,11 +65,19 @@
                                                 StretchDirection="Both">
                                                 <controls:BetterAdvancedImage.ContextFlyout>
                                                     <ui:FAMenuFlyout>
+                                                        <ui:MenuFlyoutItem 
+                                                            Command="{ReflectionBinding #ImageCarousel.DataContext.FlyoutPreviewCommand}" 
+                                                            CommandParameter="{Binding #ImageView.CurrentImage}"
+                                                            Text="Expanded Preview"
+                                                            HotKey="Space"
+                                                            IconSource="ZoomInFilled" />
+                                                        <ui:MenuFlyoutSeparator/>
                                                         <ui:MenuFlyoutItem 
                                                             IsEnabled="{OnPlatform Windows=True, Default=False}"
                                                             Command="{ReflectionBinding #ImageCarousel.DataContext.FlyoutCopyCommand}" 
                                                             CommandParameter="{Binding #ImageView.CurrentImage}"
                                                             Text="Copy"
+                                                            HotKey="Ctrl+C"
                                                             IconSource="Copy" />
                                                     </ui:FAMenuFlyout>
                                                 </controls:BetterAdvancedImage.ContextFlyout>

StabilityMatrix.Avalonia/Controls/SamplerCard.axaml

@@ -6,7 +6,7 @@
         x:DataType="vmInference:SamplerCardViewModel"
         xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData">
     <Design.PreviewWith>
-        <StackPanel MinWidth="300" Spacing="16">
+        <StackPanel MinWidth="350" Spacing="16">
             <controls:SamplerCard DataContext="{x:Static mocks:DesignData.SamplerCardViewModel}"/>
             <controls:SamplerCard DataContext="{x:Static mocks:DesignData.SamplerCardViewModelScaleMode}"/>
         </StackPanel>
@@ -21,7 +21,7 @@
                     <StackPanel 
                         Margin="8" 
                         HorizontalAlignment="{TemplateBinding HorizontalAlignment}"
-                        Spacing="16">
+                        Spacing="12">
                         <Grid ColumnDefinitions="Auto,*" RowDefinitions="*,*,*">
                             <!-- Sampler -->
                             <TextBlock
@@ -77,8 +77,7 @@
                         
                         <StackPanel>
                             <!-- Denoise Strength -->
-                            <Grid IsVisible="{Binding IsDenoiseStrengthEnabled}">
+                            <StackPanel IsVisible="{Binding IsDenoiseStrengthEnabled}">
-                                <StackPanel>
                                 <Grid ColumnDefinitions="*,Auto">
                                     <TextBlock 
                                         VerticalAlignment="Center"
@@ -91,6 +90,7 @@
                                         SimpleNumberFormat="F2"
                                         Value="{Binding DenoiseStrength}"
                                         HorizontalAlignment="Stretch"
+                                        MinWidth="70"
                                         SpinButtonPlacementMode="Compact"/>
                                 </Grid>
                                 <Slider 
@@ -98,13 +98,14 @@
                                     Maximum="1"
                                     Value="{Binding DenoiseStrength}"
                                     TickFrequency="1"
+                                    Margin="0,0,0,-4"
                                     TickPlacement="BottomRight"/>
                             </StackPanel>
-                            </Grid>
                             
                             <!-- Dimensions (Absolute) -->
                             <Grid 
                                 IsVisible="{Binding IsDimensionsEnabled}"
+                                Margin="0,4,0,0"
                                 ColumnDefinitions="*,Auto,*" 
                                 RowDefinitions="Auto,*">
                                 <TextBlock

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -524,7 +524,11 @@ public static class DesignData
     }
 
     public static ImageViewerViewModel ImageViewerViewModel 
-        => DialogFactory.Get<ImageViewerViewModel>();
+        => DialogFactory.Get<ImageViewerViewModel>(vm =>
+        {
+            vm.Source =
+                "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/a318ac1f-3ad0-48ac-98cc-79126febcc17/width=1024";
+        });
     
     public static Indexer Types => new();
 

StabilityMatrix.Avalonia/ViewModels/Inference/ImageGalleryCardViewModel.cs

@@ -11,7 +11,10 @@ using NLog;
 using StabilityMatrix.Avalonia.Controls;
 using StabilityMatrix.Avalonia.Helpers;
 using StabilityMatrix.Avalonia.Models;
+using StabilityMatrix.Avalonia.Services;
 using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Avalonia.ViewModels.Dialogs;
+using StabilityMatrix.Avalonia.Views.Dialogs;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Helper;
 
@@ -21,6 +24,7 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 public partial class ImageGalleryCardViewModel : ViewModelBase
 {
     private static readonly Logger Logger = LogManager.GetCurrentClassLogger();
+    private readonly ServiceManager<ViewModelBase> vmFactory;
     
     [ObservableProperty]
     private bool isPreviewOverlayEnabled;
@@ -41,8 +45,10 @@ public partial class ImageGalleryCardViewModel : ViewModelBase
     public bool CanNavigateBack => SelectedImageIndex > 0;
     public bool CanNavigateForward => SelectedImageIndex < ImageSources.Count - 1;
 
-    public ImageGalleryCardViewModel()
+    public ImageGalleryCardViewModel(ServiceManager<ViewModelBase> vmFactory)
     {
+        this.vmFactory = vmFactory;
+        
         ImageSources.CollectionChanged += OnImageSourcesItemsChanged;
     }
     
@@ -82,7 +88,7 @@ public partial class ImageGalleryCardViewModel : ViewModelBase
             return;
         }
 
-        Logger.Trace($"FlyoutCopy is copying {image}");
+        Logger.Trace($"FlyoutCopy is copying bitmap...");
 
         await Task.Run(() =>
         {
@@ -92,4 +98,29 @@ public partial class ImageGalleryCardViewModel : ViewModelBase
             }
         });
     }
+    
+    [RelayCommand]
+    // ReSharper disable once UnusedMember.Local
+    private async Task FlyoutPreview(IImage? image)
+    {
+        if (image is null)
+        {
+            Logger.Trace("FlyoutPreview: image is null");
+            return;
+        }
+
+        Logger.Trace($"FlyoutPreview opening...");
+
+        var viewerVm = vmFactory.Get<ImageViewerViewModel>();
+        viewerVm.Image = (Bitmap) image;
+
+        var dialog = new BetterContentDialog
+        {
+            Content = new ImageViewerDialog
+            {
+                DataContext = viewerVm,
+            }
+        };
+
+    }
 }

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -49,17 +49,15 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
     public StackCardViewModel StackCardViewModel { get; }
 
     public UpscalerCardViewModel UpscalerCardViewModel =>
-        StackCardViewModel
+        StackCardViewModel.GetCard<StackExpanderViewModel>().GetCard<UpscalerCardViewModel>();
-        .GetCard<StackExpanderViewModel>()
-        .GetCard<UpscalerCardViewModel>();
 
     public SamplerCardViewModel HiresSamplerCardViewModel =>
-        StackCardViewModel
+        StackCardViewModel.GetCard<StackExpanderViewModel>().GetCard<SamplerCardViewModel>();
-            .GetCard<StackExpanderViewModel>()
-            .GetCard<SamplerCardViewModel>();
 
     public bool IsHiresFixEnabled => StackCardViewModel.GetCard<StackExpanderViewModel>().IsEnabled;
 
+    public bool IsUpscaleEnabled => StackCardViewModel.GetCard<StackExpanderViewModel>(1).IsEnabled;
+    
     [JsonIgnore]
     public ProgressViewModel OutputProgress { get; } = new();
 
@@ -87,7 +85,8 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
 
         StackCardViewModel = vmFactory.Get<StackCardViewModel>();
 
-        StackCardViewModel.AddCards(new LoadableViewModelBase[]
+        StackCardViewModel.AddCards(
+            new LoadableViewModelBase[]
             {
                 // Model Card
                 vmFactory.Get<ModelCardViewModel>(),
@@ -97,7 +96,8 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                 vmFactory.Get<StackExpanderViewModel>(stackExpander =>
                 {
                     stackExpander.Title = "Hires Fix";
-                stackExpander.AddCards(new LoadableViewModelBase[]
+                    stackExpander.AddCards(
+                        new LoadableViewModelBase[]
                         {
                             // Hires Fix Upscaler
                             vmFactory.Get<UpscalerCardViewModel>(),
@@ -109,13 +109,25 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                                 samplerCard.IsSamplerSelectionEnabled = false;
                                 samplerCard.IsDenoiseStrengthEnabled = true;
                             })
+                        }
+                    );
+                }),
+                vmFactory.Get<StackExpanderViewModel>(stackExpander =>
+                {
+                    stackExpander.Title = "Upscale";
+                    stackExpander.AddCards(
+                        new LoadableViewModelBase[]
+                        {
+                            // Post processing upscaler
+                            vmFactory.Get<UpscalerCardViewModel>(),
                         });
                 }),
                 // Seed
                 seedCard,
                 // Batch Size
                 vmFactory.Get<BatchSizeCardViewModel>(),
-        });
+            }
+        );
 
         GenerateImageCommand.WithNotificationErrorHandler(notificationService);
     }
@@ -132,18 +144,21 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         var prompt = new NodeDictionary();
         var builder = new ComfyNodeBuilder(prompt);
 
-        var checkpointLoader = prompt.AddNamedNode(new NamedComfyNode("CheckpointLoader")
+        var checkpointLoader = prompt.AddNamedNode(
+            new NamedComfyNode("CheckpointLoader")
             {
                 ClassType = "CheckpointLoaderSimple",
                 Inputs = new Dictionary<string, object?>
                 {
                     ["ckpt_name"] = modelCard.SelectedModelName
                 }
-        });
+            }
+        );
 
         var checkpointVae = checkpointLoader.GetOutput<VAENodeConnection>(2);
 
-        var emptyLatentImage = prompt.AddNamedNode(new NamedComfyNode("EmptyLatentImage")
+        var emptyLatentImage = prompt.AddNamedNode(
+            new NamedComfyNode("EmptyLatentImage")
             {
                 ClassType = "EmptyLatentImage",
                 Inputs = new Dictionary<string, object?>
@@ -152,9 +167,11 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                     ["height"] = samplerCard.Height,
                     ["width"] = samplerCard.Width,
                 }
-        });
+            }
+        );
 
-        var positiveClip = prompt.AddNamedNode(new NamedComfyNode("PositiveCLIP")
+        var positiveClip = prompt.AddNamedNode(
+            new NamedComfyNode("PositiveCLIP")
             {
                 ClassType = "CLIPTextEncode",
                 Inputs = new Dictionary<string, object?>
@@ -162,9 +179,11 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                     ["clip"] = checkpointLoader.GetOutput(1),
                     ["text"] = PromptCardViewModel.PromptDocument.Text,
                 }
-        });
+            }
+        );
 
-        var negativeClip = prompt.AddNamedNode(new NamedComfyNode("NegativeCLIP")
+        var negativeClip = prompt.AddNamedNode(
+            new NamedComfyNode("NegativeCLIP")
             {
                 ClassType = "CLIPTextEncode",
                 Inputs = new Dictionary<string, object?>
@@ -172,32 +191,44 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                     ["clip"] = checkpointLoader.GetOutput(1),
                     ["text"] = PromptCardViewModel.NegativePromptDocument.Text,
                 }
-        });
+            }
+        );
 
-        var sampler = prompt.AddNamedNode(ComfyNodeBuilder.KSampler(
+        var sampler = prompt.AddNamedNode(
+            ComfyNodeBuilder.KSampler(
                 "Sampler",
                 checkpointLoader.GetOutput<ModelNodeConnection>(0),
                 Convert.ToUInt64(seedCard.Seed),
                 samplerCard.Steps,
                 samplerCard.CfgScale,
-            samplerCard.SelectedSampler?.Name ?? throw new InvalidOperationException("Sampler not selected"),
+                samplerCard.SelectedSampler?.Name
+                    ?? throw new InvalidOperationException("Sampler not selected"),
                 "normal",
                 positiveClip.GetOutput<ConditioningNodeConnection>(0),
                 negativeClip.GetOutput<ConditioningNodeConnection>(0),
                 emptyLatentImage.GetOutput<LatentNodeConnection>(0),
-            samplerCard.DenoiseStrength));
+                samplerCard.DenoiseStrength
+            )
+        );
+
+        var lastLatent = sampler.Output;
+        var lastLatentWidth = samplerCard.Width;
+        var lastLatentHeight = samplerCard.Height;
         
-        var vaeDecoder = prompt.AddNamedNode(new NamedComfyNode("VAEDecoder")
+        var vaeDecoder = prompt.AddNamedNode(
+            new NamedComfyNode("VAEDecoder")
             {
                 ClassType = "VAEDecode",
                 Inputs = new Dictionary<string, object?>
                 {
-                ["samples"] = sampler.GetOutput(0),
+                    ["samples"] = lastLatent,
                     ["vae"] = checkpointLoader.GetOutput(2)
                 }
-        });
+            }
+        );
 
-        var saveImage = prompt.AddNamedNode(new NamedComfyNode("SaveImage")
+        var saveImage = prompt.AddNamedNode(
+            new NamedComfyNode("SaveImage")
             {
                 ClassType = "SaveImage",
                 Inputs = new Dictionary<string, object?>
@@ -205,7 +236,8 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
                     ["filename_prefix"] = "SM-Inference",
                     ["images"] = vaeDecoder.GetOutput(0)
                 }
-        });
+            }
+        );
         
         // If hi-res fix is enabled, add the LatentUpscale node and another KSampler node
         if (IsHiresFixEnabled)
@@ -213,58 +245,73 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             var hiresUpscalerCard = UpscalerCardViewModel;
             var hiresSamplerCard = HiresSamplerCardViewModel;
 
-            // Select between latent upscale and normal upscale based on the upscale method
+            // Requested upscale to this size
-            var selectedUpscaler = hiresUpscalerCard.SelectedUpscaler;
+            var hiresWidth = (int)Math.Floor(lastLatentWidth * hiresUpscalerCard.Scale);
+            var hiresHeight = (int)Math.Floor(lastLatentHeight * hiresUpscalerCard.Scale);
             
-            LatentNodeConnection hiresOutput;
+            LatentNodeConnection hiresLatent;
             
-            if (selectedUpscaler?.Type == ComfyUpscalerType.Latent)
+            // Select between latent upscale and normal upscale based on the upscale method
-            {
+            var selectedUpscaler = hiresUpscalerCard.SelectedUpscaler!.Value;
-                hiresOutput = prompt.AddNamedNode(new NamedComfyNode("LatentUpscale")
+            
-                {
+            if (selectedUpscaler.Type == ComfyUpscalerType.None)
-                    ClassType = "LatentUpscale",
-                    Inputs = new Dictionary<string, object?>
-                    {
-                        ["upscale_method"] = hiresUpscalerCard.SelectedUpscaler?.Name,
-                        ["width"] = samplerCard.Width * hiresUpscalerCard.Scale,
-                        ["height"] = samplerCard.Height * hiresUpscalerCard.Scale,
-                        ["crop"] = "disabled",
-                        ["samples"] = sampler.Output
-                    }
-                }).GetOutput<LatentNodeConnection>(0);
-            }
-            else if (selectedUpscaler?.Type == ComfyUpscalerType.ESRGAN)
             {
-                // Convert to image space
+                // If no upscaler selected or none, just reroute the latent image
-                var samplerImage = builder.Lambda_LatentToImage(sampler.Output, checkpointVae);
+                hiresLatent = sampler.Output;
-                // Do group upscale
-                var modelUpscaler = builder.Group_UpscaleWithModel("Upscaler",
-                    selectedUpscaler.Value.Name, samplerImage);
-                // Convert back to latent space
-                hiresOutput = builder.Lambda_ImageToLatent(modelUpscaler.Output, checkpointVae);
             }
             else
             {
-                // If no upscaler selected or none, just reroute the latent image
+                // Otherwise upscale the latent image
-                hiresOutput = sampler.Output;
+                hiresLatent = builder.Group_UpscaleToLatent("HiresFix",
+                    lastLatent, checkpointVae, selectedUpscaler, hiresWidth, hiresHeight).Output;
             }
 
-            var hiresSampler = prompt.AddNamedNode(ComfyNodeBuilder.KSampler(
+            var hiresSampler = prompt.AddNamedNode(
+                ComfyNodeBuilder.KSampler(
                     "HiresSampler",
                     checkpointLoader.GetOutput<ModelNodeConnection>(0),
                     Convert.ToUInt64(seedCard.Seed),
                     hiresSamplerCard.Steps,
                     hiresSamplerCard.CfgScale,
                     // Use hires sampler name if not null, otherwise use the normal sampler name
-                hiresSamplerCard.SelectedSampler?.Name ?? samplerCard.SelectedSampler?.Name ?? throw new InvalidOperationException("Sampler not selected"),
+                    hiresSamplerCard.SelectedSampler?.Name
+                        ?? samplerCard.SelectedSampler?.Name
+                        ?? throw new InvalidOperationException("Sampler not selected"),
                     "normal",
                     positiveClip.GetOutput<ConditioningNodeConnection>(0),
                     negativeClip.GetOutput<ConditioningNodeConnection>(0),
-                hiresOutput,
+                    hiresLatent,
-                hiresSamplerCard.DenoiseStrength));
+                    hiresSamplerCard.DenoiseStrength
+                )
+            );
 
+            // Set as last latent
+            lastLatent = hiresSampler.Output;
+            lastLatentWidth = hiresWidth;
+            lastLatentHeight = hiresHeight;
             // Reroute the VAEDecoder's input to be from the hires sampler
-            vaeDecoder.Inputs["samples"] = hiresSampler.Output;
+            vaeDecoder.Inputs["samples"] = lastLatent;
+        }
+        
+        // If upscale is enabled, add another upscale group
+        if (IsUpscaleEnabled)
+        {
+            var postUpscalerCard = StackCardViewModel.GetCard<StackExpanderViewModel>(1)
+                .GetCard<UpscalerCardViewModel>();
+            
+            var upscaleWidth = (int)Math.Floor(lastLatentWidth * postUpscalerCard.Scale);
+            var upscaleHeight = (int)Math.Floor(lastLatentHeight * postUpscalerCard.Scale);
+            
+            // Build group
+            var postUpscaleGroup = builder.Group_UpscaleToImage("PostUpscale",
+                lastLatent, checkpointVae, postUpscalerCard.SelectedUpscaler!.Value,
+                upscaleWidth, upscaleHeight);
+            
+            // Remove the original vae decoder
+            prompt.Remove(vaeDecoder.Name);
+            
+            // Set as the input for save image
+            saveImage.Inputs["images"] = postUpscaleGroup.Output;
         }
 
         prompt.NormalizeConnectionTypes();
@@ -278,7 +325,8 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
         OutputProgress.Maximum = args.Maximum;
         OutputProgress.IsIndeterminate = false;
 
-        OutputProgress.Text = $"({args.Value} / {args.Maximum})" 
+        OutputProgress.Text =
+            $"({args.Value} / {args.Maximum})"
             + (args.RunningNode != null ? $" {args.RunningNode}" : "");
     }
 
@@ -323,7 +371,9 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             cancellationToken.Register(() =>
             {
                 Logger.Info("Cancelling prompt");
-                client.InterruptPromptAsync(new CancellationTokenSource(5000).Token).SafeFireAndForget();
+                client
+                    .InterruptPromptAsync(new CancellationTokenSource(5000).Token)
+                    .SafeFireAndForget();
             });
 
             try
@@ -353,7 +403,8 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             ImageGalleryCardViewModel.ImageSources.Clear();
 
             var images = imageOutputs[outputNodeNames[0]];
-            if (images is null) return;
+            if (images is null)
+                return;
 
             List<ImageSource> outputImages;
             // Use local file path if available, otherwise use remote URL
@@ -373,8 +424,9 @@ public partial class InferenceTextToImageViewModel : InferenceTabViewModelBase
             // Download all images to make grid, if multiple
             if (outputImages.Count > 1)
             {
-                var loadedImages = outputImages.Select(i =>
+                var loadedImages = outputImages
-                    SKImage.FromEncodedData(i.LocalFile?.Info.OpenRead())).ToImmutableArray();
+                    .Select(i => SKImage.FromEncodedData(i.LocalFile?.Info.OpenRead()))
+                    .ToImmutableArray();
 
                 var grid = ImageProcessor.CreateImageGrid(loadedImages);
 

StabilityMatrix.Avalonia/ViewModels/Inference/StackExpanderViewModel.cs

@@ -1,24 +1,29 @@
 using System.Linq;
 using System.Text.Json.Nodes;
 using CommunityToolkit.Mvvm.ComponentModel;
+using Newtonsoft.Json;
 using StabilityMatrix.Avalonia.Controls;
 using StabilityMatrix.Avalonia.Models.Inference;
 using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Extensions;
+#pragma warning disable CS0657 // Not a valid attribute location for this declaration
 
 namespace StabilityMatrix.Avalonia.ViewModels.Inference;
 
 [View(typeof(StackExpander))]
 public partial class StackExpanderViewModel : StackViewModelBase
 {
-    [ObservableProperty] private string? title;
+    [ObservableProperty]
-    [ObservableProperty] private bool isEnabled;
+    [property: JsonIgnore]
+    private string? title;
+    
+    [ObservableProperty] 
+    private bool isEnabled;
     
     /// <inheritdoc />
     public override void LoadStateFromJsonObject(JsonObject state)
     {
         var model = DeserializeModel<StackExpanderModel>(state);
-        Title = model.Title;
         IsEnabled = model.IsEnabled;
         
         if (model.Cards is null) return;
@@ -37,7 +42,6 @@ public partial class StackExpanderViewModel : StackViewModelBase
     {
         return SerializeModel(new StackExpanderModel
         {
-            Title = Title,
             IsEnabled = IsEnabled,
             Cards = Cards.Select(x => x.SaveStateToJsonObject()).ToList()
         });

StabilityMatrix.Avalonia/ViewModels/SettingsViewModel.cs

@@ -650,8 +650,8 @@ public partial class SettingsViewModel : PageViewModelBase
 
         var imageBox = new ImageViewerDialog()
         {
-            MinWidth = 1500,
+            Width = 1000,
-            MinHeight = 900,
+            Height = 1000,
             DataContext = new ImageViewerViewModel()
             {
                 Image = bitmap
@@ -661,9 +661,10 @@ public partial class SettingsViewModel : PageViewModelBase
         var dialog = new BetterContentDialog
         {
             MaxDialogWidth = 1000,
+            MaxDialogHeight = 1000,
             FullSizeDesired = true,
             Content = imageBox,
-            CloseButtonText = "Close",
+            IsFooterVisible = false,
             ContentVerticalScrollBarVisibility = ScrollBarVisibility.Disabled,
         };
 

StabilityMatrix.Avalonia/Views/Dialogs/ImageViewerDialog.axaml

@@ -8,23 +8,38 @@
                           xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls"
                           xmlns:system="clr-namespace:System;assembly=System.Runtime"
                           xmlns:vmDialogs="clr-namespace:StabilityMatrix.Avalonia.ViewModels.Dialogs"
+                          xmlns:icons="clr-namespace:Projektanker.Icons.Avalonia;assembly=Projektanker.Icons.Avalonia"
+                          VerticalContentAlignment="Stretch"
+                          HorizontalContentAlignment="Stretch"
                           d:DataContext="{x:Static mocks:DesignData.ImageViewerViewModel}"
                           x:DataType="vmDialogs:ImageViewerViewModel"
                           mc:Ignorable="d" d:DesignWidth="800" d:DesignHeight="450"
                           x:Class="StabilityMatrix.Avalonia.Views.Dialogs.ImageViewerDialog">
-    <Grid>
+    <Grid VerticalAlignment="Stretch" HorizontalAlignment="Stretch">
         <controls:AdvancedImageBox
             Name="MainImageBox"
             RenderOptions.BitmapInterpolationMode="None"
-            Image="{Binding Image}"/>
+            Image="{Binding Image}"
+            Source="{Binding Source}"/>
         
         <!-- The preview tracker -->
-        <Image 
+        <!--<Image 
             MinHeight="100"
             MinWidth="100"
             RenderOptions.BitmapInterpolationMode="HighQuality"
             Source="{Binding #MainImageBox.TrackerImage}"
             HorizontalAlignment="Left"
-            VerticalAlignment="Bottom"/> 
+            VerticalAlignment="Bottom"/> -->
+        
+        <!-- Close button -->
+        <Grid 
+            VerticalAlignment="Top"
+            HorizontalAlignment="Right">
+            <Button
+                Margin="0,8,8,0"
+                Classes="transparent"
+                Command="{Binding OnCloseButtonClick}"
+                icons:Attached.Icon="fa-solid fa-xmark"/>
+        </Grid>
     </Grid>
 </controls:UserControlBase>

StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs

@@ -40,7 +40,7 @@ public class ComfyNodeBuilder
         return new NamedComfyNode<ImageNodeConnection>(name)
         {
             ClassType = "VAEDecode",
-            Inputs = new Dictionary<string, object?> { ["latent"] = samples.Data, ["vae"] = vae.Data }
+            Inputs = new Dictionary<string, object?> { ["samples"] = samples.Data, ["vae"] = vae.Data }
         };
     }
     
@@ -105,6 +105,28 @@ public class ComfyNodeBuilder
         };
     }
 
+    public static NamedComfyNode<ImageNodeConnection> ImageScale(
+        string name,
+        ImageNodeConnection image,
+        string method,
+        int height,
+        int width,
+        bool crop)
+    {
+        return new NamedComfyNode<ImageNodeConnection>(name)
+        {
+            ClassType = "ImageScale",
+            Inputs = new Dictionary<string, object?>
+            {
+                ["image"] = image.Data,
+                ["upscale_method"] = method,
+                ["height"] = height,
+                ["width"] = width,
+                ["crop"] = crop ? "center" : "disabled"
+            }
+        };
+    }
+
     public ImageNodeConnection Lambda_LatentToImage(LatentNodeConnection latent, VAENodeConnection vae)
     {
         return nodes.AddNamedNode(VAEDecode($"{GetRandomPrefix()}_VAEDecode", latent, vae)).Output;
@@ -116,7 +138,7 @@ public class ComfyNodeBuilder
     }
     
     /// <summary>
-    /// Create a upscaling node based on a <see cref="ComfyUpscalerType"/>
+    /// Create a group node that upscales a given image with a given model
     /// </summary>
     public NamedComfyNode<ImageNodeConnection> Group_UpscaleWithModel(string name, string modelName, ImageNodeConnection image)
     {
@@ -128,4 +150,146 @@ public class ComfyNodeBuilder
         
         return upscaler;
     }
+
+    /// <summary>
+    /// Create a group node that scales a given image to a given size
+    /// </summary>
+    public NamedComfyNode<LatentNodeConnection> Group_UpscaleToLatent(string name,
+        LatentNodeConnection latent, VAENodeConnection vae,
+        ComfyUpscaler upscaleInfo, int width, int height)
+    {
+        if (upscaleInfo.Type == ComfyUpscalerType.Latent)
+        {
+            return nodes
+                .AddNamedNode(
+                    new NamedComfyNode<LatentNodeConnection>($"{name}_LatentUpscale")
+                    {
+                        ClassType = "LatentUpscale",
+                        Inputs = new Dictionary<string, object?>
+                        {
+                            ["upscale_method"] = upscaleInfo.Name,
+                            ["width"] = width,
+                            ["height"] = height,
+                            ["crop"] = "disabled",
+                            ["samples"] = latent.Data,
+                        }
+                    }
+                );
+        }
+        
+        if (upscaleInfo.Type == ComfyUpscalerType.ESRGAN)
+        {
+            // Convert to image space
+            var samplerImage = nodes.AddNamedNode(
+                VAEDecode(
+                    $"{name}_VAEDecode",
+                    latent,
+                    vae
+                )
+            );
+
+            // Do group upscale
+            var modelUpscaler = Group_UpscaleWithModel(
+                $"{name}_ModelUpscale",
+                upscaleInfo.Name,
+                samplerImage.Output
+            );
+
+            // Since the model upscale is fixed to model (2x/4x), scale it again to the requested size
+            var resizedScaled = nodes.AddNamedNode(
+                ImageScale(
+                    $"{name}_ImageScale",
+                    modelUpscaler.Output,
+                    "bilinear",
+                    height,
+                    width,
+                    false
+                )
+            );
+
+            // Convert back to latent space
+            return nodes
+                .AddNamedNode(
+                    VAEEncode(
+                        $"{name}_VAEEncode",
+                        resizedScaled.Output,
+                        vae
+                    )
+                );
+        }
+        
+        throw new InvalidOperationException($"Unknown upscaler type: {upscaleInfo.Type}");
+    }
+    
+    /// <summary>
+    /// Create a group node that scales a given image to image output
+    /// </summary>
+    public NamedComfyNode<ImageNodeConnection> Group_UpscaleToImage(string name,
+        LatentNodeConnection latent, VAENodeConnection vae,
+        ComfyUpscaler upscaleInfo, int width, int height)
+    {
+        if (upscaleInfo.Type == ComfyUpscalerType.Latent)
+        {
+            var latentUpscale = nodes
+                .AddNamedNode(
+                    new NamedComfyNode<LatentNodeConnection>($"{name}_LatentUpscale")
+                    {
+                        ClassType = "LatentUpscale",
+                        Inputs = new Dictionary<string, object?>
+                        {
+                            ["upscale_method"] = upscaleInfo.Name,
+                            ["width"] = width,
+                            ["height"] = height,
+                            ["crop"] = "disabled",
+                            ["samples"] = latent.Data,
+                        }
+                    }
+                );
+            
+            // Convert to image space
+            return nodes.AddNamedNode(
+                VAEDecode(
+                    $"{name}_VAEDecode",
+                    latentUpscale.Output,
+                    vae
+                )
+            );
+        }
+        
+        if (upscaleInfo.Type == ComfyUpscalerType.ESRGAN)
+        {
+            // Convert to image space
+            var samplerImage = nodes.AddNamedNode(
+                VAEDecode(
+                    $"{name}_VAEDecode",
+                    latent,
+                    vae
+                )
+            );
+
+            // Do group upscale
+            var modelUpscaler = Group_UpscaleWithModel(
+                $"{name}_ModelUpscale",
+                upscaleInfo.Name,
+                samplerImage.Output
+            );
+
+            // Since the model upscale is fixed to model (2x/4x), scale it again to the requested size
+            var resizedScaled = nodes.AddNamedNode(
+                ImageScale(
+                    $"{name}_ImageScale",
+                    modelUpscaler.Output,
+                    "bilinear",
+                    height,
+                    width,
+                    false
+                )
+            );
+
+            // No need to convert back to latent space
+            return resizedScaled;
+        }
+        
+        throw new InvalidOperationException($"Unknown upscaler type: {upscaleInfo.Type}");
+    }
 }