Merge pull request #584 from ionite34/inference-layer-diffuse

Add Inference layer diffuse

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 - Fixed Civitai model browser not showing images when "Show NSFW" is disabled
 - Fixed crash when Installed Workflows page is opened with no Workflows folder
 - Fixed progress bars not displaying properly during package installs & updates
+- Fixed ComfyUI extension updates not running install.py / updating requirements.txt
 
 ## v2.10.0-pre.1
 ### Added

StabilityMatrix.Avalonia/App.axaml

@@ -81,6 +81,7 @@
         <StyleInclude Source="Controls/Inference/FreeUCard.axaml"/>
         <StyleInclude Source="Controls/Inference/ControlNetCard.axaml"/>
         <StyleInclude Source="Controls/Inference/PromptExpansionCard.axaml"/>
+        <StyleInclude Source="Controls/Inference/LayerDiffuseCard.axaml"/>
         <labs:ControlThemes/>
         
         <Style Selector="DockControl">

StabilityMatrix.Avalonia/Controls/Inference/LayerDiffuseCard.axaml

@@ -0,0 +1,48 @@
+<Styles
+    xmlns="https://github.com/avaloniaui"
+    xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
+    xmlns:controls="using:StabilityMatrix.Avalonia.Controls"
+    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:converters="clr-namespace:StabilityMatrix.Avalonia.Converters"
+    xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData"
+    x:DataType="vmInference:LayerDiffuseCardViewModel">
+    <Design.PreviewWith>
+        <Panel Width="400" Height="200">
+            <StackPanel Width="300" VerticalAlignment="Center">
+                <controls:LayerDiffuseCard DataContext="{x:Static mocks:DesignData.LayerDiffuseCardViewModel}"/>
+            </StackPanel>
+        </Panel>
+    </Design.PreviewWith>
+
+    <Style Selector="controls|LayerDiffuseCard">
+        <Setter Property="HorizontalAlignment" Value="Stretch" />
+        <Setter Property="Template">
+            <ControlTemplate>
+                <controls:Card Padding="12">
+                    <sg:SpacedGrid
+                        ColumnDefinitions="Auto,*"
+                        ColumnSpacing="8"
+                        RowDefinitions="*,*,*,*"
+                        RowSpacing="0">
+                        <!--  Mode Selection  -->
+                        <TextBlock
+                            Grid.Column="0"
+                            VerticalAlignment="Center"
+                            Text="Mode"
+                            TextAlignment="Left" />
+
+                        <ui:FAComboBox
+                            Grid.Row="0"
+                            Grid.Column="1"
+                            HorizontalAlignment="Stretch"
+                            DisplayMemberBinding="{Binding Converter={x:Static converters:EnumAttributeConverters.DisplayName}}"
+                            ItemsSource="{Binding AvailableModes}"
+                            SelectedItem="{Binding SelectedMode}" />
+                    </sg:SpacedGrid>
+                </controls:Card>
+            </ControlTemplate>
+        </Setter>
+    </Style>
+</Styles>

StabilityMatrix.Avalonia/Controls/Inference/LayerDiffuseCard.axaml.cs

@@ -0,0 +1,7 @@
+using Avalonia.Controls.Primitives;
+using StabilityMatrix.Core.Attributes;
+
+namespace StabilityMatrix.Avalonia.Controls;
+
+[Transient]
+public class LayerDiffuseCard : TemplatedControl;

StabilityMatrix.Avalonia/DesignData/DesignData.cs

@@ -923,6 +923,9 @@ The gallery images are often inpainted, but you will get something very similar
             vm.IsBatchIndexEnabled = true;
         });
 
+    public static LayerDiffuseCardViewModel LayerDiffuseCardViewModel =>
+        DialogFactory.Get<LayerDiffuseCardViewModel>();
+  
     public static InstalledWorkflowsViewModel InstalledWorkflowsViewModel
     {
         get

StabilityMatrix.Avalonia/Models/Inference/ModuleApplyStepEventArgs.cs

@@ -28,6 +28,16 @@ public class ModuleApplyStepEventArgs : EventArgs
 
     public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = [];
 
+    public List<Action<ModuleApplyStepEventArgs>> PreOutputActions { get; init; } = [];
+
+    public void InvokeAllPreOutputActions()
+    {
+        foreach (var action in PreOutputActions)
+        {
+            action(this);
+        }
+    }
+
     /// <summary>
     /// Creates a new <see cref="ModuleApplyStepEventArgs"/> with the given <see cref="ComfyNodeBuilder"/>.
     /// </summary>

StabilityMatrix.Avalonia/ViewModels/Base/LoadableViewModelBase.cs

@@ -11,6 +11,7 @@ using NLog;
 using StabilityMatrix.Avalonia.Models;
 using StabilityMatrix.Avalonia.ViewModels.Inference;
 using StabilityMatrix.Avalonia.ViewModels.Inference.Modules;
+using StabilityMatrix.Core.Models.Inference;
 
 namespace StabilityMatrix.Avalonia.ViewModels.Base;
 
@@ -20,12 +21,14 @@ namespace StabilityMatrix.Avalonia.ViewModels.Base;
 [JsonDerivedType(typeof(UpscalerCardViewModel), UpscalerCardViewModel.ModuleKey)]
 [JsonDerivedType(typeof(ControlNetCardViewModel), ControlNetCardViewModel.ModuleKey)]
 [JsonDerivedType(typeof(PromptExpansionCardViewModel), PromptExpansionCardViewModel.ModuleKey)]
+[JsonDerivedType(typeof(LayerDiffuseCardViewModel), LayerDiffuseCardViewModel.ModuleKey)]
 [JsonDerivedType(typeof(FreeUModule))]
 [JsonDerivedType(typeof(HiresFixModule))]
 [JsonDerivedType(typeof(UpscalerModule))]
 [JsonDerivedType(typeof(ControlNetModule))]
 [JsonDerivedType(typeof(SaveImageModule))]
 [JsonDerivedType(typeof(PromptExpansionModule))]
+[JsonDerivedType(typeof(LayerDiffuseModule))]
 public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
 {
     private static readonly Logger Logger = LogManager.GetCurrentClassLogger();

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceImageToImageViewModel.cs

@@ -56,26 +56,30 @@ public class InferenceImageToImageViewModel : InferenceTextToImageViewModel
             _ => Convert.ToUInt64(SeedCardViewModel.Seed)
         };
 
-        BatchSizeCardViewModel.ApplyStep(args);
+        var applyArgs = args.ToModuleApplyStepEventArgs();
+
+        BatchSizeCardViewModel.ApplyStep(applyArgs);
 
         // Load models
-        ModelCardViewModel.ApplyStep(args);
+        ModelCardViewModel.ApplyStep(applyArgs);
 
         // Setup image latent source
-        SelectImageCardViewModel.ApplyStep(args);
+        SelectImageCardViewModel.ApplyStep(applyArgs);
 
         // Prompts and loras
-        PromptCardViewModel.ApplyStep(args);
+        PromptCardViewModel.ApplyStep(applyArgs);
 
         // Setup Sampler and Refiner if enabled
-        SamplerCardViewModel.ApplyStep(args);
+        SamplerCardViewModel.ApplyStep(applyArgs);
 
         // Apply module steps
         foreach (var module in ModulesCardViewModel.Cards.OfType<ModuleBase>())
         {
-            module.ApplyStep(args);
+            module.ApplyStep(applyArgs);
         }
 
+        applyArgs.InvokeAllPreOutputActions();
+
         builder.SetupOutputImage();
     }
 

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -131,10 +131,12 @@ public class InferenceTextToImageViewModel : InferenceGenerationViewModelBase, I
             _ => Convert.ToUInt64(SeedCardViewModel.Seed)
         };
 
-        BatchSizeCardViewModel.ApplyStep(args);
+        var applyArgs = args.ToModuleApplyStepEventArgs();
+
+        BatchSizeCardViewModel.ApplyStep(applyArgs);
 
         // Load models
-        ModelCardViewModel.ApplyStep(args);
+        ModelCardViewModel.ApplyStep(applyArgs);
 
         // Setup empty latent
         builder.SetupEmptyLatentSource(
@@ -145,17 +147,19 @@ public class InferenceTextToImageViewModel : InferenceGenerationViewModelBase, I
         );
 
         // Prompts and loras
-        PromptCardViewModel.ApplyStep(args);
+        PromptCardViewModel.ApplyStep(applyArgs);
 
         // Setup Sampler and Refiner if enabled
-        SamplerCardViewModel.ApplyStep(args);
+        SamplerCardViewModel.ApplyStep(applyArgs);
 
         // Hires fix if enabled
         foreach (var module in ModulesCardViewModel.Cards.OfType<ModuleBase>())
         {
-            module.ApplyStep(args);
+            module.ApplyStep(applyArgs);
         }
 
+        applyArgs.InvokeAllPreOutputActions();
+
         builder.SetupOutputImage();
     }
 

StabilityMatrix.Avalonia/ViewModels/Inference/LayerDiffuseCardViewModel.cs

@@ -0,0 +1,100 @@
+using System;
+using System.Collections.Generic;
+using System.ComponentModel.DataAnnotations;
+using System.Linq;
+using CommunityToolkit.Mvvm.ComponentModel;
+using KGySoft.CoreLibraries;
+using StabilityMatrix.Avalonia.Controls;
+using StabilityMatrix.Avalonia.Models.Inference;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
+using StabilityMatrix.Core.Models.Inference;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Inference;
+
+[Transient]
+[ManagedService]
+[View(typeof(LayerDiffuseCard))]
+public partial class LayerDiffuseCardViewModel : LoadableViewModelBase, IComfyStep
+{
+    public const string ModuleKey = "LayerDiffuse";
+
+    [ObservableProperty]
+    private LayerDiffuseMode selectedMode = LayerDiffuseMode.None;
+
+    public IEnumerable<LayerDiffuseMode> AvailableModes => Enum<LayerDiffuseMode>.GetValues();
+
+    [ObservableProperty]
+    [NotifyDataErrorInfo]
+    [Required]
+    [Range(-1d, 3d)]
+    private double weight = 1;
+
+    /// <inheritdoc />
+    public void ApplyStep(ModuleApplyStepEventArgs e)
+    {
+        if (SelectedMode == LayerDiffuseMode.None)
+            return;
+
+        var sdType = SelectedMode switch
+        {
+            LayerDiffuseMode.GenerateForegroundWithTransparencySD15 => "SD15",
+            LayerDiffuseMode.GenerateForegroundWithTransparencySDXL => "SDXL",
+            LayerDiffuseMode.None => throw new ArgumentOutOfRangeException(),
+            _ => throw new ArgumentOutOfRangeException()
+        };
+
+        // Choose config based on mode
+        var config = SelectedMode switch
+        {
+            LayerDiffuseMode.GenerateForegroundWithTransparencySD15
+                => "SD15, Attention Injection, attn_sharing",
+            LayerDiffuseMode.GenerateForegroundWithTransparencySDXL => "SDXL, Conv Injection",
+            LayerDiffuseMode.None => throw new ArgumentOutOfRangeException(),
+            _ => throw new ArgumentOutOfRangeException()
+        };
+
+        foreach (var modelConnections in e.Temp.Models.Values)
+        {
+            var layerDiffuseApply = e.Nodes.AddTypedNode(
+                new ComfyNodeBuilder.LayeredDiffusionApply
+                {
+                    Name = e.Nodes.GetUniqueName($"LayerDiffuseApply_{modelConnections.Name}"),
+                    Model = modelConnections.Model,
+                    Config = config,
+                    Weight = Weight,
+                }
+            );
+
+            modelConnections.Model = layerDiffuseApply.Output;
+        }
+
+        // Add pre output action
+        e.PreOutputActions.Add(applyArgs =>
+        {
+            // Use last latent for decode
+            var latent =
+                applyArgs.Builder.Connections.LastPrimaryLatent
+                ?? throw new InvalidOperationException("Connections.LastPrimaryLatent not set");
+
+            // Convert primary to image if not already
+            var primaryImage = applyArgs.Builder.GetPrimaryAsImage();
+            applyArgs.Builder.Connections.Primary = primaryImage;
+
+            // Add a Layer Diffuse Decode
+            var decode = applyArgs.Nodes.AddTypedNode(
+                new ComfyNodeBuilder.LayeredDiffusionDecodeRgba
+                {
+                    Name = applyArgs.Nodes.GetUniqueName("LayerDiffuseDecode"),
+                    Samples = latent,
+                    Images = primaryImage,
+                    SdVersion = sdType
+                }
+            );
+
+            // Set primary to decode output
+            applyArgs.Builder.Connections.Primary = decode.Output;
+        });
+    }
+}

StabilityMatrix.Avalonia/ViewModels/Inference/Modules/LayerDiffuseModule.cs

@@ -0,0 +1,26 @@
+using StabilityMatrix.Avalonia.Models.Inference;
+using StabilityMatrix.Avalonia.Services;
+using StabilityMatrix.Avalonia.ViewModels.Base;
+using StabilityMatrix.Core.Attributes;
+
+namespace StabilityMatrix.Avalonia.ViewModels.Inference.Modules;
+
+[ManagedService]
+[Transient]
+public class LayerDiffuseModule : ModuleBase
+{
+    /// <inheritdoc />
+    public LayerDiffuseModule(ServiceManager<ViewModelBase> vmFactory)
+        : base(vmFactory)
+    {
+        Title = "Layer Diffuse";
+        AddCards(vmFactory.Get<LayerDiffuseCardViewModel>());
+    }
+
+    /// <inheritdoc />
+    protected override void OnApplyStep(ModuleApplyStepEventArgs e)
+    {
+        var card = GetCard<LayerDiffuseCardViewModel>();
+        card.ApplyStep(e);
+    }
+}

StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs

@@ -109,7 +109,12 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
         ModulesCardViewModel = vmFactory.Get<StackEditableCardViewModel>(modulesCard =>
         {
             modulesCard.Title = Resources.Label_Addons;
-            modulesCard.AvailableModules = [typeof(FreeUModule), typeof(ControlNetModule)];
+            modulesCard.AvailableModules =
+            [
+                typeof(FreeUModule),
+                typeof(ControlNetModule),
+                typeof(LayerDiffuseModule)
+            ];
         });
     }
 

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

@@ -383,6 +383,45 @@ public class ComfyNodeBuilder
         public int BatchSize { get; init; } = 1;
     }
 
+    [TypedNodeOptions(
+        Name = "Inference_Core_LayeredDiffusionApply",
+        RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes >= 0.4.0"]
+    )]
+    public record LayeredDiffusionApply : ComfyTypedNodeBase<ModelNodeConnection>
+    {
+        public required ModelNodeConnection Model { get; init; }
+
+        /// <summary>
+        /// Available configs:
+        /// <para>SD15, Attention Injection, attn_sharing</para>
+        /// <para>SDXL, Conv Injection</para>
+        /// <para>SDXL, Attention Injection</para>
+        /// </summary>
+        public required string Config { get; init; }
+
+        [Range(-1d, 3d)]
+        public double Weight { get; init; } = 1.0;
+    }
+
+    [TypedNodeOptions(
+        Name = "Inference_Core_LayeredDiffusionDecodeRGBA",
+        RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes >= 0.4.0"]
+    )]
+    public record LayeredDiffusionDecodeRgba : ComfyTypedNodeBase<ImageNodeConnection>
+    {
+        public required LatentNodeConnection Samples { get; init; }
+
+        public required ImageNodeConnection Images { get; init; }
+
+        /// <summary>
+        /// Either "SD15" or "SDXL"
+        /// </summary>
+        public required string SdVersion { get; init; }
+
+        [Range(1, 4096)]
+        public int SubBatchSize { get; init; } = 16;
+    }
+
     public ImageNodeConnection Lambda_LatentToImage(LatentNodeConnection latent, VAENodeConnection vae)
     {
         var name = GetUniqueName("VAEDecode");
@@ -867,7 +906,26 @@ public class ComfyNodeBuilder
             set => SamplerTemporaryArgs["Base"] = value;
         }
 
-        public PrimaryNodeConnection? Primary { get; set; }
+        /// <summary>
+        /// The last primary set latent value, updated when <see cref="Primary"/> is set to a latent value.
+        /// </summary>
+        public LatentNodeConnection? LastPrimaryLatent { get; private set; }
+
+        private PrimaryNodeConnection? primary;
+
+        public PrimaryNodeConnection? Primary
+        {
+            get => primary;
+            set
+            {
+                if (value?.IsT0 == true)
+                {
+                    LastPrimaryLatent = value.AsT0;
+                }
+                primary = value;
+            }
+        }
+
         public VAENodeConnection? PrimaryVAE { get; set; }
         public Size PrimarySize { get; set; }
 

StabilityMatrix.Core/Models/Inference/LayerDiffuseMode.cs

@@ -0,0 +1,24 @@
+using System.ComponentModel.DataAnnotations;
+
+namespace StabilityMatrix.Core.Models.Inference;
+
+public enum LayerDiffuseMode
+{
+    /// <summary>
+    /// The layer diffuse mode is not set.
+    /// </summary>
+    [Display(Name = "None")]
+    None,
+
+    /// <summary>
+    /// Generate foreground only with transparency. SD1.5
+    /// </summary>
+    [Display(Name = "(SD 1.5) Generate Foreground with Transparency")]
+    GenerateForegroundWithTransparencySD15,
+
+    /// <summary>
+    /// Generate foreground only with transparency. SDXL
+    /// </summary>
+    [Display(Name = "(SDXL) Generate Foreground with Transparency")]
+    GenerateForegroundWithTransparencySDXL,
+}

StabilityMatrix.Core/Models/Packages/ComfyUI.cs

@@ -511,6 +511,32 @@ public class ComfyUI(
             }
         }
 
+        /// <inheritdoc />
+        public override async Task UpdateExtensionAsync(
+            InstalledPackageExtension installedExtension,
+            InstalledPackage installedPackage,
+            PackageExtensionVersion? version = null,
+            IProgress<ProgressReport>? progress = null,
+            CancellationToken cancellationToken = default
+        )
+        {
+            await base.UpdateExtensionAsync(
+                installedExtension,
+                installedPackage,
+                version,
+                progress,
+                cancellationToken
+            )
+                .ConfigureAwait(false);
+
+            cancellationToken.ThrowIfCancellationRequested();
+
+            var installedDirs = installedExtension.Paths.OfType<DirectoryPath>().Where(dir => dir.Exists);
+
+            await PostInstallAsync(installedPackage, installedDirs, progress, cancellationToken)
+                .ConfigureAwait(false);
+        }
+
         /// <inheritdoc />
         public override async Task InstallExtensionAsync(
             PackageExtension extension,
@@ -539,6 +565,20 @@ public class ComfyUI(
                 .Select(path => cloneRoot.JoinDir(path!))
                 .Where(dir => dir.Exists);
 
+            await PostInstallAsync(installedPackage, installedDirs, progress, cancellationToken)
+                .ConfigureAwait(false);
+        }
+
+        /// <summary>
+        /// Runs post install / update tasks (i.e. install.py, requirements.txt)
+        /// </summary>
+        private async Task PostInstallAsync(
+            InstalledPackage installedPackage,
+            IEnumerable<DirectoryPath> installedDirs,
+            IProgress<ProgressReport>? progress = null,
+            CancellationToken cancellationToken = default
+        )
+        {
             foreach (var installedDir in installedDirs)
             {
                 cancellationToken.ThrowIfCancellationRequested();