Add FluxGuidanceModule for using flux fp8 model with regular TextToImage inference project

Author: mohnjiles

StabilityMatrix.Avalonia/ViewModels/Base/LoadableViewModelBase.cs

@@ -32,6 +32,7 @@ namespace StabilityMatrix.Avalonia.ViewModels.Base;
 [JsonDerivedType(typeof(LoraModule))]
 [JsonDerivedType(typeof(LayerDiffuseModule))]
 [JsonDerivedType(typeof(FaceDetailerModule))]
+[JsonDerivedType(typeof(FluxGuidanceModule))]
 public abstract class LoadableViewModelBase : ViewModelBase, IJsonLoadableState
 {
     private static readonly Logger Logger = LogManager.GetCurrentClassLogger();

StabilityMatrix.Avalonia/ViewModels/CheckpointsPageViewModel.cs

@@ -318,6 +318,9 @@ or nameof(SortConnectedModelsFirst)
             )
             .Subscribe();
 
+        // make filter go
+        OnPropertyChanged(nameof(HideEmptyRootCategories));
+
         settingsManager.RelayPropertyFor(
             this,
             vm => vm.IsImportAsConnectedEnabled,
@@ -380,8 +383,6 @@ or nameof(SortConnectedModelsFirst)
 
         // make sure a sort happens
         OnPropertyChanged(nameof(SortConnectedModelsFirst));
-        // gotta do this one too i guess
-        OnPropertyChanged(nameof(HideEmptyRootCategories));
     }
 
     public void ClearSearchQuery()
@@ -852,7 +853,8 @@ private void RefreshCategories()
         SelectedCategory =
             previouslySelectedCategory
             ?? Categories.FirstOrDefault(x => x.Path == previouslySelectedCategory?.Path)
-            ?? Categories.First();
+            ?? Categories.FirstOrDefault()
+            ?? categoriesCache.Items[0];
 
         var dirPath = new DirectoryPath(SelectedCategory.Path);
 
@@ -868,12 +870,6 @@ private void RefreshCategories()
 
             dirPath = dirPath.Parent;
         }
-
-        // one more time for good measure??
-        SelectedCategory =
-            previouslySelectedCategory
-            ?? Categories.FirstOrDefault(x => x.Path == previouslySelectedCategory?.Path)
-            ?? Categories.First();
     }
 
     private ObservableCollection<CheckpointCategory> GetSubfolders(string strPath)

StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs

@@ -142,13 +142,26 @@ protected override void BuildPrompt(BuildPromptEventArgs args)
         // Load models
         ModelCardViewModel.ApplyStep(applyArgs);
 
-        // Setup empty latent
-        builder.SetupEmptyLatentSource(
-            SamplerCardViewModel.Width,
-            SamplerCardViewModel.Height,
-            BatchSizeCardViewModel.BatchSize,
-            BatchSizeCardViewModel.IsBatchIndexEnabled ? BatchSizeCardViewModel.BatchIndex : null
-        );
+        if (SamplerCardViewModel.ModulesCardViewModel.IsModuleEnabled<FluxGuidanceModule>())
+        {
+            // need SD3Latent
+            builder.SetupEmptySd3LatentSource(
+                SamplerCardViewModel.Width,
+                SamplerCardViewModel.Height,
+                BatchSizeCardViewModel.BatchSize,
+                BatchSizeCardViewModel.IsBatchIndexEnabled ? BatchSizeCardViewModel.BatchIndex : null
+            );
+        }
+        else
+        {
+            // Setup empty latent
+            builder.SetupEmptyLatentSource(
+                SamplerCardViewModel.Width,
+                SamplerCardViewModel.Height,
+                BatchSizeCardViewModel.BatchSize,
+                BatchSizeCardViewModel.IsBatchIndexEnabled ? BatchSizeCardViewModel.BatchIndex : null
+            );
+        }
 
         // Prompts and loras
         PromptCardViewModel.ApplyStep(applyArgs);

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

@@ -0,0 +1,19 @@
+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 FluxGuidanceModule : ModuleBase
+{
+    public FluxGuidanceModule(ServiceManager<ViewModelBase> vmFactory)
+        : base(vmFactory)
+    {
+        Title = "Use Flux Guidance";
+    }
+
+    protected override void OnApplyStep(ModuleApplyStepEventArgs e) { }
+}

StabilityMatrix.Avalonia/ViewModels/Inference/SamplerCardViewModel.cs

@@ -117,7 +117,8 @@ ServiceManager<ViewModelBase> vmFactory
             [
                 typeof(FreeUModule),
                 typeof(ControlNetModule),
-                typeof(LayerDiffuseModule)
+                typeof(LayerDiffuseModule),
+                typeof(FluxGuidanceModule)
             ];
         });
     }
@@ -286,6 +287,23 @@ private void ApplyStepsInitialSampler(ModuleApplyStepEventArgs e)
         var conditioning = e.Temp.Base.Conditioning.Unwrap();
         var refinerConditioning = e.Temp.Refiner.Conditioning;
 
+        var useFluxGuidance = ModulesCardViewModel.IsModuleEnabled<FluxGuidanceModule>();
+
+        if (useFluxGuidance)
+        {
+            // Flux guidance
+            var fluxGuidance = e.Nodes.AddTypedNode(
+                new ComfyNodeBuilder.FluxGuidance
+                {
+                    Name = e.Nodes.GetUniqueName("FluxGuidance"),
+                    Conditioning = conditioning.Positive,
+                    Guidance = CfgScale
+                }
+            );
+
+            conditioning = conditioning with { Positive = fluxGuidance.Output };
+        }
+
         // Use custom sampler if SDTurbo scheduler is selected
         if (e.Builder.Connections.PrimaryScheduler == ComfyScheduler.SDTurbo)
         {
@@ -320,7 +338,7 @@ private void ApplyStepsInitialSampler(ModuleApplyStepEventArgs e)
                     Model = e.Builder.Connections.Base.Model,
                     AddNoise = true,
                     NoiseSeed = e.Builder.Connections.Seed,
-                    Cfg = CfgScale,
+                    Cfg = useFluxGuidance ? 1.0d : CfgScale,
                     Positive = conditioning.Positive,
                     Negative = conditioning.Negative,
                     Sampler = kSamplerSelect.Output,
@@ -344,7 +362,7 @@ private void ApplyStepsInitialSampler(ModuleApplyStepEventArgs e)
                     SamplerName = primarySampler.Name,
                     Scheduler = primaryScheduler.Name,
                     Steps = Steps,
-                    Cfg = CfgScale,
+                    Cfg = useFluxGuidance ? 1.0d : CfgScale,
                     Positive = conditioning.Positive,
                     Negative = conditioning.Negative,
                     LatentImage = primaryLatent,
@@ -365,7 +383,7 @@ private void ApplyStepsInitialSampler(ModuleApplyStepEventArgs e)
                     AddNoise = true,
                     NoiseSeed = e.Builder.Connections.Seed,
                     Steps = TotalSteps,
-                    Cfg = CfgScale,
+                    Cfg = useFluxGuidance ? 1.0d : CfgScale,
                     SamplerName = primarySampler.Name,
                     Scheduler = primaryScheduler.Name,
                     Positive = conditioning.Positive,

StabilityMatrix.Avalonia/ViewModels/Inference/StackEditableCardViewModel.cs

@@ -117,6 +117,13 @@ private void AddModule(Type type)
         AddCards(card!);
     }
 
+    public bool IsModuleEnabled<T>(int index = 0)
+        where T : ModuleBase
+    {
+        var card = Cards.OfType<T>().ElementAtOrDefault(index);
+        return card is { IsEnabled: true };
+    }
+
     /*/// <inheritdoc />
     public override void LoadStateFromJsonObject(JsonObject state)
     {

StabilityMatrix.Avalonia/Views/CheckpointsPage.axaml

@@ -231,6 +231,11 @@
                         IconSource="Filter"
                         VerticalAlignment="Center"
                         Label="Filter">
+                        <ui:CommandBarButton.Styles>
+                            <Style Selector="FlyoutPresenter">
+                                <Setter Property="MaxHeight" Value="1000"/>
+                            </Style>
+                        </ui:CommandBarButton.Styles>
                         <ui:CommandBarButton.Flyout>
                             <Flyout>
                                 <StackPanel Orientation="Vertical">