Merge pull request #555 from ionite34/inference-reference-controlnet

Add Inference Reference-Only ControlNet

CHANGELOG.md

@@ -6,11 +6,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2.0.0.html).
 
 ## v2.10.0-dev.2
+### Added
+- Added Reference-Only mode for Inference ControlNet, used for guiding the sampler with an image without a pretrained model. Part of the latent and attention layers will be connected to the reference image, similar to Image to Image or Inpainting.
 ### Changed
 - Inference Primary Sampler Addons (i.e. ControlNet, FreeU) are now inherited by Hires Fix Samplers, this can be overriden from the Hires Fix module's settings menu by disabling the "Inherit Primary Sampler Addons" option.
 - Revisited the way images are loaded on the outputs page, with improvements to loading speed & not freezing the UI while loading
 ### Fixed
 - Fixed Outputs page not remembering where the user last was in the TreeView in certain circumstances
+- Fixed Inference extension upgrades not being added to missing extensions list for prompted install
 - Fixed "The Open Web UI button has moved" teaching tip spam
 
 ## v2.10.0-dev.1
@@ -25,12 +28,13 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
 
 ## v2.9.1
 ### Added
-- Fixed [#498](https://github.com/LykosAI/StabilityMatrix/issues/498) - Added "Pony" category to CivitAI Model Browser
+- Fixed [#498](https://github.com/LykosAI/StabilityMatrix/issues/498) Added "Pony" category to CivitAI Model Browser
 ### Fixed
 - Fixed [#502](https://github.com/LykosAI/StabilityMatrix/issues/502) - missing launch options for Forge
 - Fixed [#500](https://github.com/LykosAI/StabilityMatrix/issues/500) - missing output images in Forge when using output sharing
 - Fixed [#490](https://github.com/LykosAI/StabilityMatrix/issues/490) - `mpmath has no attribute 'rational'` error on macOS
 - Fixed incorrect progress text when deleting a checkpoint from the Checkpoints page
+- Fixed incorrect icon colors on macOS
 
 ## v2.9.0
 ### Added

StabilityMatrix.Avalonia/Extensions/ComfyNodeBuilderExtensions.cs

@@ -20,17 +20,15 @@ public static class ComfyNodeBuilderExtensions
         int? batchIndex = null
     )
     {
-        var emptyLatent = builder
-            .Nodes
-            .AddTypedNode(
-                new ComfyNodeBuilder.EmptyLatentImage
-                {
-                    Name = "EmptyLatentImage",
-                    BatchSize = batchSize,
-                    Height = height,
-                    Width = width
-                }
-            );
+        var emptyLatent = builder.Nodes.AddTypedNode(
+            new ComfyNodeBuilder.EmptyLatentImage
+            {
+                Name = "EmptyLatentImage",
+                BatchSize = batchSize,
+                Height = height,
+                Width = width
+            }
+        );
 
         builder.Connections.Primary = emptyLatent.Output;
         builder.Connections.PrimarySize = new Size(width, height);
@@ -39,15 +37,15 @@ public static class ComfyNodeBuilderExtensions
         if (batchIndex is not null)
         {
             builder.Connections.Primary = builder
-                .Nodes
-                .AddNamedNode(
-                    ComfyNodeBuilder.LatentFromBatch(
-                        "LatentFromBatch",
-                        builder.GetPrimaryAsLatent(),
+                .Nodes.AddTypedNode(
+                    new ComfyNodeBuilder.LatentFromBatch
+                    {
+                        Name = "LatentFromBatch",
+                        Samples = builder.GetPrimaryAsLatent(),
                         // remote expects a 0-based index, vm is 1-based
-                        batchIndex.Value - 1,
-                        1
-                    )
+                        BatchIndex = batchIndex.Value - 1,
+                        Length = 1
+                    }
                 )
                 .Output;
         }
@@ -67,9 +65,9 @@ public static class ComfyNodeBuilderExtensions
         var sourceImageRelativePath = Path.Combine("Inference", image.GetHashGuidFileNameCached());
 
         // Load source
-        var loadImage = builder
-            .Nodes
-            .AddTypedNode(new ComfyNodeBuilder.LoadImage { Name = "LoadImage", Image = sourceImageRelativePath });
+        var loadImage = builder.Nodes.AddTypedNode(
+            new ComfyNodeBuilder.LoadImage { Name = "LoadImage", Image = sourceImageRelativePath }
+        );
 
         builder.Connections.Primary = loadImage.Output1;
         builder.Connections.PrimarySize = imageSize;
@@ -78,15 +76,15 @@ public static class ComfyNodeBuilderExtensions
         if (batchIndex is not null)
         {
             builder.Connections.Primary = builder
-                .Nodes
-                .AddNamedNode(
-                    ComfyNodeBuilder.LatentFromBatch(
-                        "LatentFromBatch",
-                        builder.GetPrimaryAsLatent(),
+                .Nodes.AddTypedNode(
+                    new ComfyNodeBuilder.LatentFromBatch
+                    {
+                        Name = "LatentFromBatch",
+                        Samples = builder.GetPrimaryAsLatent(),
                         // remote expects a 0-based index, vm is 1-based
-                        batchIndex.Value - 1,
-                        1
-                    )
+                        BatchIndex = batchIndex.Value - 1,
+                        Length = 1
+                    }
                 )
                 .Output;
         }
@@ -97,25 +95,24 @@ public static class ComfyNodeBuilderExtensions
         if (builder.Connections.Primary is null)
             throw new ArgumentException("No Primary");
 
-        var image = builder
-            .Connections
-            .Primary
-            .Match(
-                _ =>
-                    builder.GetPrimaryAsImage(
-                        builder.Connections.PrimaryVAE
-                            ?? builder.Connections.Refiner.VAE
-                            ?? builder.Connections.Base.VAE
-                            ?? throw new ArgumentException("No Primary, Refiner, or Base VAE")
-                    ),
-                image => image
-            );
+        var image = builder.Connections.Primary.Match(
+            _ =>
+                builder.GetPrimaryAsImage(
+                    builder.Connections.PrimaryVAE
+                        ?? builder.Connections.Refiner.VAE
+                        ?? builder.Connections.Base.VAE
+                        ?? throw new ArgumentException("No Primary, Refiner, or Base VAE")
+                ),
+            image => image
+        );
 
-        var previewImage = builder
-            .Nodes
-            .AddTypedNode(
-                new ComfyNodeBuilder.PreviewImage { Name = builder.Nodes.GetUniqueName("SaveImage"), Images = image }
-            );
+        var previewImage = builder.Nodes.AddTypedNode(
+            new ComfyNodeBuilder.PreviewImage
+            {
+                Name = builder.Nodes.GetUniqueName("SaveImage"),
+                Images = image
+            }
+        );
 
         builder.Connections.OutputNodes.Add(previewImage);
 

StabilityMatrix.Avalonia/Models/Inference/ModuleApplyStepEventArgs.cs

@@ -4,7 +4,7 @@ using System.IO;
 using System.IO.Hashing;
 using System.Text;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
-using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
+using StabilityMatrix.Core.Models.Inference;
 
 namespace StabilityMatrix.Avalonia.Models.Inference;
 
@@ -17,7 +17,7 @@ public class ModuleApplyStepEventArgs : EventArgs
 
     public NodeDictionary Nodes => Builder.Nodes;
 
-    public ModuleApplyStepTemporaryArgs Temp { get; } = new();
+    public ModuleApplyStepTemporaryArgs Temp { get; set; } = new();
 
     /// <summary>
     /// Generation overrides (like hires fix generate, current seed generate, etc.)
@@ -26,6 +26,20 @@ public class ModuleApplyStepEventArgs : EventArgs
 
     public List<(string SourcePath, string DestinationRelativePath)> FilesToTransfer { get; init; } = [];
 
+    /// <summary>
+    /// Creates a new <see cref="ModuleApplyStepEventArgs"/> with the given <see cref="ComfyNodeBuilder"/>.
+    /// </summary>
+    /// <returns></returns>
+    public ModuleApplyStepTemporaryArgs CreateTempFromBuilder()
+    {
+        return new ModuleApplyStepTemporaryArgs
+        {
+            Primary = Builder.Connections.Primary,
+            PrimaryVAE = Builder.Connections.PrimaryVAE,
+            Models = Builder.Connections.Models
+        };
+    }
+
     public void AddFileTransfer(string sourcePath, string destinationRelativePath)
     {
         FilesToTransfer.Add((sourcePath, destinationRelativePath));
@@ -54,22 +68,4 @@ public class ModuleApplyStepEventArgs : EventArgs
 
         return destPath;
     }
-
-    public class ModuleApplyStepTemporaryArgs
-    {
-        /// <summary>
-        /// Temporary conditioning apply step, used by samplers to apply control net.
-        /// </summary>
-        public ConditioningConnections? Conditioning { get; set; }
-
-        /// <summary>
-        /// Temporary refiner conditioning apply step, used by samplers to apply control net.
-        /// </summary>
-        public ConditioningConnections? RefinerConditioning { get; set; }
-
-        /// <summary>
-        /// Temporary model apply step, used by samplers to apply control net.
-        /// </summary>
-        public ModelNodeConnection? Model { get; set; }
-    }
 }

StabilityMatrix.Avalonia/StabilityMatrix.Avalonia.csproj

@@ -65,8 +65,8 @@
         <PackageReference Include="Exceptionless.DateTimeExtensions" Version="3.4.3" />
         <PackageReference Include="FluentAvalonia.BreadcrumbBar" Version="2.0.2" />
         <PackageReference Include="FluentAvaloniaUI" Version="2.0.5" />
-        <PackageReference Include="FluentIcons.Avalonia" Version="1.1.230" />
-        <PackageReference Include="FluentIcons.Avalonia.Fluent" Version="1.1.230" />
+        <PackageReference Include="FluentIcons.Avalonia" Version="1.1.228" />
+        <PackageReference Include="FluentIcons.Avalonia.Fluent" Version="1.1.228" />
         <PackageReference Include="FuzzySharp" Version="2.0.2" />
         <PackageReference Include="JetBrains.Annotations" Version="2023.3.0" />
         <PackageReference Include="Markdown.Avalonia" Version="11.0.2" />

StabilityMatrix.Avalonia/ViewModels/Base/InferenceGenerationViewModelBase.cs

@@ -16,7 +16,6 @@ using Avalonia.Threading;
 using CommunityToolkit.Mvvm.Input;
 using ExifLibrary;
 using FluentAvalonia.UI.Controls;
-using Microsoft.Extensions.DependencyInjection;
 using NLog;
 using Refit;
 using Semver;
@@ -760,6 +759,7 @@ public abstract partial class InferenceGenerationViewModelBase
 
             var steps = new List<IPackageStep>();
 
+            // Add install for missing extensions
             foreach (var missingExtension in missingExtensions)
             {
                 if (!manifestExtensionsMap.TryGetValue(missingExtension.Name, out var extension))
@@ -774,6 +774,18 @@ public abstract partial class InferenceGenerationViewModelBase
                 steps.Add(new InstallExtensionStep(manager, localPackagePair.InstalledPackage, extension));
             }
 
+            // Add update for out of date extensions
+            foreach (var (specifier, installed) in outOfDateExtensions)
+            {
+                if (!manifestExtensionsMap.TryGetValue(specifier.Name, out var extension))
+                {
+                    Logger.Warn("Extension {MissingExtensionUrl} not found in manifests", specifier.Name);
+                    continue;
+                }
+
+                steps.Add(new UpdateExtensionStep(manager, localPackagePair.InstalledPackage, installed));
+            }
+
             var runner = new PackageModificationRunner
             {
                 ShowDialogOnStart = true,

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

@@ -6,6 +6,8 @@ using StabilityMatrix.Avalonia.Models.Inference;
 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.Api.Comfy;
 using StabilityMatrix.Core.Models.Api.Comfy.Nodes;
 
@@ -66,6 +68,84 @@ public class ControlNetModule : ModuleBase
             image = aioPreprocessor.Output;
         }
 
+        // If ReferenceOnly is selected, use special node
+        if (card.SelectedModel == RemoteModels.ControlNetReferenceOnlyModel)
+        {
+            // We need to rescale image to be the current primary size if it's not already
+            var originalPrimary = e.Temp.Primary!.Unwrap();
+            var originalPrimarySize = e.Builder.Connections.PrimarySize;
+
+            if (card.SelectImageCardViewModel.CurrentBitmapSize != originalPrimarySize)
+            {
+                var scaled = e.Builder.Group_Upscale(
+                    e.Nodes.GetUniqueName("ControlNet_Rescale"),
+                    image,
+                    e.Temp.GetDefaultVAE(),
+                    ComfyUpscaler.NearestExact,
+                    originalPrimarySize.Width,
+                    originalPrimarySize.Height
+                );
+                e.Temp.Primary = scaled;
+            }
+            else
+            {
+                e.Temp.Primary = image;
+            }
+
+            // Set image as new latent source, add reference only node
+            var model = e.Temp.GetRefinerOrBaseModel();
+            var controlNetReferenceOnly = e.Nodes.AddTypedNode(
+                new ComfyNodeBuilder.ReferenceOnlySimple
+                {
+                    Name = e.Nodes.GetUniqueName("ControlNet_ReferenceOnly"),
+                    Reference = e.Builder.GetPrimaryAsLatent(
+                        e.Temp.Primary,
+                        e.Builder.Connections.GetDefaultVAE()
+                    ),
+                    Model = model
+                }
+            );
+
+            var referenceOnlyModel = controlNetReferenceOnly.Output1;
+
+            // If ControlNet strength is not 1, add Model Merge
+            if (Math.Abs(card.Strength - 1) > 0.01)
+            {
+                var modelBlend = e.Nodes.AddTypedNode(
+                    new ComfyNodeBuilder.ModelMergeSimple
+                    {
+                        Name = e.Nodes.GetUniqueName("ControlNet_ReferenceOnly_ModelMerge"),
+                        Model1 = referenceOnlyModel,
+                        Model2 = e.Temp.GetRefinerOrBaseModel(),
+                        // Where 0 is full reference only, 1 is full original
+                        Ratio = 1 - card.Strength
+                    }
+                );
+
+                referenceOnlyModel = modelBlend.Output;
+            }
+
+            // Set output as new primary and model source
+            if (model == e.Temp.Refiner.Model)
+            {
+                e.Temp.Refiner.Model = referenceOnlyModel;
+            }
+            else
+            {
+                e.Temp.Base.Model = referenceOnlyModel;
+            }
+            e.Temp.Primary = controlNetReferenceOnly.Output2;
+
+            // Indicate that the Primary latent has been temp batched
+            // https://github.com/comfyanonymous/ComfyUI_experiments/issues/11
+
+            e.Temp.IsPrimaryTempBatched = true;
+            // Index 0 is the original image, index 1 is the reference only latent
+            e.Temp.PrimaryTempBatchPickIndex = 1;
+
+            return;
+        }
+
         var controlNetLoader = e.Nodes.AddTypedNode(
             new ComfyNodeBuilder.ControlNetLoader
             {
@@ -81,18 +161,18 @@ public class ControlNetModule : ModuleBase
                 Name = e.Nodes.GetUniqueName("ControlNetApply"),
                 Image = image,
                 ControlNet = controlNetLoader.Output,
-                Positive = e.Temp.Conditioning?.Positive ?? throw new ArgumentException("No Conditioning"),
-                Negative = e.Temp.Conditioning?.Negative ?? throw new ArgumentException("No Conditioning"),
+                Positive = e.Temp.Base.Conditioning!.Unwrap().Positive,
+                Negative = e.Temp.Base.Conditioning.Negative,
                 Strength = card.Strength,
                 StartPercent = card.StartPercent,
                 EndPercent = card.EndPercent,
             }
         );
 
-        e.Temp.Conditioning = (controlNetApply.Output1, controlNetApply.Output2);
+        e.Temp.Base.Conditioning = (controlNetApply.Output1, controlNetApply.Output2);
 
         // Refiner if available
-        if (e.Temp.RefinerConditioning is not null)
+        if (e.Temp.Refiner.Conditioning is not null)
         {
             var controlNetRefinerApply = e.Nodes.AddTypedNode(
                 new ComfyNodeBuilder.ControlNetApplyAdvanced
@@ -100,15 +180,15 @@ public class ControlNetModule : ModuleBase
                     Name = e.Nodes.GetUniqueName("Refiner_ControlNetApply"),
                     Image = image,
                     ControlNet = controlNetLoader.Output,
-                    Positive = e.Temp.RefinerConditioning.Positive,
-                    Negative = e.Temp.RefinerConditioning.Negative,
+                    Positive = e.Temp.Refiner.Conditioning!.Unwrap().Positive,
+                    Negative = e.Temp.Refiner.Conditioning.Negative,
                     Strength = card.Strength,
                     StartPercent = card.StartPercent,
                     EndPercent = card.EndPercent,
                 }
             );
 
-            e.Temp.RefinerConditioning = (controlNetRefinerApply.Output1, controlNetRefinerApply.Output2);
+            e.Temp.Refiner.Conditioning = (controlNetRefinerApply.Output1, controlNetRefinerApply.Output2);
         }
     }
 }

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

@@ -78,10 +78,16 @@ public partial class HiresFixModule : ModuleBase
             );
         }
 
-        // Choose conditioning based on whether to inherit primary sampler addons
-        var conditioning = samplerCard.InheritPrimarySamplerAddons
-            ? builder.Connections.GetRefinerOrBasePrimarySamplerConditioning()
-            : builder.Connections.GetRefinerOrBaseConditioning();
+        // If we need to inherit primary sampler addons, use their temp args
+        if (samplerCard.InheritPrimarySamplerAddons)
+        {
+            e.Temp = e.Builder.Connections.BaseSamplerTemporaryArgs ?? e.CreateTempFromBuilder();
+        }
+        else
+        {
+            // otherwise just use new ones
+            e.Temp = e.CreateTempFromBuilder();
+        }
 
         var hiresSampler = builder.Nodes.AddTypedNode(
             new ComfyNodeBuilder.KSampler
@@ -99,8 +105,8 @@ public partial class HiresFixModule : ModuleBase
                     samplerCard.SelectedScheduler?.Name
                     ?? e.Builder.Connections.PrimaryScheduler?.Name
                     ?? throw new ArgumentException("No PrimaryScheduler"),
-                Positive = conditioning.Positive,
-                Negative = conditioning.Negative,
+                Positive = e.Temp.GetRefinerOrBaseConditioning().Positive,
+                Negative = e.Temp.GetRefinerOrBaseConditioning().Negative,
                 LatentImage = builder.GetPrimaryAsLatent(),
                 Denoise = samplerCard.DenoiseStrength
             }

StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs

@@ -135,8 +135,7 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
         }
 
         // Provide temp values
-        e.Temp.Conditioning = e.Builder.Connections.Base.Conditioning;
-        e.Temp.RefinerConditioning = e.Builder.Connections.Refiner.Conditioning;
+        e.Temp = e.CreateTempFromBuilder();
 
         // Apply steps from our addons
         ApplyAddonSteps(e);
@@ -148,9 +147,8 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
         {
             ApplyStepsInitialSampler(e);
 
-            // Save temp conditioning for primary
-            e.Builder.Connections.Base.PrimarySamplerConditioning = e.Temp.Conditioning;
-            e.Builder.Connections.Refiner.PrimarySamplerConditioning = e.Temp.RefinerConditioning;
+            // Save temp
+            e.Builder.Connections.BaseSamplerTemporaryArgs = e.Temp;
         }
         else
         {
@@ -164,7 +162,10 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
     private void ApplyStepsInitialSampler(ModuleApplyStepEventArgs e)
     {
         // Get primary as latent using vae
-        var primaryLatent = e.Builder.GetPrimaryAsLatent();
+        var primaryLatent = e.Builder.GetPrimaryAsLatent(
+            e.Temp.Primary!.Unwrap(),
+            e.Builder.Connections.GetDefaultVAE()
+        );
 
         // Set primary sampler and scheduler
         var primarySampler = SelectedSampler ?? throw new ValidationException("Sampler not selected");
@@ -174,8 +175,8 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
         e.Builder.Connections.PrimaryScheduler = primaryScheduler;
 
         // Use Temp Conditioning that may be modified by addons
-        var conditioning = e.Temp.Conditioning.Unwrap();
-        var refinerConditioning = e.Temp.RefinerConditioning;
+        var conditioning = e.Temp.Base.Conditioning.Unwrap();
+        var refinerConditioning = e.Temp.Refiner.Conditioning;
 
         // Use custom sampler if SDTurbo scheduler is selected
         if (e.Builder.Connections.PrimaryScheduler == ComfyScheduler.SDTurbo)
@@ -221,21 +222,16 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
             );
 
             e.Builder.Connections.Primary = sampler.Output1;
-
-            return;
         }
-
         // Use KSampler if no refiner, otherwise need KSamplerAdvanced
-        if (e.Builder.Connections.Refiner.Model is null)
+        else if (e.Builder.Connections.Refiner.Model is null)
         {
-            var baseConditioning = e.Builder.Connections.Base.Conditioning.Unwrap();
-
             // No refiner
             var sampler = e.Nodes.AddTypedNode(
                 new ComfyNodeBuilder.KSampler
                 {
                     Name = "Sampler",
-                    Model = e.Builder.Connections.Base.Model.Unwrap(),
+                    Model = e.Temp.Base.Model!.Unwrap(),
                     Seed = e.Builder.Connections.Seed,
                     SamplerName = primarySampler.Name,
                     Scheduler = primaryScheduler.Name,
@@ -257,7 +253,7 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
                 new ComfyNodeBuilder.KSamplerAdvanced
                 {
                     Name = "Sampler",
-                    Model = e.Builder.Connections.Base.Model.Unwrap(),
+                    Model = e.Temp.Base.Model!.Unwrap(),
                     AddNoise = true,
                     NoiseSeed = e.Builder.Connections.Seed,
                     Steps = TotalSteps,
@@ -273,8 +269,30 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
                 }
             );
 
+            e.Builder.Connections.Primary = sampler.Output;
+        }
+
+        // If temp batched, add a LatentFromBatch to pick the temp batch right after first sampler
+        if (e.Temp.IsPrimaryTempBatched)
+        {
+            e.Builder.Connections.Primary = e.Nodes.AddTypedNode(
+                new ComfyNodeBuilder.LatentFromBatch
+                {
+                    Name = e.Nodes.GetUniqueName("ControlNet_LatentFromBatch"),
+                    Samples = e.Builder.GetPrimaryAsLatent(),
+                    BatchIndex = e.Temp.PrimaryTempBatchPickIndex,
+                    // Use max length here as recommended
+                    // https://github.com/comfyanonymous/ComfyUI_experiments/issues/11
+                    Length = 64
+                }
+            ).Output;
+        }
+
+        // Refiner
+        if (e.Builder.Connections.Refiner.Model is not null)
+        {
             // Add refiner sampler
-            var refinerSampler = e.Nodes.AddTypedNode(
+            e.Builder.Connections.Primary = e.Nodes.AddTypedNode(
                 new ComfyNodeBuilder.KSamplerAdvanced
                 {
                     Name = "Sampler_Refiner",
@@ -288,14 +306,12 @@ public partial class SamplerCardViewModel : LoadableViewModelBase, IParametersLo
                     Positive = refinerConditioning!.Positive,
                     Negative = refinerConditioning.Negative,
                     // Connect to previous sampler
-                    LatentImage = sampler.Output,
+                    LatentImage = e.Builder.GetPrimaryAsLatent(),
                     StartAtStep = Steps,
                     EndAtStep = TotalSteps,
                     ReturnWithLeftoverNoise = false
                 }
-            );
-
-            e.Builder.Connections.Primary = refinerSampler.Output;
+            ).Output;
         }
     }
 

StabilityMatrix.Core/Extensions/NullableExtensions.cs

@@ -1,7 +1,7 @@
 using System.ComponentModel;
 using System.Diagnostics;
-using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using JetBrains.Annotations;
 
 namespace StabilityMatrix.Core.Extensions;
 
@@ -16,7 +16,12 @@ public static class NullableExtensions
     [DebuggerStepThrough]
     [EditorBrowsable(EditorBrowsableState.Never)]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static T Unwrap<T>([NotNull] this T? obj, [CallerArgumentExpression("obj")] string? paramName = null)
+    [ContractAnnotation("obj:null => halt; obj:notnull => notnull")]
+    [return: System.Diagnostics.CodeAnalysis.NotNull]
+    public static T Unwrap<T>(
+        [System.Diagnostics.CodeAnalysis.NotNull] this T? obj,
+        [CallerArgumentExpression("obj")] string? paramName = null
+    )
         where T : class
     {
         if (obj is null)
@@ -35,7 +40,11 @@ public static class NullableExtensions
     [DebuggerStepThrough]
     [EditorBrowsable(EditorBrowsableState.Never)]
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static T Unwrap<T>([NotNull] this T? obj, [CallerArgumentExpression("obj")] string? paramName = null)
+    [ContractAnnotation("obj:null => halt")]
+    public static T Unwrap<T>(
+        [System.Diagnostics.CodeAnalysis.NotNull] this T? obj,
+        [CallerArgumentExpression("obj")] string? paramName = null
+    )
         where T : struct
     {
         if (obj is null)

StabilityMatrix.Core/Helper/RemoteModels.cs

@@ -167,8 +167,14 @@ public static class RemoteModels
             )
         };
 
+    public static HybridModelFile ControlNetReferenceOnlyModel { get; } =
+        HybridModelFile.FromRemote("@ReferenceOnly");
+
     public static IReadOnlyList<HybridModelFile> ControlNetModels { get; } =
-        ControlNets.Select(HybridModelFile.FromDownloadable).ToImmutableArray();
+        ControlNets
+            .Select(HybridModelFile.FromDownloadable)
+            .Concat([ControlNetReferenceOnlyModel])
+            .ToImmutableArray();
 
     private static IEnumerable<RemoteResource> PromptExpansions =>
         [

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

@@ -7,6 +7,7 @@ using StabilityMatrix.Core.Attributes;
 using StabilityMatrix.Core.Extensions;
 using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
 using StabilityMatrix.Core.Models.Database;
+using StabilityMatrix.Core.Models.Inference;
 
 namespace StabilityMatrix.Core.Models.Api.Comfy.Nodes;
 
@@ -131,23 +132,35 @@ public class ComfyNodeBuilder
         public int StopAtClipLayer { get; init; } = -1;
     }
 
-    public static NamedComfyNode<LatentNodeConnection> LatentFromBatch(
-        string name,
-        LatentNodeConnection samples,
-        int batchIndex,
-        int length
-    )
+    public record LatentFromBatch : ComfyTypedNodeBase<LatentNodeConnection>
     {
-        return new NamedComfyNode<LatentNodeConnection>(name)
-        {
-            ClassType = "LatentFromBatch",
-            Inputs = new Dictionary<string, object?>
-            {
-                ["samples"] = samples.Data,
-                ["batch_index"] = batchIndex,
-                ["length"] = length,
-            }
-        };
+        public required LatentNodeConnection Samples { get; init; }
+
+        [Range(0, 63)]
+        public int BatchIndex { get; init; } = 0;
+
+        [Range(1, 64)]
+        public int Length { get; init; } = 1;
+    }
+
+    public record LatentBlend : ComfyTypedNodeBase<LatentNodeConnection>
+    {
+        public required LatentNodeConnection Samples1 { get; init; }
+
+        public required LatentNodeConnection Samples2 { get; init; }
+
+        [Range(0d, 1d)]
+        public double BlendFactor { get; init; } = 0.5;
+    }
+
+    public record ModelMergeSimple : ComfyTypedNodeBase<ModelNodeConnection>
+    {
+        public required ModelNodeConnection Model1 { get; init; }
+
+        public required ModelNodeConnection Model2 { get; init; }
+
+        [Range(0d, 1d)]
+        public double Ratio { get; init; } = 1;
     }
 
     public static NamedComfyNode<ImageNodeConnection> ImageUpscaleWithModel(
@@ -356,6 +369,20 @@ public class ComfyNodeBuilder
         public int Resolution { get; init; } = 512;
     }
 
+    [TypedNodeOptions(
+        Name = "Inference_Core_ReferenceOnlySimple",
+        RequiredExtensions = ["https://github.com/LykosAI/ComfyUI-Inference-Core-Nodes >= 0.3.0"]
+    )]
+    public record ReferenceOnlySimple : ComfyTypedNodeBase<ModelNodeConnection, LatentNodeConnection>
+    {
+        public required ModelNodeConnection Model { get; init; }
+
+        public required LatentNodeConnection Reference { get; init; }
+
+        [Range(1, 64)]
+        public int BatchSize { get; init; } = 1;
+    }
+
     public ImageNodeConnection Lambda_LatentToImage(LatentNodeConnection latent, VAENodeConnection vae)
     {
         var name = GetUniqueName("VAEDecode");
@@ -832,6 +859,14 @@ public class ComfyNodeBuilder
         public ModelConnections Base => Models["Base"];
         public ModelConnections Refiner => Models["Refiner"];
 
+        public Dictionary<string, ModuleApplyStepTemporaryArgs?> SamplerTemporaryArgs { get; } = new();
+
+        public ModuleApplyStepTemporaryArgs? BaseSamplerTemporaryArgs
+        {
+            get => SamplerTemporaryArgs.GetValueOrDefault("Base");
+            set => SamplerTemporaryArgs["Base"] = value;
+        }
+
         public PrimaryNodeConnection? Primary { get; set; }
         public VAENodeConnection? PrimaryVAE { get; set; }
         public Size PrimarySize { get; set; }
@@ -857,13 +892,6 @@ public class ComfyNodeBuilder
                 ?? throw new NullReferenceException("No Refiner or Base Conditioning");
         }
 
-        public ConditioningConnections GetRefinerOrBasePrimarySamplerConditioning()
-        {
-            return Refiner.PrimarySamplerConditioning
-                ?? Base.PrimarySamplerConditioning
-                ?? throw new NullReferenceException("No Refiner or Base PrimarySampler Conditioning");
-        }
-
         public VAENodeConnection GetDefaultVAE()
         {
             return PrimaryVAE ?? Refiner.VAE ?? Base.VAE ?? throw new NullReferenceException("No VAE");

StabilityMatrix.Core/Models/HybridModelFile.cs

@@ -1,5 +1,6 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json.Serialization;
+using StabilityMatrix.Core.Helper;
 using StabilityMatrix.Core.Models.Database;
 
 namespace StabilityMatrix.Core.Models;
@@ -67,6 +68,11 @@ public record HybridModelFile
                 return "Default";
             }
 
+            if (ReferenceEquals(this, RemoteModels.ControlNetReferenceOnlyModel))
+            {
+                return "Reference Only";
+            }
+
             var fileName = Path.GetFileNameWithoutExtension(RelativePath);
 
             if (

StabilityMatrix.Core/Models/Inference/ModuleApplyStepTemporaryArgs.cs

@@ -0,0 +1,46 @@
+using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes;
+
+namespace StabilityMatrix.Core.Models.Inference;
+
+public class ModuleApplyStepTemporaryArgs
+{
+    /// <summary>
+    /// Temporary Primary apply step, used by ControlNet ReferenceOnly which changes the latent.
+    /// </summary>
+    public PrimaryNodeConnection? Primary { get; set; }
+
+    public VAENodeConnection? PrimaryVAE { get; set; }
+
+    /// <summary>
+    /// Used by Reference-Only ControlNet to indicate that <see cref="Primary"/> has been batched.
+    /// </summary>
+    public bool IsPrimaryTempBatched { get; set; }
+
+    /// <summary>
+    /// When <see cref="IsPrimaryTempBatched"/> is true, this is the index of the temp batch to pick after sampling.
+    /// </summary>
+    public int PrimaryTempBatchPickIndex { get; set; }
+
+    public Dictionary<string, ModelConnections> Models { get; set; } =
+        new() { ["Base"] = new ModelConnections("Base"), ["Refiner"] = new ModelConnections("Refiner") };
+
+    public ModelConnections Base => Models["Base"];
+    public ModelConnections Refiner => Models["Refiner"];
+
+    public ConditioningConnections GetRefinerOrBaseConditioning()
+    {
+        return Refiner.Conditioning
+            ?? Base.Conditioning
+            ?? throw new NullReferenceException("No Refiner or Base Conditioning");
+    }
+
+    public ModelNodeConnection GetRefinerOrBaseModel()
+    {
+        return Refiner.Model ?? Base.Model ?? throw new NullReferenceException("No Refiner or Base Model");
+    }
+
+    public VAENodeConnection GetDefaultVAE()
+    {
+        return PrimaryVAE ?? Refiner.VAE ?? Base.VAE ?? throw new NullReferenceException("No VAE");
+    }
+}