diff --git a/CHANGELOG.md b/CHANGELOG.md index bf6f950a0..e0697ded6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,39 @@ All notable changes to Stability Matrix will be documented in this file. 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.16.1 +### Added +- Added **automatic text encoder and VAE selection** to the Inference Model card. Selecting a model now fills any empty encoder slots and the default VAE with the matching local files for the detected workflow, so you don't need to know which files pair with which architecture (e.g. `qwen_3_4b` or `qwen_3_8b` + Flux.2 VAE for Flux.2 Klein, `clip_l` + `t5xxl` for Flux, `qwen_3_06b` + `qwen_image_vae` for Anima). Anything you pick manually is never overridden +- Added a **misplaced-model warning** to the Inference Model card with a one-click **Move** button. If a model sits in a folder that can't work with the selected workflow (like a Z-Image or Anima file in the StableDiffusion folder), a compact warning explains the problem instead of letting generation fail with a cryptic ComfyUI error. The Move button relocates the file and its metadata to the right folder, then re-selects it. Dismissible per model +- Added a **native macOS menu bar** with the standard application menu — About Stability Matrix, **Settings… (⌘,)**, and the usual Services / Hide / **Quit (⌘Q)** items +### Changed +- The Inference **Workflow** selector now switches the model loader to match the chosen profile, showing or hiding the separate encoder and VAE fields as appropriate. It will never switch to a loader that can't load the selected file; you get the warning above instead +- Renamed the "Anima / SD" workflow profile to **"Anima"**. Anima has no all-in-one version, so it's now handled like Z-Image: standalone model in DiffusionModels with a separate text encoder and VAE +- Image Lab's Flux.2 Klein model checks now match the text encoder to your selected UNET variant (4B vs 9B), and switching variants updates the status banner immediately +- Image Lab's Flux Kontext and Qwen Image Edit providers now show ComfyUI's actual workflow rejection message instead of a generic "Generation failed" (Flux.2 Klein already did this) +- Image Lab now opens the same ComfyUI error detail dialog that Inference uses when a workflow is rejected or a node fails mid-generation, showing the full error JSON instead of a truncated toast +- Keyboard shortcuts now use **⌘ (Command)** instead of Ctrl on macOS — save/open, Generate, undo/redo, copy/paste/cut, tab navigation, and the mask editor all follow the platform convention, and context menus show the ⌘ glyph +### Fixed +- Fixed [#1659](https://github.com/LykosAI/StabilityMatrix/issues/1659) - Z-Image and Anima workflows hiding the Text Encoder selectors and passing an invalid `None` CLIP input to ComfyUI; standalone workflows now expose and automatically fill compatible text encoders and VAEs +- Fixed [#1654](https://github.com/LykosAI/StabilityMatrix/issues/1654) and [#1658](https://github.com/LykosAI/StabilityMatrix/issues/1658) - Inference Image-to-Image masks could disappear in Linux AppImage builds, fail to appear in the image-card preview, or crash the app when saving an enabled clipping mask because Avalonia retained pixels owned by a disposed Skia image; converted mask bitmaps now own their pixel data and previews refresh after editing, restoring a project, or loading image dimensions +- Fixed [#1660](https://github.com/LykosAI/StabilityMatrix/issues/1660) - CivArchive downloads always saving to the model folder root; the primary Download button now offers the inferred folder, existing subdirectories, and a custom folder picker, while filename patterns containing path separators create missing nested subfolders and keep downloads inside the selected models directory +- Fixed [#1661](https://github.com/LykosAI/StabilityMatrix/issues/1661) - Custom UNet workflows missing current ComfyUI encoder types such as `sdxl`; the selector now includes current CLIP loader values while retaining single-encoder compatibility and migrating legacy `HiDream` casing +- Fixed [#1664](https://github.com/LykosAI/StabilityMatrix/issues/1664) - Gemini failures with a saved key being misreported as "API key not configured"; Image Lab now shows the actual invalid-key, billing, quota, or permission error and explains that Nano Banana image generation requires a paid-tier API key from a Google AI project with billing enabled +- Fixed **"No text encoders configured"** errors when generating with an all-in-one checkpoint after a UNet model had been selected in the same tab +- Fixed Qwen Image Edit in Image Lab failing mid-generation when a wrong-size Qwen2.5-VL text encoder was installed. The **7B** encoder is now required, and the correct download is offered when it's missing +- Fixed Image Lab reporting "all models present" for Flux.2 Klein 9B setups that only had the 4B text encoder (and vice versa). The matching encoder download is now offered +- Fixed Image Lab model and LoRA dropdowns hiding files whose CivitAI base model tag is unrecognized (commonly "Other"), even when the filename clearly matches the provider +- Fixed Animagine XL and other SDXL models with "anima" in the name being misdetected as the Anima architecture +- Fixed the running-package exit confirmation not appearing when quitting via **⌘Q, the macOS app menu, or the Dock** — it previously only showed when closing the window directly, so those paths could tear down running packages without warning. All quit paths now prompt +- Fixed an empty strip of space appearing below the native window title bar on macOS and Linux; the Windows-only caption area (icon, title, and min/max/close buttons) is now collapsed on those platforms so content sits directly under the system title bar +### Performance +- Optimized SKBitmap-to-WriteableBitmap conversion used throughout Inference and Image Lab; pixel data is now copied directly native-to-native (single bulk copy when strides match, per-row otherwise) instead of round-tripping every scanline through a temporary managed buffer, reducing allocations and GC pressure on full-resolution images +### Supporters +#### 🌟 Visionaries +A massive thank you to our brilliant Visionaries: **Waterclouds**, **bluepopsicle**, **Ibixat**, **Droolguy**, **snotty**, **LG**, **whudunit**, **MrMxyzptlk12836**, **Psilocyfer18731**, **KalAbaddon**, and **moon_milky2843**! There's a little of your support behind every fix and refinement in this update, and we're grateful for all of it. A warm welcome to our newest Visionary, **cusalapapen1481**; it's wonderful to have you with us! 💛 +#### 🚀 Pioneers +And an equally big thank you to our fantastic Pioneer crew, all familiar faces this time around: **Szir777**, **[USA]TechDude**, **SinthCore**, **Jisuren**, **Tigon**, **jweg79**, **rwx14662**, **Hurbie53**, **ahnhj.al**, **drew.lukas**, **Tuskaruho**, **Cjloha**, **Alligator1907**, **Bitti**, **Ghislain G**, **CommissarGiygas16050**, **qob97515211**, **bastardofbethlehem**, and **Zombop**! You keep showing up for us, and that steadiness is a big part of how this project keeps moving forward. Thank you, truly, every one of you. 💛 + ## v2.16.0 ### Added #### New Feature: 🧪 Image Lab - Conversational Image Generation for ComfyUI diff --git a/StabilityMatrix.Avalonia/App.axaml.cs b/StabilityMatrix.Avalonia/App.axaml.cs index 0b3d14f8f..7f7be2d5f 100644 --- a/StabilityMatrix.Avalonia/App.axaml.cs +++ b/StabilityMatrix.Avalonia/App.axaml.cs @@ -3,6 +3,7 @@ using System.Net; using System.Net.Http.Headers; using System.Reflection; +using System.Runtime.Versioning; using System.Text.Json; using System.Text.Json.Serialization; using Apizr; @@ -12,6 +13,7 @@ using Avalonia.Controls; using Avalonia.Controls.ApplicationLifetimes; using Avalonia.Data.Core.Plugins; +using Avalonia.Input; using Avalonia.Input.Platform; using Avalonia.Markup.Xaml; using Avalonia.Media; @@ -100,6 +102,18 @@ public sealed class App : Application private bool isOnExitComplete; + /// + /// True once the user has confirmed exiting while packages are running, so + /// doesn't prompt again on the follow-up shutdown. + /// + private bool isExitConfirmed; + + /// + /// True while the exit confirmation dialog is open, to avoid stacking dialogs if more + /// shutdown requests arrive (e.g. ⌘Q pressed repeatedly). + /// + private bool isConfirmingExit; + private ServiceProvider? serviceProvider; [NotNull] @@ -140,6 +154,13 @@ public override void Initialize() SetFontFamily(GetPlatformDefaultFontFamily()); + // macOS app menu must be set here (before AppBuilder's AfterSetup creates the menu + // exporter) so it becomes the app menu and Avalonia appends the standard Services/Hide/Quit + if (Compat.IsMacOS && !Design.IsDesignMode) + { + SetupMacOsApplicationMenu(); + } + // Set design theme if (Design.IsDesignMode) { @@ -220,6 +241,47 @@ public override void OnFrameworkInitializationCompleted() } } + /// + /// Sets the native macOS application menu. Avalonia's native backend reads the app menu (the + /// bold app-name menu) from attached to the + /// during AppBuilder's AfterSetup phase, then appends the standard Services / Hide / Quit (⌘Q) + /// items. We add About and Settings… (⌘,) above those. Called from , + /// which runs just before that phase, so it must not be moved any later or Avalonia falls back + /// to its default "About Avalonia" menu. + /// + [SupportedOSPlatform("macos")] + private void SetupMacOsApplicationMenu() + { + var aboutItem = new NativeMenuItem("About Stability Matrix"); + aboutItem.Click += (_, _) => ShowAboutDialog(); + + var settingsItem = new NativeMenuItem("Settings…") + { + Gesture = new KeyGesture(Key.OemComma, KeyModifiers.Meta), + }; + settingsItem.Click += (_, _) => + Services + .GetRequiredService>() + .NavigateTo(); + + var appMenu = new NativeMenu { Items = { aboutItem, new NativeMenuItemSeparator(), settingsItem } }; + + NativeMenu.SetMenu(this, appMenu); + } + + private static void ShowAboutDialog() + { + var dialog = DialogHelper.CreateTaskDialog( + "Stability Matrix", + $"Version {Compat.AppVersion.ToDisplayString()}" + ); + dialog.ShowProgressBar = false; + dialog.FooterVisibility = TaskDialogFooterVisibility.Never; + dialog.Buttons = new List { TaskDialogButton.CloseButton }; + + dialog.ShowAsync(true).SafeFireAndForget(); + } + /// /// Set the default font family for the application. /// @@ -899,6 +961,26 @@ public static void Shutdown(int exitCode = 0) } } + private static TaskDialog CreateExitConfirmDialog() + { + var dialog = DialogHelper.CreateTaskDialog( + Languages.Resources.Label_ConfirmExit, + Languages.Resources.Label_ConfirmExitDetail + ); + + dialog.ShowProgressBar = false; + dialog.FooterVisibility = TaskDialogFooterVisibility.Never; + + dialog.Buttons = new List + { + new("Exit", TaskDialogStandardResult.Yes), + TaskDialogButton.CancelButton, + }; + dialog.Buttons[0].IsDefault = true; + + return dialog; + } + private void OnShutdownRequested(object? sender, ShutdownRequestedEventArgs e) { Logger.Trace("Start OnShutdownRequested"); @@ -906,9 +988,54 @@ private void OnShutdownRequested(object? sender, ShutdownRequestedEventArgs e) if (e.Cancel) return; - // Skip if Async Dispose already started, shutdown will be handled by it + // Confirm exit while packages are running. This covers every quit path — the window + // close button, ⌘Q, the app menu Quit, and dock Quit — since they all end up here. + if (!isExitConfirmed && !isAsyncDisposeStarted && serviceProvider is not null) + { + var runningPackageService = serviceProvider.GetRequiredService(); + if (runningPackageService.RunningPackages.Count > 0) + { + e.Cancel = true; + + // Avoid stacking dialogs if another shutdown request arrives while this one is open + if (isConfirmingExit) + return; + + isConfirmingExit = true; + Dispatcher + .UIThread.InvokeAsync(async () => + { + try + { + var dialog = CreateExitConfirmDialog(); + if ( + (TaskDialogStandardResult)await dialog.ShowAsync(true) + == TaskDialogStandardResult.Yes + ) + { + isExitConfirmed = true; + DesktopLifetime?.MainWindow?.Hide(); + Shutdown(); + } + } + finally + { + isConfirmingExit = false; + } + }) + .SafeFireAndForget(); + return; + } + } + + // If an async dispose is already running, cancel until it completes so we don't + // Environment.Exit before settings/database flushes finish if (isAsyncDisposeStarted) + { + if (!isAsyncDisposeComplete) + e.Cancel = true; return; + } // Cancel shutdown for now to dispose e.Cancel = true; diff --git a/StabilityMatrix.Avalonia/Behaviors/ResizeBehavior.cs b/StabilityMatrix.Avalonia/Behaviors/ResizeBehavior.cs index cf47cd17b..2ee7c89a9 100644 --- a/StabilityMatrix.Avalonia/Behaviors/ResizeBehavior.cs +++ b/StabilityMatrix.Avalonia/Behaviors/ResizeBehavior.cs @@ -7,6 +7,7 @@ using Avalonia.Media; using Avalonia.Xaml.Interactivity; using FluentAvalonia.UI.Controls; +using StabilityMatrix.Avalonia.Helpers; namespace StabilityMatrix.Avalonia.Behaviors; @@ -109,7 +110,7 @@ protected override void OnDetaching() private void OnPointerWheelChanged(object? sender, PointerWheelEventArgs e) { - if (e.KeyModifiers != KeyModifiers.Control) + if (e.KeyModifiers != PlatformKeyModifiers.CommandModifier) return; if (!UseMouseWheelResize) diff --git a/StabilityMatrix.Avalonia/Controls/AdvancedImageBoxView.axaml b/StabilityMatrix.Avalonia/Controls/AdvancedImageBoxView.axaml index 692de0ecd..cf7531ff1 100644 --- a/StabilityMatrix.Avalonia/Controls/AdvancedImageBoxView.axaml +++ b/StabilityMatrix.Avalonia/Controls/AdvancedImageBoxView.axaml @@ -4,30 +4,27 @@ xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml" xmlns:controls="clr-namespace:StabilityMatrix.Avalonia.Controls" xmlns:d="http://schemas.microsoft.com/expression/blend/2008" + xmlns:gif="clr-namespace:Avalonia.Gif;assembly=Avalonia.Gif" + xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages" + xmlns:markupExtensions="clr-namespace:StabilityMatrix.Avalonia.MarkupExtensions" xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006" xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData" xmlns:models="clr-namespace:StabilityMatrix.Avalonia.Models" xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia" - xmlns:gif="clr-namespace:Avalonia.Gif;assembly=Avalonia.Gif" - xmlns:lang="clr-namespace:StabilityMatrix.Avalonia.Languages" d:DataContext="{x:Static mocks:DesignData.SampleImageSource}" d:DesignHeight="450" d:DesignWidth="800" x:DataType="models:ImageSource" mc:Ignorable="d"> - - + + - + - + - + - + Text="Unsupported Format" /> - + + CornerRadius="16" + IsVisible="{Binding Label, Converter={x:Static StringConverters.IsNotNullOrEmpty}}"> diff --git a/StabilityMatrix.Avalonia/Controls/Inference/ImageFolderCard.axaml.cs b/StabilityMatrix.Avalonia/Controls/Inference/ImageFolderCard.axaml.cs index 430cb4442..897359735 100644 --- a/StabilityMatrix.Avalonia/Controls/Inference/ImageFolderCard.axaml.cs +++ b/StabilityMatrix.Avalonia/Controls/Inference/ImageFolderCard.axaml.cs @@ -2,6 +2,7 @@ using Avalonia.Controls.Primitives; using Avalonia.Input; using Injectio.Attributes; +using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.ViewModels.Inference; using StabilityMatrix.Core.Models.Settings; @@ -34,7 +35,7 @@ protected override void DragOverHandler(object? sender, DragEventArgs e) protected override void OnPointerWheelChanged(PointerWheelEventArgs e) { - if (e.KeyModifiers != KeyModifiers.Control) + if (e.KeyModifiers != PlatformKeyModifiers.CommandModifier) return; if (DataContext is not ImageFolderCardViewModel vm) return; diff --git a/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml b/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml index bb3ba422c..a227f40af 100644 --- a/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml +++ b/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml @@ -182,6 +182,54 @@ IsVisible="{Binding ShowWorkflowProfileStatus}" Text="{Binding WorkflowProfileStatusText}" /> + + + + + + + + + diff --git a/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml b/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml index 59b32a321..07b498084 100644 --- a/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml +++ b/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml @@ -5,6 +5,7 @@ xmlns:converters="clr-namespace:StabilityMatrix.Avalonia.Converters" xmlns:fluentIcons="clr-namespace:FluentIcons.Avalonia.Fluent;assembly=FluentIcons.Avalonia.Fluent" xmlns:input="clr-namespace:FluentAvalonia.UI.Input;assembly=FluentAvalonia" + xmlns:markupExtensions="clr-namespace:StabilityMatrix.Avalonia.MarkupExtensions" xmlns:mocks="clr-namespace:StabilityMatrix.Avalonia.DesignData" xmlns:models="clr-namespace:StabilityMatrix.Avalonia.Models" xmlns:ui="clr-namespace:FluentAvalonia.UI.Controls;assembly=FluentAvalonia" @@ -32,7 +33,7 @@ diff --git a/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml.cs b/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml.cs index 2a960ed0f..a902bbff8 100644 --- a/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml.cs +++ b/StabilityMatrix.Avalonia/Controls/Painting/PaintCanvas.axaml.cs @@ -18,6 +18,7 @@ using DynamicData.Binding; using SkiaSharp; using StabilityMatrix.Avalonia.Controls.Models; +using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models; using StabilityMatrix.Avalonia.ViewModels.Controls; @@ -410,7 +411,7 @@ protected override void OnKeyDown(KeyEventArgs e) return; // Check for modifier keys - var isCtrl = e.KeyModifiers.HasFlag(KeyModifiers.Control); + var isCtrl = e.KeyModifiers.HasFlag(PlatformKeyModifiers.CommandModifier); // Ctrl+Z: Undo if (isCtrl && e.Key == Key.Z && !e.KeyModifiers.HasFlag(KeyModifiers.Shift)) diff --git a/StabilityMatrix.Avalonia/Controls/Scroll/BetterScrollContentPresenter.cs b/StabilityMatrix.Avalonia/Controls/Scroll/BetterScrollContentPresenter.cs index a5f30ff15..53c8880e4 100644 --- a/StabilityMatrix.Avalonia/Controls/Scroll/BetterScrollContentPresenter.cs +++ b/StabilityMatrix.Avalonia/Controls/Scroll/BetterScrollContentPresenter.cs @@ -1,5 +1,6 @@ using Avalonia.Controls.Presenters; using Avalonia.Input; +using StabilityMatrix.Avalonia.Helpers; namespace StabilityMatrix.Avalonia.Controls.Scroll; @@ -7,7 +8,7 @@ public class BetterScrollContentPresenter : ScrollContentPresenter { protected override void OnPointerWheelChanged(PointerWheelEventArgs e) { - if (e.KeyModifiers == KeyModifiers.Control) + if (e.KeyModifiers == PlatformKeyModifiers.CommandModifier) return; base.OnPointerWheelChanged(e); } diff --git a/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs b/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs index 1b91c5d28..09504217e 100644 --- a/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs +++ b/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs @@ -110,7 +110,7 @@ public static Bitmap ToAvaloniaBitmap(this SKBitmap bitmap, Vector dpi) { SKColorType.Rgba8888 => PixelFormat.Rgba8888, SKColorType.Bgra8888 => PixelFormat.Bgra8888, - _ => throw new NotSupportedException($"Unsupported SKColorType: {bitmap.ColorType}") + _ => throw new NotSupportedException($"Unsupported SKColorType: {bitmap.ColorType}"), }; // ReSharper disable once SwitchExpressionHandlesSomeKnownEnumValuesWithExceptionInDefault @@ -119,19 +119,25 @@ public static Bitmap ToAvaloniaBitmap(this SKBitmap bitmap, Vector dpi) SKAlphaType.Opaque => AlphaFormat.Opaque, SKAlphaType.Premul => AlphaFormat.Premul, SKAlphaType.Unpremul => AlphaFormat.Unpremul, - _ => throw new NotSupportedException($"Unsupported SKAlphaType: {bitmap.AlphaType}") + _ => throw new NotSupportedException($"Unsupported SKAlphaType: {bitmap.AlphaType}"), }; - var dataPointer = bitmap.GetPixels(); - - return new Bitmap( - avaloniaColorFormat, - avaloniaAlphaFormat, - dataPointer, + var result = new WriteableBitmap( new PixelSize(bitmap.Width, bitmap.Height), dpi, - bitmap.RowBytes + avaloniaColorFormat, + avaloniaAlphaFormat ); + + using var framebuffer = result.Lock(); + CopyPixelRows( + bitmap.GetPixels(), + bitmap.RowBytes, + framebuffer.Address, + framebuffer.RowBytes, + bitmap.Height + ); + return result; } public static Bitmap ToAvaloniaBitmap(this SKImage image) @@ -143,12 +149,20 @@ public static Bitmap ToAvaloniaBitmap(this SKImage image, Vector dpi) { ArgumentNullException.ThrowIfNull(image, nameof(image)); + // PeekPixels returns null for GPU-backed images; only then do we need the extra copy. + using var pixmap = image.PeekPixels(); + if (pixmap is null) + { + using var bitmap = SKBitmap.FromImage(image); + return bitmap.ToAvaloniaBitmap(dpi); + } + // ReSharper disable once SwitchExpressionHandlesSomeKnownEnumValuesWithExceptionInDefault var avaloniaColorFormat = image.ColorType switch { SKColorType.Rgba8888 => PixelFormat.Rgba8888, SKColorType.Bgra8888 => PixelFormat.Bgra8888, - _ => throw new NotSupportedException($"Unsupported SKColorType: {image.ColorType}") + _ => throw new NotSupportedException($"Unsupported SKColorType: {image.ColorType}"), }; // ReSharper disable once SwitchExpressionHandlesSomeKnownEnumValuesWithExceptionInDefault @@ -157,20 +171,56 @@ public static Bitmap ToAvaloniaBitmap(this SKImage image, Vector dpi) SKAlphaType.Opaque => AlphaFormat.Opaque, SKAlphaType.Premul => AlphaFormat.Premul, SKAlphaType.Unpremul => AlphaFormat.Unpremul, - _ => throw new NotSupportedException($"Unsupported SKAlphaType: {image.AlphaType}") + _ => throw new NotSupportedException($"Unsupported SKAlphaType: {image.AlphaType}"), }; - var pixmap = image.PeekPixels(); - var dataPointer = pixmap.GetPixels(); - - return new Bitmap( - avaloniaColorFormat, - avaloniaAlphaFormat, - dataPointer, + var result = new WriteableBitmap( new PixelSize(image.Width, image.Height), dpi, - pixmap.RowBytes + avaloniaColorFormat, + avaloniaAlphaFormat ); + + using var framebuffer = result.Lock(); + CopyPixelRows( + pixmap.GetPixels(), + pixmap.RowBytes, + framebuffer.Address, + framebuffer.RowBytes, + image.Height + ); + return result; + } + + private static unsafe void CopyPixelRows( + IntPtr source, + int sourceRowBytes, + IntPtr destination, + int destinationRowBytes, + int height + ) + { + var bytesPerRow = Math.Min(sourceRowBytes, destinationRowBytes); + var src = (byte*)source; + var dest = (byte*)destination; + + // When strides match the buffer is contiguous, so copy it all in one shot. + if (sourceRowBytes == destinationRowBytes) + { + var totalBytes = (long)sourceRowBytes * height; + Buffer.MemoryCopy(src, dest, totalBytes, totalBytes); + return; + } + + for (var row = 0; row < height; row++) + { + Buffer.MemoryCopy( + src + ((long)row * sourceRowBytes), + dest + ((long)row * destinationRowBytes), + bytesPerRow, + bytesPerRow + ); + } } public static PixelFormat ToAvaloniaPixelFormat(this SKColorType colorType) @@ -180,7 +230,7 @@ public static PixelFormat ToAvaloniaPixelFormat(this SKColorType colorType) { SKColorType.Rgba8888 => PixelFormat.Rgba8888, SKColorType.Bgra8888 => PixelFormat.Bgra8888, - _ => throw new NotSupportedException($"Unsupported SKColorType: {colorType}") + _ => throw new NotSupportedException($"Unsupported SKColorType: {colorType}"), }; } } diff --git a/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs b/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs new file mode 100644 index 000000000..4c033395b --- /dev/null +++ b/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs @@ -0,0 +1,118 @@ +using Microsoft.Extensions.Logging; +using Refit; +using StabilityMatrix.Core.Exceptions; +using StabilityMatrix.Core.Models.Api.Comfy; +using StabilityMatrix.Core.Services.ImageGeneration; + +namespace StabilityMatrix.Avalonia.Helpers; + +/// +/// Shared logic for Image Lab providers that generate via a local ComfyUI backend +/// +public static class ComfyGenerationHelper +{ + /// + /// Builds the failure response for a workflow ComfyUI rejected at queue time + /// (Refit with the validation JSON in the body). The raw + /// body is carried in so the UI + /// can show it in a detail dialog. + /// + public static ImageGenerationResponse CreateWorkflowRejectedResponse( + ApiException apiEx, + string providerName, + ILogger logger + ) + { + var detail = !string.IsNullOrWhiteSpace(apiEx.Content) ? apiEx.Content : apiEx.Message; + logger.LogError( + apiEx, + "ComfyUI rejected {Provider} workflow ({StatusCode}): {Detail}", + providerName, + apiEx.StatusCode, + detail + ); + + return new ImageGenerationResponse + { + IsSuccess = false, + ErrorMessage = + $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {Truncate(detail, 800)}", + ErrorDetailJson = string.IsNullOrWhiteSpace(apiEx.Content) ? null : apiEx.Content, + }; + } + + /// + /// Builds the failure response for a node that failed during execution + /// (, e.g. a tensor-shape mismatch from a wrong + /// encoder pairing). The full error JSON is carried in + /// for the detail dialog. + /// + public static ImageGenerationResponse CreateNodeErrorResponse( + ComfyNodeException nodeEx, + string providerName, + ILogger logger + ) + { + logger.LogError( + nodeEx, + "ComfyUI node execution failed for {Provider}: {Json}", + providerName, + nodeEx.JsonData + ); + + var nodeType = nodeEx.ErrorData.NodeType; + var exceptionMessage = nodeEx.ErrorData.ExceptionMessage ?? "Unknown error"; + + return new ImageGenerationResponse + { + IsSuccess = false, + ErrorMessage = string.IsNullOrEmpty(nodeType) + ? $"ComfyUI node execution failed: {Truncate(exceptionMessage, 400)}" + : $"ComfyUI node '{nodeType}' failed: {Truncate(exceptionMessage, 400)}", + ErrorDetailJson = nodeEx.JsonData, + }; + } + + /// + /// Selects the output images from an executed prompt's outputs: prefers the + /// "SaveImage" node deterministically, otherwise the first non-empty output + /// by ordinal key order. Returns (null, null) if no output has images. + /// + public static (string? OutputKey, List? Images) SelectOutputImages( + Dictionary?> outputImages + ) + { + const string preferredOutputKey = "SaveImage"; + + if ( + outputImages.TryGetValue(preferredOutputKey, out var preferredImages) + && preferredImages is { Count: > 0 } + ) + { + return (preferredOutputKey, preferredImages); + } + + var selected = outputImages + .OrderBy(kvp => kvp.Key, StringComparer.Ordinal) + .FirstOrDefault(kvp => kvp.Value is { Count: > 0 }); + + return (string.IsNullOrEmpty(selected.Key) ? null : selected.Key, selected.Value); + } + + /// + /// Gets the MIME type for a generated image filename by extension, defaulting to PNG + /// + public static string GetMimeTypeForFileName(string fileName) => + fileName.EndsWith(".png", StringComparison.OrdinalIgnoreCase) ? "image/png" + : fileName.EndsWith(".jpg", StringComparison.OrdinalIgnoreCase) + || fileName.EndsWith(".jpeg", StringComparison.OrdinalIgnoreCase) + ? "image/jpeg" + : fileName.EndsWith(".webp", StringComparison.OrdinalIgnoreCase) ? "image/webp" + : "image/png"; + + /// + /// Truncates a string to at most characters, appending "..." + /// + public static string Truncate(string value, int maxLength) => + value.Length <= maxLength ? value : value[..maxLength] + "..."; +} diff --git a/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs b/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs new file mode 100644 index 000000000..09b5e9493 --- /dev/null +++ b/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs @@ -0,0 +1,82 @@ +using StabilityMatrix.Core.Inference; +using StabilityMatrix.Core.Services.ImageGeneration; + +namespace StabilityMatrix.Avalonia.Helpers; + +/// +/// Forwards a ComfyTask's progress events to an Image Lab generation request as +/// updates, deduplicated by (percent, running node). +/// Reports a "Queued" stage immediately on construction and subscribes to the task's +/// events; dispose to unsubscribe. +/// +public sealed class ComfyProgressReporter : IDisposable +{ + private readonly ComfyTask task; + private readonly string providerId; + private readonly IProgress? progress; + + private int? lastPercent; + private string? lastRunningNode; + + public ComfyProgressReporter( + ComfyTask task, + string providerId, + IProgress? progress + ) + { + ArgumentNullException.ThrowIfNull(task); + ArgumentNullException.ThrowIfNull(providerId); + + this.task = task; + this.providerId = providerId; + this.progress = progress; + + progress?.Report( + new ImageGenerationProgress( + providerId, + task.Id, + Value: null, + Maximum: null, + RunningNode: null, + Stage: "Queued" + ) + ); + + task.ProgressUpdate += OnProgressUpdate; + task.RunningNodeChanged += OnRunningNodeChanged; + } + + private void Report(int? value, int? maximum, string? runningNode, string? stage) + { + var update = new ImageGenerationProgress(providerId, task.Id, value, maximum, runningNode, stage); + + if ( + update.Percent == lastPercent + && string.Equals(lastRunningNode, runningNode, StringComparison.Ordinal) + ) + { + return; + } + + lastPercent = update.Percent; + lastRunningNode = runningNode; + + progress?.Report(update); + } + + private void OnProgressUpdate(object? sender, ComfyProgressUpdateEventArgs args) + { + Report(args.Value, args.Maximum, args.RunningNode, "Generating"); + } + + private void OnRunningNodeChanged(object? sender, string? node) + { + Report(task.LastProgressUpdate?.Value, task.LastProgressUpdate?.Maximum, node, "Generating"); + } + + public void Dispose() + { + task.ProgressUpdate -= OnProgressUpdate; + task.RunningNodeChanged -= OnRunningNodeChanged; + } +} diff --git a/StabilityMatrix.Avalonia/Helpers/ComfyWorkflowHelper.cs b/StabilityMatrix.Avalonia/Helpers/ComfyWorkflowHelper.cs new file mode 100644 index 000000000..23b9b88fd --- /dev/null +++ b/StabilityMatrix.Avalonia/Helpers/ComfyWorkflowHelper.cs @@ -0,0 +1,88 @@ +using StabilityMatrix.Avalonia.Models.BananaVision; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; +using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes; +using StabilityMatrix.Core.Services.ImageGeneration; + +namespace StabilityMatrix.Avalonia.Helpers; + +/// +/// Shared logic for Image Lab ComfyUI workflow builders +/// +public static class ComfyWorkflowHelper +{ + /// + /// Get reference image filenames from the request (input images + conversation history). + /// Note: Images uploaded via ComfyClient.UploadImageAsync go to the "Inference" subfolder, + /// using names matching what uploads. + /// + /// Generation request + /// Maximum number of reference images supported by the provider + /// Prefix for uploaded filenames (e.g. "flux_kontext") + public static List GetReferenceImageNames( + ImageGenerationRequest request, + int maxImages, + string providerPrefix + ) + { + var imageNames = new List(); + + // Priority 1: Current input images (uploaded with known names by the provider) + if (request.InputImages?.Count > 0) + { + for (var i = 0; i < Math.Min(request.InputImages.Count, maxImages); i++) + { + imageNames.Add($"Inference/{providerPrefix}_input_{i}.png"); + } + } + + // Priority 2: Most recent image from conversation history (previous generation). + // Only include if we have room and we actually have the image content + // (the provider will have uploaded it). + if (imageNames.Count < maxImages && request.ConversationHistory != null) + { + var lastAssistantImage = request.ConversationHistory.LastOrDefault(m => + m is { Role: MessageRole.Assistant, ImageContent: not null } + ); + + if (lastAssistantImage?.ImageContent != null) + { + imageNames.Add($"Inference/{providerPrefix}_history_latest.png"); + } + } + + return imageNames; + } + + /// + /// Chains a LoraLoader node for each selected LoRA onto the model and clip connections, + /// returning the final connections (unchanged if no LoRAs are selected) + /// + public static (ModelNodeConnection Model, ClipNodeConnection Clip) ApplyLoras( + NodeDictionary nodes, + IEnumerable? loras, + ModelNodeConnection model, + ClipNodeConnection clip + ) + { + var loraList = loras?.ToList() ?? []; + + for (var i = 0; i < loraList.Count; i++) + { + var lora = loraList[i]; + var loraLoader = nodes.AddNamedNode( + ComfyNodeBuilder.LoraLoader( + nodes.GetUniqueName($"LoraLoader_{i + 1}"), + model, + clip, + lora.Model.RelativePath, + (double)lora.ModelWeight, + (double)lora.ClipWeight + ) + ); + model = loraLoader.Output1; + clip = loraLoader.Output2; + } + + return (model, clip); + } +} diff --git a/StabilityMatrix.Avalonia/Helpers/PlatformKeyModifiers.cs b/StabilityMatrix.Avalonia/Helpers/PlatformKeyModifiers.cs new file mode 100644 index 000000000..82a362c09 --- /dev/null +++ b/StabilityMatrix.Avalonia/Helpers/PlatformKeyModifiers.cs @@ -0,0 +1,16 @@ +using Avalonia.Input; +using StabilityMatrix.Core.Helper; + +namespace StabilityMatrix.Avalonia.Helpers; + +/// +/// Platform-aware keyboard modifiers for code-behind shortcut handling. +/// +public static class PlatformKeyModifiers +{ + /// + /// The primary command modifier for the current platform: + /// (⌘) on macOS, elsewhere. + /// + public static KeyModifiers CommandModifier => Compat.IsMacOS ? KeyModifiers.Meta : KeyModifiers.Control; +} diff --git a/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs index 853d33248..fe9ce0b73 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs +++ b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs @@ -103,7 +103,16 @@ public static string Action_Cancel { return ResourceManager.GetString("Action_Cancel", resourceCulture); } } - + + /// + /// Looks up a localized string similar to Custom.... + /// + public static string Action_CustomFolderEllipsis { + get { + return ResourceManager.GetString("Action_CustomFolderEllipsis", resourceCulture); + } + } + /// /// Looks up a localized string similar to Check for Updates. /// @@ -4834,5 +4843,203 @@ public static string Watermark_EnterPackageName { return ResourceManager.GetString("Watermark_EnterPackageName", resourceCulture); } } + + /// + /// Looks up a localized string similar to The {0} workflow expects a model from DiffusionModels, but this file is in {1}. It may not load correctly.. + /// + public static string TextTemplate_WorkflowNeedsDiffusionModelsFile { + get { + return ResourceManager.GetString("TextTemplate_WorkflowNeedsDiffusionModelsFile", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to This file is in DiffusionModels and can't load as an all-in-one checkpoint. It may not load correctly.. + /// + public static string Text_FileInDiffusionModelsNotCheckpoint { + get { + return ResourceManager.GetString("Text_FileInDiffusionModelsNotCheckpoint", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to This looks like {0}, which loads from DiffusionModels, but the file is in {1}. It may not load correctly.. + /// + public static string TextTemplate_FileLooksLikeStandaloneModel { + get { + return ResourceManager.GetString("TextTemplate_FileLooksLikeStandaloneModel", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Move to {0}. + /// + public static string TextTemplate_MoveToFolder { + get { + return ResourceManager.GetString("TextTemplate_MoveToFolder", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Move. + /// + public static string Action_Move { + get { + return ResourceManager.GetString("Action_Move", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Could not move model. + /// + public static string Label_CouldNotMoveModel { + get { + return ResourceManager.GetString("Label_CouldNotMoveModel", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to A file named '{0}' already exists in the {1} folder.. + /// + public static string TextTemplate_FileAlreadyExistsInFolder { + get { + return ResourceManager.GetString("TextTemplate_FileAlreadyExistsInFolder", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Model moved. + /// + public static string Label_ModelMoved { + get { + return ResourceManager.GetString("Label_ModelMoved", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Moved '{0}' to the {1} folder.. + /// + public static string TextTemplate_MovedFileToFolder { + get { + return ResourceManager.GetString("TextTemplate_MovedFileToFolder", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Comfy Error. + /// + public static string Label_ComfyError { + get { + return ResourceManager.GetString("Label_ComfyError", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to ComfyUI reported an error during generation. + /// + public static string Text_ComfyReportedGenerationError { + get { + return ResourceManager.GetString("Text_ComfyReportedGenerationError", resourceCulture); + } + } + + /// + /// Looks up a localized string similar to Generation Failed. + /// + public static string Label_GenerationFailed { + get { + return ResourceManager.GetString("Label_GenerationFailed", resourceCulture); + } + } + + public static string Label_Success { + get { + return ResourceManager.GetString("Label_Success", resourceCulture); + } + } + + public static string Label_GeminiApiKey { + get { + return ResourceManager.GetString("Label_GeminiApiKey", resourceCulture); + } + } + + public static string Validation_ApiKeyRequired { + get { + return ResourceManager.GetString("Validation_ApiKeyRequired", resourceCulture); + } + } + + public static string Label_SetGeminiApiKey { + get { + return ResourceManager.GetString("Label_SetGeminiApiKey", resourceCulture); + } + } + + public static string Text_SetGeminiApiKeyDescription { + get { + return ResourceManager.GetString("Text_SetGeminiApiKeyDescription", resourceCulture); + } + } + + public static string Text_GeminiApiKeySaved { + get { + return ResourceManager.GetString("Text_GeminiApiKeySaved", resourceCulture); + } + } + + public static string Label_RemoveGeminiApiKey { + get { + return ResourceManager.GetString("Label_RemoveGeminiApiKey", resourceCulture); + } + } + + public static string Text_ConfirmRemoveGeminiApiKey { + get { + return ResourceManager.GetString("Text_ConfirmRemoveGeminiApiKey", resourceCulture); + } + } + + public static string Text_GeminiApiKeyRemoved { + get { + return ResourceManager.GetString("Text_GeminiApiKeyRemoved", resourceCulture); + } + } + + public static string Label_GeminiApiSetupRequired { + get { + return ResourceManager.GetString("Label_GeminiApiSetupRequired", resourceCulture); + } + } + + public static string Text_GeminiPaidTierRequired { + get { + return ResourceManager.GetString("Text_GeminiPaidTierRequired", resourceCulture); + } + } + + public static string Error_GeminiApiKeyNotConfigured { + get { + return ResourceManager.GetString("Error_GeminiApiKeyNotConfigured", resourceCulture); + } + } + + public static string Error_GeminiQuotaExceeded { + get { + return ResourceManager.GetString("Error_GeminiQuotaExceeded", resourceCulture); + } + } + + public static string Error_GeminiInvalidApiKey { + get { + return ResourceManager.GetString("Error_GeminiInvalidApiKey", resourceCulture); + } + } + + public static string Error_GeminiAccessForbidden { + get { + return ResourceManager.GetString("Error_GeminiAccessForbidden", resourceCulture); + } + } } } diff --git a/StabilityMatrix.Avalonia/Languages/Resources.de.resx b/StabilityMatrix.Avalonia/Languages/Resources.de.resx index cb9538ec2..b48614e7b 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.de.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.de.resx @@ -1087,7 +1087,17 @@ Bildgenerierungs-APIs - Füge deinen Gemini-API-Schlüssel hinzu, um Nano Banana im Image Lab zur Bildgenerierung zu nutzen + Füge einen kostenpflichtigen Gemini-API-Schlüssel (Abrechnung aktiviert) hinzu, um Nano Banana im Image Lab zu nutzen + + + Die Nano-Banana-Modelle im Image Lab benötigen einen **kostenpflichtigen** Gemini-API-Schlüssel – Schlüssel der kostenlosen Stufe können keine Bilder generieren. + +1. Erstelle einen Schlüssel im [Google AI Studio](https://aistudio.google.com/apikey). +2. Aktiviere die Abrechnung für das zugehörige Google-Cloud-Projekt in der [Cloud Console](https://console.cloud.google.com/billing). +3. Füge den Schlüssel unten ein. + + + Benutzerdefiniert... Über Konto-Anmeldeinformationen diff --git a/StabilityMatrix.Avalonia/Languages/Resources.fr-FR.resx b/StabilityMatrix.Avalonia/Languages/Resources.fr-FR.resx index 8d6d7778d..3e67c967f 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.fr-FR.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.fr-FR.resx @@ -1179,7 +1179,17 @@ APIs de génération d'images - Ajoutez votre clé API Gemini pour utiliser Nano Banana dans l'Image Lab pour la génération d'images + Ajoutez une clé API Gemini payante (facturation activée) pour utiliser Nano Banana dans Image Lab + + + Les modèles Nano Banana de l'Image Lab nécessitent une clé API Gemini **payante** — les clés de l'offre gratuite ne peuvent pas générer d'images. + +1. Créez une clé sur [Google AI Studio](https://aistudio.google.com/apikey). +2. Activez la facturation pour son projet Google Cloud dans la [Cloud Console](https://console.cloud.google.com/billing). +3. Collez la clé ci-dessous. + + + Personnalisé... À propos des identifiants de compte diff --git a/StabilityMatrix.Avalonia/Languages/Resources.ja-JP.resx b/StabilityMatrix.Avalonia/Languages/Resources.ja-JP.resx index a702e8caa..2aac03425 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.ja-JP.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.ja-JP.resx @@ -1500,7 +1500,17 @@ Sparkは兆レベルのパラメータを持つ基盤モデルで、膨大なパ 画像生成API - Gemini APIキーを追加して、Image LabでNano Bananaを使用して画像を生成します + Image LabでNano Bananaを使用するには、課金が有効な有料Gemini APIキーを追加してください + + + Image Lab の Nano Banana モデルでは、**有料版**の Gemini API キーが必要です。無料版のキーでは画像を生成できません。 + +1. [Google AI Studio](https://aistudio.google.com/apikey) でキーを作成します。 +2. [Cloud Console](https://console.cloud.google.com/billing) で、そのキーの Google Cloud プロジェクトの課金を有効にします。 +3. 下にキーを貼り付けます。 + + + カスタム... アカウント認証情報について diff --git a/StabilityMatrix.Avalonia/Languages/Resources.ko-KR.resx b/StabilityMatrix.Avalonia/Languages/Resources.ko-KR.resx index 847b18e25..7fbe06ac6 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.ko-KR.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.ko-KR.resx @@ -1326,7 +1326,17 @@ 이미지 생성 API - Image Lab에서 Nano Banana를 이미지 생성에 사용하려면 Gemini API 키를 추가하세요 + Image Lab에서 Nano Banana를 사용하려면 결제가 활성화된 유료 Gemini API 키를 추가하세요 + + + Image Lab의 Nano Banana 모델에는 **유료** Gemini API 키가 필요합니다. 무료 등급 키로는 이미지를 생성할 수 없습니다. + +1. [Google AI Studio](https://aistudio.google.com/apikey)에서 키를 생성하세요. +2. [Cloud Console](https://console.cloud.google.com/billing)에서 해당 키의 Google Cloud 프로젝트에 결제를 활성화하세요. +3. 아래에 키를 붙여넣으세요. + + + 사용자 지정... 계정 자격 증명 정보 diff --git a/StabilityMatrix.Avalonia/Languages/Resources.resx b/StabilityMatrix.Avalonia/Languages/Resources.resx index 8048d9ea9..d172e3bcb 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.resx @@ -1622,7 +1622,7 @@ We prioritize your privacy ([Gen AI Terms](<https://lykos.ai/gen-ai-terms> Image Generation APIs - Add your Gemini API key to use Nano Banana in the Image Lab for image generation + Add a paid-tier Gemini API key (billing enabled) to use Nano Banana in the Image Lab About Account Credentials @@ -1744,4 +1744,92 @@ Account tokens will not be viewable after saving, please make a note of them if Example: {0} + + The {0} workflow expects a model from DiffusionModels, but this file is in {1}. It may not load correctly. + + + This file is in DiffusionModels and can't load as an all-in-one checkpoint. It may not load correctly. + + + This looks like {0}, which loads from DiffusionModels, but the file is in {1}. It may not load correctly. + + + Move to {0} + + + Move + + + Could not move model + + + A file named '{0}' already exists in the {1} folder. + + + Model moved + + + Moved '{0}' to the {1} folder. + + + Comfy Error + + + ComfyUI reported an error during generation + + + Generation Failed + + + Success + + + Gemini API Key + + + API key is required + + + Set Gemini API Key + + + Image Lab's Nano Banana models need a **paid-tier** Gemini API key — free-tier keys cannot generate images. + +1. Create a key at [Google AI Studio](https://aistudio.google.com/apikey). +2. Enable billing for its Google Cloud project in the [Cloud Console](https://console.cloud.google.com/billing). +3. Paste the key below. + + + Gemini API key saved + + + Remove Gemini API Key + + + Are you sure you want to remove your Gemini API key? + + + Gemini API key removed + + + Gemini API Setup Required + + + Nano Banana image generation requires a Gemini API key from a Google AI project with billing enabled. Free-tier keys do not support these image-generation models. + + + Gemini API key not configured. Please add it in Account Settings. + + + Rate limit or quota exceeded. Gemini image-generation models require a Google AI project with billing enabled; free-tier API keys cannot generate images. If billing is enabled, wait for the paid tier to activate or try again after the rate limit resets. + + + Invalid Gemini API key. Please check the key saved in Account Settings. + + + Gemini image generation was denied. Confirm that billing is enabled for the key's Google AI project and that the Gemini API is available to that project. + + + Custom... + diff --git a/StabilityMatrix.Avalonia/Languages/Resources.zh-Hans.resx b/StabilityMatrix.Avalonia/Languages/Resources.zh-Hans.resx index a73c5bb40..d437113d7 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.zh-Hans.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.zh-Hans.resx @@ -1450,7 +1450,17 @@ Spark 模型规模庞大，算力需求堪比万亿参数的基础模型。我 图像生成API - 添加您的 Gemini API 密钥，以便在 Image Lab 中使用 Nano Banana 生成图像 + 添加已开通结算功能的付费 Gemini API 密钥，以便在 Image Lab 中使用 Nano Banana + + + Image Lab 的 Nano Banana 模型需要**付费版** Gemini API 密钥，免费版密钥无法生成图像。 + +1. 在 [Google AI Studio](https://aistudio.google.com/apikey) 创建密钥。 +2. 在 [Cloud Console](https://console.cloud.google.com/billing) 中为该密钥所属的 Google Cloud 项目开通结算功能。 +3. 将密钥粘贴到下方。 + + + 自定义... 关于账号凭据 diff --git a/StabilityMatrix.Avalonia/MarkupExtensions/PlatformGestureExtension.cs b/StabilityMatrix.Avalonia/MarkupExtensions/PlatformGestureExtension.cs new file mode 100644 index 000000000..a9baa258c --- /dev/null +++ b/StabilityMatrix.Avalonia/MarkupExtensions/PlatformGestureExtension.cs @@ -0,0 +1,47 @@ +using Avalonia.Input; +using Avalonia.Markup.Xaml; +using StabilityMatrix.Core.Helper; + +namespace StabilityMatrix.Avalonia.MarkupExtensions; + +/// +/// Provides a whose primary modifier follows the platform convention: +/// the Ctrl token resolves to Cmd (⌘ / Meta) on macOS and stays Ctrl elsewhere. +/// Usage: Gesture="{markupExtensions:PlatformGesture Ctrl+S}" +/// +public class PlatformGestureExtension : MarkupExtension +{ + /// + /// The gesture string, e.g. Ctrl+S or Ctrl+Shift+Tab. This is the default + /// constructor argument. The Ctrl/Control token is swapped for the platform + /// command key on macOS. + /// + public string? Gesture { get; set; } + + public PlatformGestureExtension() { } + + public PlatformGestureExtension(string gesture) + { + Gesture = gesture; + } + + public override object ProvideValue(IServiceProvider serviceProvider) + { + if (string.IsNullOrWhiteSpace(Gesture)) + { + throw new InvalidOperationException("PlatformGesture requires a gesture string."); + } + + var gesture = Gesture; + + if (Compat.IsMacOS) + { + // KeyGesture.Parse maps the "Cmd" token to KeyModifiers.Meta (⌘). + gesture = gesture + .Replace("Ctrl", "Cmd", StringComparison.OrdinalIgnoreCase) + .Replace("Control", "Cmd", StringComparison.OrdinalIgnoreCase); + } + + return KeyGesture.Parse(gesture); + } +} diff --git a/StabilityMatrix.Avalonia/Models/Inference/InferenceWorkflowProfile.cs b/StabilityMatrix.Avalonia/Models/Inference/InferenceWorkflowProfile.cs index ceff46172..fd669d180 100644 --- a/StabilityMatrix.Avalonia/Models/Inference/InferenceWorkflowProfile.cs +++ b/StabilityMatrix.Avalonia/Models/Inference/InferenceWorkflowProfile.cs @@ -22,7 +22,7 @@ public enum InferenceWorkflowProfile [StringValue("Z-Image Turbo")] ZImageTurbo, - [StringValue("Anima / SD")] + [StringValue("Anima")] Anima, [StringValue("HiDream")] diff --git a/StabilityMatrix.Avalonia/Services/ComfyImageGenerationProviderBase.cs b/StabilityMatrix.Avalonia/Services/ComfyImageGenerationProviderBase.cs new file mode 100644 index 000000000..fa94d66cb --- /dev/null +++ b/StabilityMatrix.Avalonia/Services/ComfyImageGenerationProviderBase.cs @@ -0,0 +1,310 @@ +using AsyncAwaitBestPractices; +using Microsoft.Extensions.Logging; +using Refit; +using StabilityMatrix.Avalonia.Helpers; +using StabilityMatrix.Avalonia.Models.BananaVision; +using StabilityMatrix.Core.Exceptions; +using StabilityMatrix.Core.Inference; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Api.Comfy; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; +using StabilityMatrix.Core.Services.ImageGeneration; + +namespace StabilityMatrix.Avalonia.Services; + +/// +/// Base class for Image Lab providers that generate via a local ComfyUI backend. +/// Implements the shared generation flow as a template method — connection check, +/// model validation, image upload, prompt queue + progress reporting, cancellation +/// interrupt, output download, and error handling — so subclasses supply only the +/// provider-specific model requirements and workflow node graph. +/// +public abstract class ComfyImageGenerationProviderBase(ILogger logger, IInferenceClientManager clientManager) + : IImageGenerationProvider +{ + protected ILogger Logger { get; } = logger; + protected IInferenceClientManager ClientManager { get; } = clientManager; + + public abstract string ProviderId { get; } + public abstract string ProviderName { get; } + + public virtual bool SupportsImageInput => true; + public virtual bool SupportsMultiTurn => true; + public virtual bool RequiresThoughtSignatures => false; + + /// Short name used in log messages and error responses (e.g. "Flux Kontext"). + protected abstract string LogName { get; } + + /// Maximum number of input images uploaded to ComfyUI for this provider. + protected abstract int MaxInputImages { get; } + + /// Filename prefix for uploaded images (e.g. "flux_kontext"). + protected abstract string ProviderPrefix { get; } + + /// + /// Returns the names of any required models that are missing. An empty list means all + /// required models are available and generation may proceed. + /// + protected abstract IReadOnlyList GetMissingModels(ImageGenerationRequest request); + + /// + /// Extracts provider options from the request and builds the ComfyUI workflow node graph. + /// + protected abstract Dictionary BuildWorkflow(ImageGenerationRequest request); + + public async Task GenerateAsync( + ImageGenerationRequest request, + CancellationToken cancellationToken = default + ) + { + try + { + if (!ClientManager.IsConnected) + { + Logger.LogWarning("ComfyUI is not connected"); + return new ImageGenerationResponse + { + IsSuccess = false, + ErrorMessage = + "ComfyUI is not connected. Use the Launch button in the header to start and connect to ComfyUI.", + }; + } + + var missingModels = GetMissingModels(request); + if (missingModels.Count > 0) + { + var modelsList = string.Join(", ", missingModels); + Logger.LogWarning("Required models not found: {Models}", modelsList); + return new ImageGenerationResponse + { + IsSuccess = false, + ErrorMessage = + $"Required models not found: {modelsList}. Please download them from the HuggingFace model browser.", + }; + } + + await ComfyImageUploadHelper.UploadImagesAsync( + ClientManager, + request, + MaxInputImages, + ProviderPrefix, + Logger, + cancellationToken + ); + + Logger.LogInformation("Building {Provider} workflow", LogName); + var nodes = BuildWorkflow(request); + + Logger.LogInformation("Queuing prompt to ComfyUI"); + var task = await ClientManager.Client.QueuePromptAsync(nodes, cancellationToken); + + // Reports "Queued" now, then deduplicated progress updates until disposed + using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); + + // Interrupt the running ComfyUI prompt if the user cancels + await using var promptInterrupt = RegisterPromptInterrupt(task, cancellationToken); + + Logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); + await task.Task.WaitAsync(cancellationToken); + + var outputImages = await ClientManager.Client.GetImagesForExecutedPromptAsync( + task.Id, + cancellationToken + ); + + var (selectedOutputKey, candidateImages) = ComfyGenerationHelper.SelectOutputImages(outputImages); + + if (candidateImages is null || candidateImages.Count == 0) + { + Logger.LogWarning("No output images found from generation"); + return new ImageGenerationResponse + { + IsSuccess = false, + ErrorMessage = "No output images were generated", + }; + } + + var generatedImages = await DownloadImagesAsync(candidateImages, cancellationToken); + + Logger.LogInformation( + "Successfully generated {Count} image(s) with {Provider}", + generatedImages.Count, + LogName + ); + + return new ImageGenerationResponse + { + IsSuccess = true, + Images = generatedImages, + TextResponse = null, + Metadata = new Dictionary + { + ["promptId"] = task.Id, + ["outputNode"] = selectedOutputKey ?? "unknown", + }, + }; + } + catch (OperationCanceledException) + { + Logger.LogInformation("Image generation was cancelled"); + throw; // Propagate cancellation to the ViewModel for proper handling + } + catch (ComfyNodeException nodeEx) + { + // Execution-time node failure (e.g. tensor-shape mismatch from a wrong encoder + // pairing). Carries ComfyUI's full error JSON for the detail dialog. + return ComfyGenerationHelper.CreateNodeErrorResponse(nodeEx, LogName, Logger); + } + catch (ApiException apiEx) + { + // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. + return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, LogName, Logger); + } + catch (Exception ex) + { + Logger.LogError(ex, "Failed to generate image with {Provider}", LogName); + return new ImageGenerationResponse + { + IsSuccess = false, + ErrorMessage = $"Generation failed: {ex.Message}", + }; + } + } + + /// + /// Registers a cancellation callback that interrupts the running ComfyUI prompt. + /// + private CancellationTokenRegistration RegisterPromptInterrupt( + ComfyTask task, + CancellationToken cancellationToken + ) => + cancellationToken.Register(() => + { + Logger.LogInformation("Cancellation requested, interrupting ComfyUI prompt {PromptId}", task.Id); + // CTS holds an internal timer that needs disposing; chain dispose onto the + // fire-and-forget interrupt so it cleans up once the request settles. + var interruptCts = new CancellationTokenSource(5000); + ClientManager + .Client.InterruptPromptAsync(interruptCts.Token) + .ContinueWith( + t => + { + if (t.IsFaulted) + { + Logger.LogWarning( + t.Exception, + "Failed to interrupt ComfyUI prompt {PromptId}", + task.Id + ); + } + interruptCts.Dispose(); + }, + TaskScheduler.Default + ) + .SafeFireAndForget(); + }); + + /// + /// Downloads each output image from ComfyUI and converts it to a base64 GeneratedImage. + /// + private async Task> DownloadImagesAsync( + IReadOnlyList images, + CancellationToken cancellationToken + ) + { + var generatedImages = new List(); + + foreach (var comfyImage in images) + { + Logger.LogInformation("Downloading generated image: {FileName}", comfyImage.FileName); + + await using var imageStream = await ClientManager.Client.GetImageStreamAsync( + comfyImage, + cancellationToken + ); + using var memoryStream = new MemoryStream(); + await imageStream.CopyToAsync(memoryStream, cancellationToken); + var base64Image = Convert.ToBase64String(memoryStream.ToArray()); + + generatedImages.Add( + new GeneratedImage + { + Base64Data = base64Image, + MimeType = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName), + } + ); + } + + return generatedImages; + } + + /// + /// Reads the optional custom UNET model selection from the request options. + /// + /// Generation request + /// + /// Whether to log the selection. Variant-aware providers resolve the UNET during model + /// validation as well as workflow build, and pass false on the validation pass to + /// avoid logging the same selection twice. + /// + protected HybridModelFile? GetCustomUnetModel(ImageGenerationRequest request, bool logSelection = true) + { + if ( + request.ProviderOptions?.TryGetValue("CustomUnetModel", out var modelObj) == true + && modelObj is HybridModelFile model + ) + { + if (logSelection) + { + Logger.LogInformation("Using custom UNet model: {ModelPath}", model.RelativePath); + } + return model; + } + + return null; + } + + /// + /// Reads the optional LoRA selections from the request options. + /// + protected IReadOnlyList? GetSelectedLoras(ImageGenerationRequest request) + { + if ( + request.ProviderOptions?.TryGetValue("SelectedLoras", out var lorasObj) == true + && lorasObj is IEnumerable loraList + ) + { + var loras = loraList.ToList(); + Logger.LogInformation("Using {Count} LoRAs", loras.Count); + return loras; + } + + return null; + } + + /// + /// Reads the optional custom output dimensions from the request options. + /// + protected (int? Width, int? Height) GetDimensions(ImageGenerationRequest request) + { + int? width = null; + int? height = null; + + if (request.ProviderOptions?.TryGetValue("Width", out var widthObj) == true && widthObj is int w) + { + width = w; + } + + if (request.ProviderOptions?.TryGetValue("Height", out var heightObj) == true && heightObj is int h) + { + height = h; + } + + if (width.HasValue && height.HasValue) + { + Logger.LogInformation("Using custom resolution: {Width}x{Height}", width, height); + } + + return (width, height); + } +} diff --git a/StabilityMatrix.Avalonia/Services/Flux2KleinModelManager.cs b/StabilityMatrix.Avalonia/Services/Flux2KleinModelManager.cs index b49beeacb..011eda8af 100644 --- a/StabilityMatrix.Avalonia/Services/Flux2KleinModelManager.cs +++ b/StabilityMatrix.Avalonia/Services/Flux2KleinModelManager.cs @@ -18,23 +18,56 @@ public class Flux2KleinModelManager : ILocalProviderModelManager public string DownloadDialogDescription => "Flux.2 Klein requires the following models to run. You can deselect any you already have installed elsewhere:"; - public bool AreModelsAvailable(IInferenceClientManager clientManager) + public bool AreModelsAvailable(IInferenceClientManager clientManager) => + AreModelsAvailable(clientManager, preferredUnet: null); + + /// + /// Variant-aware availability check. The text encoder only counts as available when its + /// size matches the UNET it will be paired with ( when the + /// caller knows the user's selection, otherwise the installed Klein UNET) — a qwen_3_4b + /// encoder next to a 9B UNET loads fine but fails at sampling with a tensor-shape + /// mismatch, so it must be reported as missing rather than "ready". + /// + public bool AreModelsAvailable(IInferenceClientManager clientManager, HybridModelFile? preferredUnet) { var hasUnet = clientManager.UnetModels.Any(m => m.Local != null && IsKleinUnet(m.FileName)); var hasVae = clientManager.VaeModels.Any(m => m.Local != null && IsFlux2Vae(m.FileName)); - var hasClip = clientManager.ClipModels.Any(m => m.Local != null && IsKleinTextEncoder(m.FileName)); + + var encoderSize = GetInstalledKleinVariant(clientManager, preferredUnet); + var hasClip = clientManager.ClipModels.Any(m => + m.Local != null && IsKleinTextEncoder(m.FileName) && MatchesEncoderSize(m.FileName, encoderSize) + ); return hasUnet && hasVae && hasClip; } - public IEnumerable GetMissingModels(IInferenceClientManager clientManager) + public IEnumerable GetMissingModels(IInferenceClientManager clientManager) => + GetMissingModels(clientManager, preferredUnet: null); + + public IEnumerable GetMissingModels( + IInferenceClientManager clientManager, + HybridModelFile? preferredUnet + ) { - var allModels = RemoteModels.Flux2KleinModels; var missing = new List(); + // Pick the variant to offer downloads for based on the UNET the encoder will pair with + // (the user's dropdown selection when provided, otherwise whatever is installed). + // A 9B UNET pairs with the qwen_3_8b encoder, a 4B UNET with qwen_3_4b — offering the + // wrong size just wastes an 8-16 GB download and still fails to generate (the sampler + // hits a tensor-shape mismatch), so we match the encoder/VAE to that UNET. + var encoderSize = GetInstalledKleinVariant(clientManager, preferredUnet); + var variantModels = + encoderSize == "8b" ? RemoteModels.Flux2Klein9BModels : RemoteModels.Flux2KleinModels; + + // The 9B UNET is gated and can't be redistributed for one-click download, so we only + // ever auto-offer the Apache 2.0 4B UNET — and only when no Klein UNET is present at all. + // A user who already has the 9B UNET keeps it; we just fill in the encoder/VAE around it. if (!clientManager.UnetModels.Any(m => m.Local != null && IsKleinUnet(m.FileName))) { - var unet = allModels.FirstOrDefault(m => m.ContextType is SharedFolderType.DiffusionModels); + var unet = RemoteModels.Flux2KleinModels.FirstOrDefault(m => + m.ContextType is SharedFolderType.DiffusionModels + ); if (unet.Url != null) { missing.Add(unet); @@ -43,16 +76,24 @@ public IEnumerable GetMissingModels(IInferenceClientManager clie if (!clientManager.VaeModels.Any(m => m.Local != null && IsFlux2Vae(m.FileName))) { - var vae = allModels.FirstOrDefault(m => m.ContextType is SharedFolderType.VAE); + var vae = variantModels.FirstOrDefault(m => m.ContextType is SharedFolderType.VAE); if (vae.Url != null) { missing.Add(vae); } } - if (!clientManager.ClipModels.Any(m => m.Local != null && IsKleinTextEncoder(m.FileName))) + // Size-aware: an installed encoder of the WRONG size (e.g. qwen_3_4b next to a 9B + // UNET) still means the matching encoder is missing and should be offered. + if ( + !clientManager.ClipModels.Any(m => + m.Local != null + && IsKleinTextEncoder(m.FileName) + && MatchesEncoderSize(m.FileName, encoderSize) + ) + ) { - var clip = allModels.FirstOrDefault(m => m.ContextType is SharedFolderType.TextEncoders); + var clip = variantModels.FirstOrDefault(m => m.ContextType is SharedFolderType.TextEncoders); if (clip.Url != null) { missing.Add(clip); @@ -62,21 +103,51 @@ public IEnumerable GetMissingModels(IInferenceClientManager clie return missing; } - public IEnumerable GetMissingModelNames(IInferenceClientManager clientManager) + public IEnumerable GetMissingModelNames(IInferenceClientManager clientManager) => + GetMissingModelNames(clientManager, preferredUnet: null); + + public IEnumerable GetMissingModelNames( + IInferenceClientManager clientManager, + HybridModelFile? preferredUnet + ) { - foreach (var model in GetMissingModels(clientManager)) + var encoderSize = GetInstalledKleinVariant(clientManager, preferredUnet); + + foreach (var model in GetMissingModels(clientManager, preferredUnet)) { var name = model.ContextType switch { + // Only the 4B UNET is ever auto-offered (the 9B UNET is gated), so this label + // is always accurate. SharedFolderType.DiffusionModels => "Flux.2 Klein 4B UNET", SharedFolderType.VAE => "Flux.2 VAE", - SharedFolderType.TextEncoders => "Qwen3 4B text encoder", + SharedFolderType.TextEncoders => encoderSize == "8b" + ? "Qwen3 8B text encoder" + : "Qwen3 4B text encoder", _ => model.FileName, }; yield return name; } } + /// + /// Returns the encoder size ("8b" or "4b") implied by + /// (the user's dropdown selection, when the caller has one) or by the Klein UNET that's + /// currently installed, or "4b" when neither is present (the Apache 2.0 default for fresh + /// installs). Used to offer the matching text encoder rather than blindly pushing 4B. + /// + private static string GetInstalledKleinVariant( + IInferenceClientManager clientManager, + HybridModelFile? preferredUnet = null + ) + { + var unet = + preferredUnet + ?? clientManager.UnetModels.FirstOrDefault(m => m.Local != null && IsKleinUnet(m.FileName)); + + return unet != null ? GetExpectedEncoderSize(unet) : "4b"; + } + /// /// Select the best available models for Flux.2 Klein (only LOCAL models). /// When the selected UNET is the 9B variant, prefers the matching qwen_3_8b text encoder; @@ -125,7 +196,7 @@ internal SelectedModels SelectModels( /// community merges often don't include a "9b" / "4b" hint. Falls back to the filename, /// then defaults to "4b" (matches the auto-downloaded Apache 2.0 Klein 4B variant). /// - private static string GetExpectedEncoderSize(HybridModelFile unetModel) + internal static string GetExpectedEncoderSize(HybridModelFile unetModel) { var info = unetModel.Local?.ConnectedModelInfo; @@ -175,7 +246,7 @@ private static bool IsFourBSignal(string text) => || text.Contains("klein-4", StringComparison.OrdinalIgnoreCase) || text.Contains("klein_4", StringComparison.OrdinalIgnoreCase); - private static bool MatchesEncoderSize(string encoderFileName, string size) => + internal static bool MatchesEncoderSize(string encoderFileName, string size) => size switch { "8b" => encoderFileName.Contains("qwen_3_8b", StringComparison.OrdinalIgnoreCase) @@ -189,7 +260,7 @@ private static bool IsKleinUnet(string fileName) => fileName.Contains("flux-2-klein", StringComparison.OrdinalIgnoreCase) || fileName.Contains("flux2-klein", StringComparison.OrdinalIgnoreCase); - private static bool IsFlux2Vae(string fileName) => + internal static bool IsFlux2Vae(string fileName) => // Distilled variants use flux2-vae.safetensors; base variants use // full_encoder_small_decoder.safetensors. Both are valid Flux.2 VAEs. fileName.Contains("flux2-vae", StringComparison.OrdinalIgnoreCase) @@ -199,7 +270,7 @@ private static bool IsFlux2Vae(string fileName) => // Match the Qwen3 text encoders that Klein uses (qwen_3_4b for 4B model, qwen_3_8b for 9B). // Note: deliberately excludes Qwen 2.5 VL encoders used by Qwen Image Edit // (those have "vl" in the filename and use the older "_2.5_" version tag). - private static bool IsKleinTextEncoder(string fileName) => + internal static bool IsKleinTextEncoder(string fileName) => ( fileName.Contains("qwen_3_4b", StringComparison.OrdinalIgnoreCase) || fileName.Contains("qwen_3_8b", StringComparison.OrdinalIgnoreCase) diff --git a/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs index 8dabd48f5..b58e0dbd8 100644 --- a/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs +++ b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs @@ -1,10 +1,6 @@ -using AsyncAwaitBestPractices; using Microsoft.Extensions.Logging; -using Refit; -using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; -using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.Api.Comfy; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -15,339 +11,84 @@ namespace StabilityMatrix.Avalonia.Services; /// making it well-suited to conversational, iterative editing. /// public class Flux2KleinProvider(ILogger logger, IInferenceClientManager clientManager) - : IImageGenerationProvider + : ComfyImageGenerationProviderBase(logger, clientManager) { - public string ProviderId => BananaVisionProviderIds.Flux2Klein; - public string ProviderName => "Flux.2 Klein (Local)"; - public bool SupportsImageInput => true; - public bool SupportsMultiTurn => true; - public bool RequiresThoughtSignatures => false; + public override string ProviderId => BananaVisionProviderIds.Flux2Klein; + public override string ProviderName => "Flux.2 Klein (Local)"; - public async Task GenerateAsync( - ImageGenerationRequest request, - CancellationToken cancellationToken = default - ) - { - try - { - if (!clientManager.IsConnected) - { - logger.LogWarning("ComfyUI is not connected"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - "ComfyUI is not connected. Use the Launch button in the header to start and connect to ComfyUI.", - }; - } - - var modelManager = new Flux2KleinModelManager(); - if (!modelManager.AreModelsAvailable(clientManager)) - { - var modelsList = string.Join(", ", modelManager.GetMissingModelNames(clientManager)); - - logger.LogWarning("Required models not found: {Models}", modelsList); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - $"Required models not found: {modelsList}. Please download them from the HuggingFace model browser.", - }; - } - - // Klein supports multi-reference editing; cap at 4 for predictable VRAM use. - await ComfyImageUploadHelper.UploadImagesAsync( - clientManager, - request, - maxInputImages: 4, - providerPrefix: "flux2_klein", - logger, - cancellationToken - ); - - HybridModelFile? customUnetModel = null; - IEnumerable? loras = null; - int? width = null; - int? height = null; - int? steps = null; - double? cfg = null; - var explicitDimensions = false; - - if (request.ProviderOptions != null) - { - if ( - request.ProviderOptions.TryGetValue("CustomUnetModel", out var modelObj) - && modelObj is HybridModelFile model - ) - { - customUnetModel = model; - logger.LogInformation("Using custom UNet model: {ModelPath}", model.RelativePath); - } - - if ( - request.ProviderOptions.TryGetValue("SelectedLoras", out var lorasObj) - && lorasObj is IEnumerable loraList - ) - { - loras = loraList; - logger.LogInformation("Using {Count} LoRAs", loraList.Count()); - } - - if (request.ProviderOptions.TryGetValue("Width", out var widthObj) && widthObj is int w) - { - width = w; - } - - if (request.ProviderOptions.TryGetValue("Height", out var heightObj) && heightObj is int h) - { - height = h; - } - - if ( - request.ProviderOptions.TryGetValue("ExplicitDimensions", out var explicitObj) - && explicitObj is bool eb - ) - { - explicitDimensions = eb; - } - - if (request.ProviderOptions.TryGetValue("Steps", out var stepsObj) && stepsObj is int s) - { - steps = s; - } - - if (request.ProviderOptions.TryGetValue("CfgScale", out var cfgObj)) - { - cfg = cfgObj switch - { - double d => d, - float f => f, - int i => i, - _ => null, - }; - } - - if (width.HasValue && height.HasValue) - { - logger.LogInformation("Using custom resolution: {Width}x{Height}", width, height); - } - if (steps.HasValue || cfg.HasValue) - { - logger.LogInformation("Using Klein overrides: Steps={Steps}, Cfg={Cfg}", steps, cfg); - } - } - - logger.LogInformation("Building Flux.2 Klein workflow"); - var nodes = Flux2KleinWorkflowBuilder.Build( - request, - clientManager, - customUnetModel, - loras, - width, - height, - steps, - cfg, - explicitDimensions - ); - - logger.LogInformation("Queuing prompt to ComfyUI"); - var task = await clientManager.Client.QueuePromptAsync(nodes, cancellationToken); - - request.Progress?.Report( - new ImageGenerationProgress( - ProviderId, - task.Id, - Value: null, - Maximum: null, - RunningNode: null, - Stage: "Queued" - ) - ); - - int? lastPercent = null; - string? lastRunningNode = null; - - void ReportProgress(int? value, int? maximum, string? runningNode, string? stage) - { - int? percent = value is >= 0 && maximum is > 0 ? (value.Value * 100) / maximum.Value : null; - - if ( - percent == lastPercent - && string.Equals(lastRunningNode, runningNode, StringComparison.Ordinal) - ) - { - return; - } + protected override string LogName => "Flux.2 Klein"; - lastPercent = percent; - lastRunningNode = runningNode; + // Klein supports multi-reference editing; cap at 4 for predictable VRAM use. + protected override int MaxInputImages => 4; + protected override string ProviderPrefix => "flux2_klein"; - request.Progress?.Report( - new ImageGenerationProgress(ProviderId, task.Id, value, maximum, runningNode, stage) - ); - } - - void OnProgressUpdate( - object? sender, - StabilityMatrix.Core.Inference.ComfyProgressUpdateEventArgs args - ) - { - ReportProgress(args.Value, args.Maximum, args.RunningNode, "Generating"); - } - - void OnRunningNodeChanged(object? sender, string? node) - { - ReportProgress( - task.LastProgressUpdate?.Value, - task.LastProgressUpdate?.Maximum, - node, - "Generating" - ); - } - - task.ProgressUpdate += OnProgressUpdate; - task.RunningNodeChanged += OnRunningNodeChanged; - - await using var promptInterrupt = cancellationToken.Register(() => - { - logger.LogInformation( - "Cancellation requested, interrupting ComfyUI prompt {PromptId}", - task.Id - ); - // CTS holds an internal timer that needs disposing; chain dispose onto - // the fire-and-forget interrupt so it cleans up once the request settles. - var interruptCts = new CancellationTokenSource(5000); - clientManager - .Client.InterruptPromptAsync(interruptCts.Token) - .ContinueWith(_ => interruptCts.Dispose(), TaskScheduler.Default) - .SafeFireAndForget(); - }); + protected override IReadOnlyList GetMissingModels(ImageGenerationRequest request) + { + // Resolve the user's UNET selection first — the availability check is variant-aware + // (a 9B UNET needs the qwen_3_8b encoder, 4B needs qwen_3_4b), so it has to know which + // UNET the workflow will actually use. Logging happens on the build pass instead. + var customUnetModel = GetCustomUnetModel(request, logSelection: false); - try - { - logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); - await task.Task.WaitAsync(cancellationToken); - } - finally - { - task.ProgressUpdate -= OnProgressUpdate; - task.RunningNodeChanged -= OnRunningNodeChanged; - } + var modelManager = new Flux2KleinModelManager(); + if (modelManager.AreModelsAvailable(ClientManager, customUnetModel)) + { + return []; + } + return modelManager.GetMissingModelNames(ClientManager, customUnetModel).ToList(); + } - var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( - task.Id, - cancellationToken - ); + protected override Dictionary BuildWorkflow(ImageGenerationRequest request) + { + var customUnetModel = GetCustomUnetModel(request); + var loras = GetSelectedLoras(request); + var (width, height) = GetDimensions(request); - var preferredOutputKey = "SaveImage"; - string? selectedOutputKey = null; - List? candidateImages = null; + int? steps = null; + double? cfg = null; + var explicitDimensions = false; + if (request.ProviderOptions != null) + { if ( - outputImages.TryGetValue(preferredOutputKey, out var preferredImages) - && preferredImages is { Count: > 0 } + request.ProviderOptions.TryGetValue("ExplicitDimensions", out var explicitObj) + && explicitObj is bool eb ) { - selectedOutputKey = preferredOutputKey; - candidateImages = preferredImages; + explicitDimensions = eb; } - else - { - var selected = outputImages - .OrderBy(kvp => kvp.Key, StringComparer.Ordinal) - .FirstOrDefault(kvp => kvp.Value is { Count: > 0 }); - selectedOutputKey = string.IsNullOrEmpty(selected.Key) ? null : selected.Key; - candidateImages = selected.Value; + if (request.ProviderOptions.TryGetValue("Steps", out var stepsObj) && stepsObj is int s) + { + steps = s; } - if (candidateImages is null || candidateImages.Count == 0) + if (request.ProviderOptions.TryGetValue("CfgScale", out var cfgObj)) { - logger.LogWarning("No output images found from generation"); - return new ImageGenerationResponse + cfg = cfgObj switch { - IsSuccess = false, - ErrorMessage = "No output images were generated", + double d => d, + float f => f, + int i => i, + _ => null, }; } - var generatedImages = new List(); - - foreach (var comfyImage in candidateImages) + if (steps.HasValue || cfg.HasValue) { - logger.LogInformation("Downloading generated image: {FileName}", comfyImage.FileName); - - await using var imageStream = await clientManager.Client.GetImageStreamAsync( - comfyImage, - cancellationToken - ); - using var memoryStream = new MemoryStream(); - await imageStream.CopyToAsync(memoryStream, cancellationToken); - var imageBytes = memoryStream.ToArray(); - var base64Image = Convert.ToBase64String(imageBytes); - - var mimeType = - comfyImage.FileName.EndsWith(".png", StringComparison.OrdinalIgnoreCase) ? "image/png" - : comfyImage.FileName.EndsWith(".jpg", StringComparison.OrdinalIgnoreCase) - || comfyImage.FileName.EndsWith(".jpeg", StringComparison.OrdinalIgnoreCase) - ? "image/jpeg" - : comfyImage.FileName.EndsWith(".webp", StringComparison.OrdinalIgnoreCase) ? "image/webp" - : "image/png"; - - generatedImages.Add(new GeneratedImage { Base64Data = base64Image, MimeType = mimeType }); + Logger.LogInformation("Using Klein overrides: Steps={Steps}, Cfg={Cfg}", steps, cfg); } - - logger.LogInformation( - "Successfully generated {Count} image(s) with Flux.2 Klein", - generatedImages.Count - ); - - return new ImageGenerationResponse - { - IsSuccess = true, - Images = generatedImages, - TextResponse = null, - Metadata = new Dictionary - { - ["promptId"] = task.Id, - ["outputNode"] = selectedOutputKey ?? "unknown", - }, - }; } - catch (OperationCanceledException) - { - logger.LogInformation("Image generation was cancelled"); - throw; - } - catch (ApiException apiEx) - { - // ComfyUI returns a JSON body explaining which node validation failed when the - // workflow is rejected; surfacing that beats a generic 400 message by miles. - var detail = !string.IsNullOrWhiteSpace(apiEx.Content) ? apiEx.Content : apiEx.Message; - logger.LogError( - apiEx, - "ComfyUI rejected Flux.2 Klein workflow ({StatusCode}): {Detail}", - apiEx.StatusCode, - detail - ); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {Truncate(detail, 800)}", - }; - } - catch (Exception ex) - { - logger.LogError(ex, "Failed to generate image with Flux.2 Klein"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = $"Generation failed: {ex.Message}", - }; - } - } - private static string Truncate(string value, int maxLength) => - value.Length <= maxLength ? value : value[..maxLength] + "..."; + return Flux2KleinWorkflowBuilder.Build( + request, + ClientManager, + customUnetModel, + loras, + width, + height, + steps, + cfg, + explicitDimensions + ); + } } diff --git a/StabilityMatrix.Avalonia/Services/Flux2KleinWorkflowBuilder.cs b/StabilityMatrix.Avalonia/Services/Flux2KleinWorkflowBuilder.cs index f6a3d5750..656e42df9 100644 --- a/StabilityMatrix.Avalonia/Services/Flux2KleinWorkflowBuilder.cs +++ b/StabilityMatrix.Avalonia/Services/Flux2KleinWorkflowBuilder.cs @@ -1,4 +1,5 @@ using SkiaSharp; +using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.Api.Comfy.Nodes; @@ -138,29 +139,12 @@ public static Dictionary Build( ); // Apply LoRAs if any - var currentModel = unetOutput; - var currentClip = clipLoader.Output; - - var loraList = loras?.ToList() ?? []; - if (loraList.Count > 0) - { - for (var i = 0; i < loraList.Count; i++) - { - var lora = loraList[i]; - var loraLoader = nodes.AddNamedNode( - ComfyNodeBuilder.LoraLoader( - nodes.GetUniqueName($"LoraLoader_{i + 1}"), - currentModel, - currentClip, - lora.Model.RelativePath, - (double)lora.ModelWeight, - (double)lora.ClipWeight - ) - ); - currentModel = loraLoader.Output1; - currentClip = loraLoader.Output2; - } - } + var (currentModel, currentClip) = ComfyWorkflowHelper.ApplyLoras( + nodes, + loras, + unetOutput, + clipLoader.Output + ); // 2. Encode the positive prompt var positivePrompt = request.TextPrompt ?? "a beautiful image"; @@ -193,7 +177,13 @@ public static Dictionary Build( var positiveConditioning = positiveTextEncode.Output; var negativeConditioning = negativeTextEncode.Output; - var referenceImageNames = GetReferenceImageNames(request); + // Klein supports multi-reference editing; cap at 4 to keep prompt build time + // and VRAM use predictable + var referenceImageNames = ComfyWorkflowHelper.GetReferenceImageNames( + request, + maxImages: 4, + providerPrefix: "flux2_klein" + ); for (var i = 0; i < referenceImageNames.Count; i++) { var imageName = referenceImageNames[i]; @@ -397,37 +387,4 @@ int step return (w, h); } - - /// - /// Get reference image filenames (uploaded inputs + most recent history image). - /// Klein supports multi-reference editing; we cap at 4 to keep prompt build time - /// and VRAM use predictable. - /// - private static List GetReferenceImageNames(ImageGenerationRequest request) - { - const int maxRefs = 4; - var imageNames = new List(); - - if (request.InputImages?.Count > 0) - { - for (var i = 0; i < Math.Min(request.InputImages.Count, maxRefs); i++) - { - imageNames.Add($"Inference/flux2_klein_input_{i}.png"); - } - } - - if (imageNames.Count < maxRefs && request.ConversationHistory != null) - { - var lastAssistantImage = request.ConversationHistory.LastOrDefault(m => - m is { Role: MessageRole.Assistant, ImageContent: not null } - ); - - if (lastAssistantImage?.ImageContent != null) - { - imageNames.Add("Inference/flux2_klein_history_latest.png"); - } - } - - return imageNames; - } } diff --git a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs index 00aaf3672..0ee87e744 100644 --- a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs +++ b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs @@ -1,9 +1,6 @@ -using AsyncAwaitBestPractices; using Microsoft.Extensions.Logging; -using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; -using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.Api.Comfy; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -12,293 +9,38 @@ namespace StabilityMatrix.Avalonia.Services; /// Image generation provider for Flux Kontext using local ComfyUI backend /// public class FluxKontextProvider(ILogger logger, IInferenceClientManager clientManager) - : IImageGenerationProvider + : ComfyImageGenerationProviderBase(logger, clientManager) { - public string ProviderId => BananaVisionProviderIds.FluxKontext; - public string ProviderName => "Flux Kontext (Local)"; - public bool SupportsImageInput => true; - public bool SupportsMultiTurn => true; - public bool RequiresThoughtSignatures => false; + public override string ProviderId => BananaVisionProviderIds.FluxKontext; + public override string ProviderName => "Flux Kontext (Local)"; - public async Task GenerateAsync( - ImageGenerationRequest request, - CancellationToken cancellationToken = default - ) - { - try - { - // Check if ComfyUI is connected - if (!clientManager.IsConnected) - { - logger.LogWarning("ComfyUI is not connected"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - "ComfyUI is not connected. Use the Launch button in the header to start and connect to ComfyUI.", - }; - } - - // Validate models are available - var modelManager = new FluxKontextModelManager(); - if (!modelManager.AreModelsAvailable(clientManager)) - { - var modelsList = string.Join(", ", modelManager.GetMissingModelNames(clientManager)); - - logger.LogWarning("Required models not found: {Models}", modelsList); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - $"Required models not found: {modelsList}. Please download them from the HuggingFace model browser.", - }; - } - - // Upload images using helper - await ComfyImageUploadHelper.UploadImagesAsync( - clientManager, - request, - maxInputImages: 2, - providerPrefix: "flux_kontext", - logger, - cancellationToken - ); - - // Extract custom model, LoRA selections, and resolution from provider options - HybridModelFile? customUnetModel = null; - IEnumerable? loras = null; - int? width = null; - int? height = null; - - if (request.ProviderOptions != null) - { - if ( - request.ProviderOptions.TryGetValue("CustomUnetModel", out var modelObj) - && modelObj is HybridModelFile model - ) - { - customUnetModel = model; - logger.LogInformation("Using custom UNet model: {ModelPath}", model.RelativePath); - } - - if ( - request.ProviderOptions.TryGetValue("SelectedLoras", out var lorasObj) - && lorasObj is IEnumerable loraList - ) - { - loras = loraList; - logger.LogInformation("Using {Count} LoRAs", loraList.Count()); - } - - if (request.ProviderOptions.TryGetValue("Width", out var widthObj) && widthObj is int w) - { - width = w; - } - - if (request.ProviderOptions.TryGetValue("Height", out var heightObj) && heightObj is int h) - { - height = h; - } - - if (width.HasValue && height.HasValue) - { - logger.LogInformation("Using custom resolution: {Width}x{Height}", width, height); - } - } - - // Build workflow nodes - logger.LogInformation("Building Flux Kontext workflow"); - var nodes = FluxKontextWorkflowBuilder.Build( - request, - clientManager, - customUnetModel, - loras, - width, - height - ); - - // Queue the prompt - logger.LogInformation("Queuing prompt to ComfyUI"); - var task = await clientManager.Client.QueuePromptAsync(nodes, cancellationToken); - - request.Progress?.Report( - new ImageGenerationProgress( - ProviderId, - task.Id, - Value: null, - Maximum: null, - RunningNode: null, - Stage: "Queued" - ) - ); - - int? lastPercent = null; - string? lastRunningNode = null; - - void ReportProgress(int? value, int? maximum, string? runningNode, string? stage) - { - int? percent = value is >= 0 && maximum is > 0 ? (value.Value * 100) / maximum.Value : null; - - if ( - percent == lastPercent - && string.Equals(lastRunningNode, runningNode, StringComparison.Ordinal) - ) - { - return; - } - - lastPercent = percent; - lastRunningNode = runningNode; - - request.Progress?.Report( - new ImageGenerationProgress(ProviderId, task.Id, value, maximum, runningNode, stage) - ); - } - - void OnProgressUpdate( - object? sender, - StabilityMatrix.Core.Inference.ComfyProgressUpdateEventArgs args - ) - { - ReportProgress(args.Value, args.Maximum, args.RunningNode, "Generating"); - } - - void OnRunningNodeChanged(object? sender, string? node) - { - ReportProgress( - task.LastProgressUpdate?.Value, - task.LastProgressUpdate?.Maximum, - node, - "Generating" - ); - } + protected override string LogName => "Flux Kontext"; + protected override int MaxInputImages => 2; + protected override string ProviderPrefix => "flux_kontext"; - task.ProgressUpdate += OnProgressUpdate; - task.RunningNodeChanged += OnRunningNodeChanged; - - // Register cancellation to interrupt ComfyUI if the user cancels - await using var promptInterrupt = cancellationToken.Register(() => - { - logger.LogInformation( - "Cancellation requested, interrupting ComfyUI prompt {PromptId}", - task.Id - ); - // CTS holds an internal timer that needs disposing; chain dispose onto - // the fire-and-forget interrupt so it cleans up once the request settles. - var interruptCts = new CancellationTokenSource(5000); - clientManager - .Client.InterruptPromptAsync(interruptCts.Token) - .ContinueWith(_ => interruptCts.Dispose(), TaskScheduler.Default) - .SafeFireAndForget(); - }); - - try - { - // Wait for completion - logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); - await task.Task.WaitAsync(cancellationToken); - } - finally - { - task.ProgressUpdate -= OnProgressUpdate; - task.RunningNodeChanged -= OnRunningNodeChanged; - } - - // Get the output images - var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( - task.Id, - cancellationToken - ); - - // Prefer the "SaveImage" output node deterministically. - var preferredOutputKey = "SaveImage"; - string? selectedOutputKey = null; - List? candidateImages = null; - - if ( - outputImages.TryGetValue(preferredOutputKey, out var preferredImages) - && preferredImages is { Count: > 0 } - ) - { - selectedOutputKey = preferredOutputKey; - candidateImages = preferredImages; - } - else - { - var selected = outputImages - .OrderBy(kvp => kvp.Key, StringComparer.Ordinal) - .FirstOrDefault(kvp => kvp.Value is { Count: > 0 }); - - selectedOutputKey = string.IsNullOrEmpty(selected.Key) ? null : selected.Key; - candidateImages = selected.Value; - } - - if (candidateImages is null || candidateImages.Count == 0) - { - logger.LogWarning("No output images found from generation"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = "No output images were generated", - }; - } - - var generatedImages = new List(); - - foreach (var comfyImage in candidateImages) - { - logger.LogInformation("Downloading generated image: {FileName}", comfyImage.FileName); - - await using var imageStream = await clientManager.Client.GetImageStreamAsync( - comfyImage, - cancellationToken - ); - using var memoryStream = new MemoryStream(); - await imageStream.CopyToAsync(memoryStream, cancellationToken); - var imageBytes = memoryStream.ToArray(); - var base64Image = Convert.ToBase64String(imageBytes); - - var mimeType = - comfyImage.FileName.EndsWith(".png", StringComparison.OrdinalIgnoreCase) ? "image/png" - : comfyImage.FileName.EndsWith(".jpg", StringComparison.OrdinalIgnoreCase) - || comfyImage.FileName.EndsWith(".jpeg", StringComparison.OrdinalIgnoreCase) - ? "image/jpeg" - : comfyImage.FileName.EndsWith(".webp", StringComparison.OrdinalIgnoreCase) ? "image/webp" - : "image/png"; - - generatedImages.Add(new GeneratedImage { Base64Data = base64Image, MimeType = mimeType }); - } - - logger.LogInformation( - "Successfully generated {Count} image(s) with Flux Kontext", - generatedImages.Count - ); - - return new ImageGenerationResponse - { - IsSuccess = true, - Images = generatedImages, - TextResponse = null, - Metadata = new Dictionary - { - ["promptId"] = task.Id, - ["outputNode"] = selectedOutputKey ?? "unknown", - }, - }; - } - catch (OperationCanceledException) - { - logger.LogInformation("Image generation was cancelled"); - throw; // Propagate cancellation to ViewModel for proper handling - } - catch (Exception ex) + protected override IReadOnlyList GetMissingModels(ImageGenerationRequest request) + { + var modelManager = new FluxKontextModelManager(); + if (modelManager.AreModelsAvailable(ClientManager)) { - logger.LogError(ex, "Failed to generate image with Flux Kontext"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = $"Generation failed: {ex.Message}", - }; + return []; } + return modelManager.GetMissingModelNames(ClientManager).ToList(); + } + + protected override Dictionary BuildWorkflow(ImageGenerationRequest request) + { + var customUnetModel = GetCustomUnetModel(request); + var loras = GetSelectedLoras(request); + var (width, height) = GetDimensions(request); + + return FluxKontextWorkflowBuilder.Build( + request, + ClientManager, + customUnetModel, + loras, + width, + height + ); } } diff --git a/StabilityMatrix.Avalonia/Services/FluxKontextWorkflowBuilder.cs b/StabilityMatrix.Avalonia/Services/FluxKontextWorkflowBuilder.cs index 4bca4a82a..1491142cc 100644 --- a/StabilityMatrix.Avalonia/Services/FluxKontextWorkflowBuilder.cs +++ b/StabilityMatrix.Avalonia/Services/FluxKontextWorkflowBuilder.cs @@ -1,3 +1,4 @@ +using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.Api.Comfy.Nodes; @@ -89,29 +90,12 @@ public static Dictionary Build( ); // Apply LoRAs if any - var currentModel = unetOutput; - var currentClip = clipLoader.Output; - - var loraList = loras?.ToList() ?? []; - if (loraList.Count > 0) - { - for (var i = 0; i < loraList.Count; i++) - { - var lora = loraList[i]; - var loraLoader = nodes.AddNamedNode( - ComfyNodeBuilder.LoraLoader( - nodes.GetUniqueName($"LoraLoader_{i + 1}"), - currentModel, - currentClip, - lora.Model.RelativePath, - (double)lora.ModelWeight, - (double)lora.ClipWeight - ) - ); - currentModel = loraLoader.Output1; - currentClip = loraLoader.Output2; - } - } + var (currentModel, currentClip) = ComfyWorkflowHelper.ApplyLoras( + nodes, + loras, + unetOutput, + clipLoader.Output + ); // 2. Encode text prompt var positivePrompt = request.TextPrompt ?? "a beautiful image"; @@ -130,7 +114,11 @@ public static Dictionary Build( ConditioningNodeConnection conditioningForGuidance; // Collect reference images (from input and conversation history) - var referenceImageNames = GetReferenceImageNames(request); + var referenceImageNames = ComfyWorkflowHelper.GetReferenceImageNames( + request, + maxImages: 2, + providerPrefix: "flux_kontext" + ); if (referenceImageNames.Count > 0) { @@ -324,44 +312,6 @@ public static Dictionary Build( return nodes; } - /// - /// Get reference image filenames from the request (input images + conversation history) - /// Note: Images uploaded via ComfyClient.UploadImageAsync go to the "Inference" subfolder - /// - private static List GetReferenceImageNames(ImageGenerationRequest request) - { - var imageNames = new List(); - - // Priority 1: Current input images (will be uploaded with known names) - if (request.InputImages?.Count > 0) - { - for (var i = 0; i < Math.Min(request.InputImages.Count, 2); i++) // Max 2 images - { - // Include the "Inference/" subfolder prefix since that's where UploadImageAsync uploads to - imageNames.Add($"Inference/flux_kontext_input_{i}.png"); - } - } - - // Priority 2: Most recent image from conversation history (previous generation) - // Only include if we actually have the image content to upload - if (imageNames.Count < 2 && request.ConversationHistory != null) - { - var lastAssistantImage = request.ConversationHistory.LastOrDefault(m => - m is { Role: MessageRole.Assistant, ImageContent: not null } - ); - - // Only add the filename if we found an image with content - // (the provider will have uploaded it) - if (lastAssistantImage?.ImageContent != null) - { - // Include the "Inference/" subfolder prefix - imageNames.Add("Inference/flux_kontext_history_latest.png"); - } - } - - return imageNames; - } - /// /// Selected models for Flux Kontext /// diff --git a/StabilityMatrix.Avalonia/Services/ModelImportService.cs b/StabilityMatrix.Avalonia/Services/ModelImportService.cs index f782f4356..280122fbd 100644 --- a/StabilityMatrix.Avalonia/Services/ModelImportService.cs +++ b/StabilityMatrix.Avalonia/Services/ModelImportService.cs @@ -274,27 +274,37 @@ public async Task DoCustomImport( Action? configureDownload = null ) { - // Subfolder support: if the file name contains a path separator (e.g. from a - // user-defined FileNameFormat pattern like "{base_model}/{file_name}"), split off - // the directory portion and join it onto the download folder. Matches DoImport. - if (modelFileName.Contains('/') || modelFileName.Contains('\\')) + // Subfolder support for user-defined patterns such as + // "{base_model}/{model_name}/{file_name}". Treat every component as relative so + // rooted or traversal input cannot escape the selected models folder. + var pathSegments = modelFileName + .Split(['/', '\\'], StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries) + .Where(segment => segment is not "." and not "..") + .Select(SanitizePathSegment) + .Where(segment => !string.IsNullOrWhiteSpace(segment)) + .ToArray(); + + if (pathSegments.Length == 0) { - var lastIndex = modelFileName.LastIndexOfAny(['/', '\\']); - if (lastIndex >= 0) - { - var folderPath = modelFileName.Substring(0, lastIndex); - modelFileName = modelFileName.Substring(lastIndex + 1); - downloadFolder = downloadFolder.JoinDir(folderPath); - } + throw new ArgumentException( + "Model file name must contain a valid file name.", + nameof(modelFileName) + ); + } + + modelFileName = pathSegments[^1]; + if (pathSegments.Length > 1) + { + downloadFolder = new DirectoryPath( + [downloadFolder.FullPath, .. pathSegments.Take(pathSegments.Length - 1)] + ); } // Folders might be missing if user didn't install any packages yet downloadFolder.Create(); // Fix invalid chars in FileName - var modelBaseFileName = Path.GetFileNameWithoutExtension(modelFileName); - modelBaseFileName = Path.GetInvalidFileNameChars() - .Aggregate(modelBaseFileName, (current, c) => current.Replace(c, '_')); + var modelBaseFileName = SanitizePathSegment(Path.GetFileNameWithoutExtension(modelFileName)); var modelFileExtension = Path.GetExtension(modelFileName); var downloadPath = downloadFolder.JoinFile(modelBaseFileName + modelFileExtension); @@ -357,6 +367,9 @@ public async Task DoCustomImport( } } + private static string SanitizePathSegment(string segment) => + Path.GetInvalidFileNameChars().Aggregate(segment, (current, c) => current.Replace(c, '_')); + private async Task DownloadPreviewImageAsync(Uri previewImageUri, FilePath previewImageDownloadPath) { await notificationService.TryAsync( diff --git a/StabilityMatrix.Avalonia/Services/ModelOrganizationService.cs b/StabilityMatrix.Avalonia/Services/ModelOrganizationService.cs index a0b2efbac..d1d7cd250 100644 --- a/StabilityMatrix.Avalonia/Services/ModelOrganizationService.cs +++ b/StabilityMatrix.Avalonia/Services/ModelOrganizationService.cs @@ -68,6 +68,28 @@ public ModelOrganizationPlan BuildPlan( }; } + /// + /// Moves a single model file into , keeping its file + /// name and bringing the .cm-info.json / preview / .yaml sidecars along. Rolls back any + /// already-moved sidecars on failure. Throws when + /// a destination file already exists. + /// + public async Task MoveModelFileAsync(LocalModelFile model, string modelsRoot, string destinationDirectory) + { + var sourcePath = model.GetFullPath(modelsRoot); + if (!File.Exists(sourcePath)) + { + throw new FileNotFoundException("Model file no longer exists", sourcePath); + } + + Directory.CreateDirectory(destinationDirectory); + + var targetPath = Path.Combine(destinationDirectory, Path.GetFileName(sourcePath)); + var moves = BuildFileMoves(sourcePath, targetPath); + + await ApplyFileMovesAsync(moves).ConfigureAwait(false); + } + public async Task ApplyPlan(ModelOrganizationPlan plan) { var movedCount = 0; diff --git a/StabilityMatrix.Avalonia/Services/QwenImageEditModelManager.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditModelManager.cs index db9f240b8..de7f64ed8 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditModelManager.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditModelManager.cs @@ -39,8 +39,7 @@ public bool AreModelsAvailable(IInferenceClientManager clientManager) var hasClip = clientManager.ClipModels.Any(m => m.Local != null && // Only check LOCAL models - m.FileName.Contains("qwen", StringComparison.OrdinalIgnoreCase) - && m.FileName.Contains("vl", StringComparison.OrdinalIgnoreCase) + IsUsableQwenVlEncoder(m.FileName) ); return hasUnet && hasVae && hasClip; @@ -86,14 +85,9 @@ public IEnumerable GetMissingModels(IInferenceClientManager clie } } - // Check for Qwen CLIP model (only LOCAL models) - if ( - !clientManager.ClipModels.Any(m => - m.Local != null - && m.FileName.Contains("qwen", StringComparison.OrdinalIgnoreCase) - && m.FileName.Contains("vl", StringComparison.OrdinalIgnoreCase) - ) - ) + // Check for Qwen CLIP model (only LOCAL models). Wrong-size VL encoders (2B/3B) + // don't count — they fail at sampling — so the 7B download is still offered. + if (!clientManager.ClipModels.Any(m => m.Local != null && IsUsableQwenVlEncoder(m.FileName))) { var clipModel = allModels.FirstOrDefault(m => m.ContextType is SharedFolderType.TextEncoders); if (clipModel.Url != null) @@ -144,16 +138,45 @@ internal SelectedModels SelectModels(IInferenceClientManager clientManager) m.Local != null && m.FileName.Contains("qwen_image_vae", StringComparison.OrdinalIgnoreCase) ) ?? throw new InvalidOperationException("Qwen Image VAE model not found"); + // Prefer an explicit 7B match, then any VL encoder without a size hint (could be a + // renamed 7B). Smaller VL encoders (e.g. the 3B, hidden size 2048) load fine but die + // mid-sampling with "expected input with shape [*, 3584]", so fail fast with a clear + // message instead of silently picking one. var clipModel = clientManager.ClipModels.FirstOrDefault(m => m.Local != null - && m.FileName.Contains("qwen", StringComparison.OrdinalIgnoreCase) - && m.FileName.Contains("vl", StringComparison.OrdinalIgnoreCase) - ) ?? throw new InvalidOperationException("Qwen 2.5 VL CLIP model not found"); + && IsQwenVlEncoder(m.FileName) + && m.FileName.Contains("7b", StringComparison.OrdinalIgnoreCase) + ) + ?? clientManager.ClipModels.FirstOrDefault(m => + m.Local != null && IsUsableQwenVlEncoder(m.FileName) + ) + ?? throw new InvalidOperationException( + clientManager.ClipModels.Any(m => m.Local != null && IsQwenVlEncoder(m.FileName)) + ? "Qwen Image Edit requires the Qwen2.5-VL 7B text encoder, but only a different-size " + + "VL encoder was found (smaller variants fail with a tensor shape mismatch). " + + "Download qwen_2.5_vl_7b_fp8_scaled.safetensors from " + + "huggingface.co/Comfy-Org/Qwen-Image_ComfyUI and place it in your TextEncoders folder." + : "Qwen 2.5 VL CLIP model not found" + ); return new SelectedModels(unetModel, vaeModel, clipModel); } + private static bool IsQwenVlEncoder(string fileName) => + fileName.Contains("qwen", StringComparison.OrdinalIgnoreCase) + && fileName.Contains("vl", StringComparison.OrdinalIgnoreCase); + + // Qwen Image Edit pairs with the Qwen2.5-VL **7B** encoder (hidden size 3584). Other + // sizes produce "Given normalized_shape=[3584], expected input with shape [*, 3584]" + // deep in the sampler, so they are treated as not installed. Files without any size + // hint are accepted (likely a renamed 7B). + private static readonly string[] WrongVlEncoderSizeHints = ["2b", "3b", "32b", "72b"]; + + private static bool IsUsableQwenVlEncoder(string fileName) => + IsQwenVlEncoder(fileName) + && !WrongVlEncoderSizeHints.Any(hint => fileName.Contains(hint, StringComparison.OrdinalIgnoreCase)); + /// /// Selected models for Qwen Image Edit /// diff --git a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs index b23f32fc0..bc98d85ed 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs @@ -1,9 +1,6 @@ -using AsyncAwaitBestPractices; using Microsoft.Extensions.Logging; -using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; -using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.Api.Comfy; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -14,293 +11,38 @@ namespace StabilityMatrix.Avalonia.Services; public class QwenImageEditProvider( ILogger logger, IInferenceClientManager clientManager -) : IImageGenerationProvider +) : ComfyImageGenerationProviderBase(logger, clientManager) { - public string ProviderId => BananaVisionProviderIds.QwenImageEdit; - public string ProviderName => "Qwen Image Edit (Local)"; - public bool SupportsImageInput => true; - public bool SupportsMultiTurn => true; - public bool RequiresThoughtSignatures => false; + public override string ProviderId => BananaVisionProviderIds.QwenImageEdit; + public override string ProviderName => "Qwen Image Edit (Local)"; - public async Task GenerateAsync( - ImageGenerationRequest request, - CancellationToken cancellationToken = default - ) - { - try - { - // Check if ComfyUI is connected - if (!clientManager.IsConnected) - { - logger.LogWarning("ComfyUI is not connected"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - "ComfyUI is not connected. Use the Launch button in the header to start and connect to ComfyUI.", - }; - } - - // Validate models are available - var modelManager = new QwenImageEditModelManager(); - if (!modelManager.AreModelsAvailable(clientManager)) - { - var modelsList = string.Join(", ", modelManager.GetMissingModelNames(clientManager)); - - logger.LogWarning("Required models not found: {Models}", modelsList); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - $"Required models not found: {modelsList}. Please download them from the HuggingFace model browser.", - }; - } - - // Upload images using helper - await ComfyImageUploadHelper.UploadImagesAsync( - clientManager, - request, - maxInputImages: 3, - providerPrefix: "qwen_image_edit", - logger, - cancellationToken - ); - - // Extract custom model, LoRA selections, and resolution from provider options - HybridModelFile? customUnetModel = null; - IEnumerable? loras = null; - int? width = null; - int? height = null; - - if (request.ProviderOptions != null) - { - if ( - request.ProviderOptions.TryGetValue("CustomUnetModel", out var modelObj) - && modelObj is HybridModelFile model - ) - { - customUnetModel = model; - logger.LogInformation("Using custom UNet model: {ModelPath}", model.RelativePath); - } - - if ( - request.ProviderOptions.TryGetValue("SelectedLoras", out var lorasObj) - && lorasObj is IEnumerable loraList - ) - { - loras = loraList; - logger.LogInformation("Using {Count} LoRAs", loraList.Count()); - } - - if (request.ProviderOptions.TryGetValue("Width", out var widthObj) && widthObj is int w) - { - width = w; - } - - if (request.ProviderOptions.TryGetValue("Height", out var heightObj) && heightObj is int h) - { - height = h; - } - - if (width.HasValue && height.HasValue) - { - logger.LogInformation("Using custom resolution: {Width}x{Height}", width, height); - } - } - - // Build workflow nodes - logger.LogInformation("Building Qwen Image Edit workflow"); - var nodes = QwenImageEditWorkflowBuilder.Build( - request, - clientManager, - customUnetModel, - loras, - width, - height - ); - - // Queue the prompt - logger.LogInformation("Queuing prompt to ComfyUI"); - var task = await clientManager.Client.QueuePromptAsync(nodes, cancellationToken); - - request.Progress?.Report( - new ImageGenerationProgress( - ProviderId, - task.Id, - Value: null, - Maximum: null, - RunningNode: null, - Stage: "Queued" - ) - ); - - int? lastPercent = null; - string? lastRunningNode = null; - - void ReportProgress(int? value, int? maximum, string? runningNode, string? stage) - { - int? percent = value is >= 0 && maximum is > 0 ? (value.Value * 100) / maximum.Value : null; - - if ( - percent == lastPercent - && string.Equals(lastRunningNode, runningNode, StringComparison.Ordinal) - ) - { - return; - } - - lastPercent = percent; - lastRunningNode = runningNode; - - request.Progress?.Report( - new ImageGenerationProgress(ProviderId, task.Id, value, maximum, runningNode, stage) - ); - } - - void OnProgressUpdate( - object? sender, - StabilityMatrix.Core.Inference.ComfyProgressUpdateEventArgs args - ) - { - ReportProgress(args.Value, args.Maximum, args.RunningNode, "Generating"); - } - - void OnRunningNodeChanged(object? sender, string? node) - { - ReportProgress( - task.LastProgressUpdate?.Value, - task.LastProgressUpdate?.Maximum, - node, - "Generating" - ); - } + protected override string LogName => "Qwen Image Edit"; + protected override int MaxInputImages => 3; + protected override string ProviderPrefix => "qwen_image_edit"; - task.ProgressUpdate += OnProgressUpdate; - task.RunningNodeChanged += OnRunningNodeChanged; - - // Register cancellation to interrupt ComfyUI if the user cancels - await using var promptInterrupt = cancellationToken.Register(() => - { - logger.LogInformation( - "Cancellation requested, interrupting ComfyUI prompt {PromptId}", - task.Id - ); - // CTS holds an internal timer that needs disposing; chain dispose onto - // the fire-and-forget interrupt so it cleans up once the request settles. - var interruptCts = new CancellationTokenSource(5000); - clientManager - .Client.InterruptPromptAsync(interruptCts.Token) - .ContinueWith(_ => interruptCts.Dispose(), TaskScheduler.Default) - .SafeFireAndForget(); - }); - - try - { - // Wait for completion - logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); - await task.Task.WaitAsync(cancellationToken); - } - finally - { - task.ProgressUpdate -= OnProgressUpdate; - task.RunningNodeChanged -= OnRunningNodeChanged; - } - - // Get the output images - var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( - task.Id, - cancellationToken - ); - - // Prefer the "SaveImage" output node deterministically. - var preferredOutputKey = "SaveImage"; - string? selectedOutputKey = null; - List? candidateImages = null; - - if ( - outputImages.TryGetValue(preferredOutputKey, out var preferredImages) - && preferredImages is { Count: > 0 } - ) - { - selectedOutputKey = preferredOutputKey; - candidateImages = preferredImages; - } - else - { - var selected = outputImages - .OrderBy(kvp => kvp.Key, StringComparer.Ordinal) - .FirstOrDefault(kvp => kvp.Value is { Count: > 0 }); - - selectedOutputKey = string.IsNullOrEmpty(selected.Key) ? null : selected.Key; - candidateImages = selected.Value; - } - - if (candidateImages is null || candidateImages.Count == 0) - { - logger.LogWarning("No output images found from generation"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = "No output images were generated", - }; - } - - var generatedImages = new List(); - - foreach (var comfyImage in candidateImages) - { - logger.LogInformation("Downloading generated image: {FileName}", comfyImage.FileName); - - await using var imageStream = await clientManager.Client.GetImageStreamAsync( - comfyImage, - cancellationToken - ); - using var memoryStream = new MemoryStream(); - await imageStream.CopyToAsync(memoryStream, cancellationToken); - var imageBytes = memoryStream.ToArray(); - var base64Image = Convert.ToBase64String(imageBytes); - - var mimeType = - comfyImage.FileName.EndsWith(".png", StringComparison.OrdinalIgnoreCase) ? "image/png" - : comfyImage.FileName.EndsWith(".jpg", StringComparison.OrdinalIgnoreCase) - || comfyImage.FileName.EndsWith(".jpeg", StringComparison.OrdinalIgnoreCase) - ? "image/jpeg" - : comfyImage.FileName.EndsWith(".webp", StringComparison.OrdinalIgnoreCase) ? "image/webp" - : "image/png"; - - generatedImages.Add(new GeneratedImage { Base64Data = base64Image, MimeType = mimeType }); - } - - logger.LogInformation( - "Successfully generated {Count} image(s) with Qwen Image Edit", - generatedImages.Count - ); - - return new ImageGenerationResponse - { - IsSuccess = true, - Images = generatedImages, - TextResponse = null, - Metadata = new Dictionary - { - ["promptId"] = task.Id, - ["outputNode"] = selectedOutputKey ?? "unknown", - }, - }; - } - catch (OperationCanceledException) - { - logger.LogInformation("Image generation was cancelled"); - throw; // Propagate cancellation to ViewModel for proper handling - } - catch (Exception ex) + protected override IReadOnlyList GetMissingModels(ImageGenerationRequest request) + { + var modelManager = new QwenImageEditModelManager(); + if (modelManager.AreModelsAvailable(ClientManager)) { - logger.LogError(ex, "Failed to generate image with Qwen Image Edit"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = $"Generation failed: {ex.Message}", - }; + return []; } + return modelManager.GetMissingModelNames(ClientManager).ToList(); + } + + protected override Dictionary BuildWorkflow(ImageGenerationRequest request) + { + var customUnetModel = GetCustomUnetModel(request); + var loras = GetSelectedLoras(request); + var (width, height) = GetDimensions(request); + + return QwenImageEditWorkflowBuilder.Build( + request, + ClientManager, + customUnetModel, + loras, + width, + height + ); } } diff --git a/StabilityMatrix.Avalonia/Services/QwenImageEditWorkflowBuilder.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditWorkflowBuilder.cs index a423983f7..d29599a19 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditWorkflowBuilder.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditWorkflowBuilder.cs @@ -1,3 +1,4 @@ +using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.Api.Comfy.Nodes; @@ -101,29 +102,14 @@ public static Dictionary Build( var currentClip = clipLoader.Output; // Apply LoRAs if any - var loraList = loras?.ToList() ?? []; - if (loraList.Count > 0) - { - for (var i = 0; i < loraList.Count; i++) - { - var lora = loraList[i]; - var loraLoader = nodes.AddNamedNode( - ComfyNodeBuilder.LoraLoader( - nodes.GetUniqueName($"LoraLoader_{i + 1}"), - currentModel, - currentClip, - lora.Model.RelativePath, - (double)lora.ModelWeight, - (double)lora.ClipWeight - ) - ); - currentModel = loraLoader.Output1; - currentClip = loraLoader.Output2; - } - } + (currentModel, currentClip) = ComfyWorkflowHelper.ApplyLoras(nodes, loras, currentModel, currentClip); // Get reference image filenames (up to 3 for Qwen) - var referenceImageNames = GetReferenceImageNames(request); + var referenceImageNames = ComfyWorkflowHelper.GetReferenceImageNames( + request, + maxImages: 3, + providerPrefix: "qwen_image_edit" + ); // Load reference images if any ImageNodeConnection? image1 = null; @@ -248,43 +234,4 @@ public static Dictionary Build( // Return the node dictionary directly return nodes; } - - /// - /// Get reference image filenames from the request (input images + conversation history) - /// Note: Images uploaded via ComfyClient.UploadImageAsync go to the "Inference" subfolder - /// - private static List GetReferenceImageNames(ImageGenerationRequest request) - { - var imageNames = new List(); - - // Priority 1: Current input images (will be uploaded with known names) - // Qwen supports up to 3 images natively - if (request.InputImages?.Count > 0) - { - for (var i = 0; i < Math.Min(request.InputImages.Count, 3); i++) // Max 3 images for Qwen - { - // Include the "Inference/" subfolder prefix since that's where UploadImageAsync uploads to - imageNames.Add($"Inference/qwen_image_edit_input_{i}.png"); - } - } - - // Priority 2: Most recent image from conversation history (previous generation) - // Only include if we have room (less than 3 images) - if (imageNames.Count < 3 && request.ConversationHistory != null) - { - var lastAssistantImage = request.ConversationHistory.LastOrDefault(m => - m is { Role: MessageRole.Assistant, ImageContent: not null } - ); - - // Only add the filename if we found an image with content - // (the provider will have uploaded it) - if (lastAssistantImage?.ImageContent != null) - { - // Include the "Inference/" subfolder prefix - imageNames.Add("Inference/qwen_image_edit_history_latest.png"); - } - } - - return imageNames; - } } diff --git a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Downloads.cs b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Downloads.cs index 39f6c96ce..1ca768cd8 100644 --- a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Downloads.cs +++ b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Downloads.cs @@ -93,7 +93,7 @@ private async Task ShowMissingModelsDialogAsync() return; } - var missingModels = modelManager.GetMissingModels(ClientManager).ToList(); + var missingModels = GetProviderMissingModels(modelManager).ToList(); if (missingModels.Count == 0) { diff --git a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Models.cs b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Models.cs index 67bf7aafe..12208b5a3 100644 --- a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Models.cs +++ b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.Models.cs @@ -9,13 +9,47 @@ using Microsoft.Extensions.Logging; using StabilityMatrix.Avalonia.Controls; using StabilityMatrix.Avalonia.Models.BananaVision; +using StabilityMatrix.Avalonia.Services; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.ViewModels; +internal enum BananaVisionModelTermMatch +{ + Excluded, + Primary, + Secondary, + Untagged, +} + public partial class BananaVisionPageViewModel { + /// + /// Provider-aware availability check. Klein's check is encoder-variant aware (a 9B UNET + /// needs the 8B encoder), so the user's UNET selection is passed through for it. + /// + private bool AreProviderModelsAvailable(ILocalProviderModelManager manager) => + manager is Flux2KleinModelManager klein + ? klein.AreModelsAvailable(ClientManager, SelectedKleinModel) + : manager.AreModelsAvailable(ClientManager); + + /// + /// Provider-aware missing-model list. See . + /// + private IEnumerable GetProviderMissingModels(ILocalProviderModelManager manager) => + manager is Flux2KleinModelManager klein + ? klein.GetMissingModels(ClientManager, SelectedKleinModel) + : manager.GetMissingModels(ClientManager); + + /// + /// Provider-aware missing-model display names. See . + /// + private IEnumerable GetProviderMissingModelNames(ILocalProviderModelManager manager) => + manager is Flux2KleinModelManager klein + ? klein.GetMissingModelNames(ClientManager, SelectedKleinModel) + : manager.GetMissingModelNames(ClientManager); + /// /// Sorts models by connected status first, then alphabetically by display name /// @@ -69,47 +103,64 @@ List Untagged foreach (var model in modelIndexService.FindByModelType(folderType).Select(HybridModelFile.FromLocal)) { - var baseModel = model.Local?.ConnectedModelInfo?.BaseModel; - - // Check primary terms first - if ( - primaryTerms.Any(term => - baseModel?.Contains(term, StringComparison.OrdinalIgnoreCase) == true - ) - ) + switch (GetModelTermMatch(model, primaryTerms, secondaryTerms)) { - primaryModels.Add(model); - } - // Check secondary terms - else if ( - secondaryTerms?.Any(term => - baseModel?.Contains(term, StringComparison.OrdinalIgnoreCase) == true - ) == true - ) - { - secondaryModels.Add(model); - } - // Check filename fallback for untagged models - else if (string.IsNullOrEmpty(baseModel)) - { - if ( - primaryTerms.Any(term => - model.FileName.Contains(term, StringComparison.OrdinalIgnoreCase) - ) - ) - { + case BananaVisionModelTermMatch.Primary: primaryModels.Add(model); - } - else - { + break; + case BananaVisionModelTermMatch.Secondary: + secondaryModels.Add(model); + break; + case BananaVisionModelTermMatch.Untagged: untaggedModels.Add(model); - } + break; } } return (primaryModels, secondaryModels, untaggedModels); } + internal static BananaVisionModelTermMatch GetModelTermMatch( + HybridModelFile model, + string[] primaryTerms, + string[]? secondaryTerms = null + ) + { + var baseModel = model.Local?.ConnectedModelInfo?.BaseModel; + + if (primaryTerms.Any(term => baseModel?.Contains(term, StringComparison.OrdinalIgnoreCase) == true)) + { + return BananaVisionModelTermMatch.Primary; + } + + if ( + secondaryTerms?.Any(term => baseModel?.Contains(term, StringComparison.OrdinalIgnoreCase) == true) + == true + ) + { + return BananaVisionModelTermMatch.Secondary; + } + + // Filename fallback applies even when metadata is present but unrecognized, which + // is common for CivitAI uploads tagged "Other". + if (primaryTerms.Any(term => model.FileName.Contains(term, StringComparison.OrdinalIgnoreCase))) + { + return BananaVisionModelTermMatch.Primary; + } + + if ( + secondaryTerms?.Any(term => model.FileName.Contains(term, StringComparison.OrdinalIgnoreCase)) + == true + ) + { + return BananaVisionModelTermMatch.Secondary; + } + + return string.IsNullOrEmpty(baseModel) + ? BananaVisionModelTermMatch.Untagged + : BananaVisionModelTermMatch.Excluded; + } + /// /// Loads available Flux Kontext models from the DiffusionModels folder using local model index /// @@ -340,6 +391,10 @@ partial void OnSelectedKleinModelChanged(HybridModelFile? value) var (recommendedSteps, recommendedCfg) = DetectKleinDefaults(value); KleinSteps = recommendedSteps; KleinCfg = recommendedCfg; + + // The availability check is encoder-variant aware, so switching between a 4B and a + // 9B UNET can change which text encoder counts as missing — re-evaluate the banner. + UpdateProviderStatus(); } /// diff --git a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs index b78c4d6c1..a3f610c19 100644 --- a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs @@ -19,6 +19,7 @@ using Microsoft.Extensions.Logging; using StabilityMatrix.Avalonia.Controls; using StabilityMatrix.Avalonia.Helpers; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Models; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Avalonia.Services; @@ -260,9 +261,6 @@ await Dispatcher.UIThread.InvokeAsync(() => [ObservableProperty] public partial string? ProviderStatusMessage { get; set; } - [ObservableProperty] - public partial bool IsFluxKontextAvailable { get; set; } - [ObservableProperty] public partial bool CanRetryLastMessage { get; set; } @@ -1109,15 +1107,14 @@ private async Task SendMessageAsync(CancellationToken cancellationToken) logger.LogWarning("Image generation failed: {Message}", ex.Message); // Check if this is an API key error - if (ex.Message.Contains("API key", StringComparison.OrdinalIgnoreCase)) + if (IsGeminiSetupError(ex)) { - await ShowApiKeyRequiredDialogAsync(); + await ShowApiKeyRequiredDialogAsync(ex); CanRetryLastMessage = true; } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -1143,18 +1140,74 @@ private async Task SendMessageAsync(CancellationToken cancellationToken) } } + /// + /// Surfaces a failed generation: short message in the chat error banner, plus the same + /// JSON detail dialog Inference uses when the failure carries ComfyUI's error body + /// (workflow rejection or node execution error), or a plain toast otherwise. + /// + private async Task ShowGenerationFailedAsync(ImageGenerationException ex) + { + var errorMessage = GetLocalizedGenerationError(ex); + ErrorMessage = errorMessage; + + if (!string.IsNullOrWhiteSpace(ex.DetailJson)) + { + await DialogHelper + .CreateJsonDialog( + ex.DetailJson, + Resources.Label_ComfyError, + Resources.Text_ComfyReportedGenerationError + ) + .ShowAsync(); + } + else + { + notificationService.Show( + Resources.Label_GenerationFailed, + errorMessage, + NotificationType.Warning + ); + } + } + /// /// Shows a dialog prompting the user to add their Gemini API key in settings /// - private async Task ShowApiKeyRequiredDialogAsync() + private static bool IsGeminiSetupError(ImageGenerationException exception) => + exception.ErrorCode + is ImageGenerationErrorCode.GeminiApiKeyNotConfigured + or ImageGenerationErrorCode.GeminiQuotaExceeded + or ImageGenerationErrorCode.GeminiInvalidApiKey + or ImageGenerationErrorCode.GeminiAccessForbidden + || exception.Message.Contains("API key", StringComparison.OrdinalIgnoreCase) + || exception.Message.Contains("billing", StringComparison.OrdinalIgnoreCase) + || exception.Message.Contains("quota", StringComparison.OrdinalIgnoreCase); + + private static string GetLocalizedGenerationError(ImageGenerationException exception) => + exception.ErrorCode switch + { + ImageGenerationErrorCode.GeminiApiKeyNotConfigured => Resources.Error_GeminiApiKeyNotConfigured, + ImageGenerationErrorCode.GeminiQuotaExceeded => Resources.Error_GeminiQuotaExceeded, + ImageGenerationErrorCode.GeminiInvalidApiKey => Resources.Error_GeminiInvalidApiKey, + ImageGenerationErrorCode.GeminiAccessForbidden => Resources.Error_GeminiAccessForbidden, + _ => exception.Message, + }; + + private async Task ShowApiKeyRequiredDialogAsync(ImageGenerationException exception) { + var errorMessage = GetLocalizedGenerationError(exception); + var content = exception.ErrorCode + is ImageGenerationErrorCode.GeminiQuotaExceeded + or ImageGenerationErrorCode.GeminiAccessForbidden + ? errorMessage + : errorMessage + "\n\n" + Resources.Text_GeminiPaidTierRequired; + var dialog = new ContentDialog { - Title = "API Key Required", - Content = - "Gemini API key not configured. Please add your Gemini API key in Account Settings to use cloud providers.", - PrimaryButtonText = "Open Settings", - CloseButtonText = "Cancel", + Title = Resources.Label_GeminiApiSetupRequired, + Content = content, + PrimaryButtonText = Resources.Action_GoToSettings, + CloseButtonText = Resources.Action_Cancel, DefaultButton = ContentDialogButton.Primary, }; @@ -1171,6 +1224,24 @@ private async Task ShowApiKeyRequiredDialogAsync() } } + /// + /// Debug-only: routes a synthetic through the same + /// handlers a real failed generation uses, so the error banner, setup dialog, and + /// notifications can be previewed without triggering an actual API failure. A null + /// exercises the generic (non-Gemini) failure path. + /// + public Task DebugSimulateGenerationErrorAsync(ImageGenerationErrorCode? errorCode) + { + var exception = new ImageGenerationException("Debug simulated image generation error") + { + ErrorCode = errorCode, + }; + + return IsGeminiSetupError(exception) + ? ShowApiKeyRequiredDialogAsync(exception) + : ShowGenerationFailedAsync(exception); + } + [RelayCommand(IncludeCancelCommand = true)] private async Task RetryLastMessageAsync(CancellationToken cancellationToken) { @@ -1260,15 +1331,14 @@ private async Task RetryLastMessageAsync(CancellationToken cancellationToken) logger.LogWarning("Retry generation failed: {Message}", ex.Message); // Check if this is an API key error - if (ex.Message.Contains("API key", StringComparison.OrdinalIgnoreCase)) + if (IsGeminiSetupError(ex)) { - await ShowApiKeyRequiredDialogAsync(); + await ShowApiKeyRequiredDialogAsync(ex); CanRetryLastMessage = true; } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -1452,15 +1522,14 @@ private async Task RegenerateLastResponseAsync(CancellationToken cancellationTok logger.LogWarning("Regenerate failed: {Message}", ex.Message); // Check if this is an API key error - if (ex.Message.Contains("API key", StringComparison.OrdinalIgnoreCase)) + if (IsGeminiSetupError(ex)) { - await ShowApiKeyRequiredDialogAsync(); + await ShowApiKeyRequiredDialogAsync(ex); CanRetryLastMessage = true; } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -1694,15 +1763,14 @@ string editedText { logger.LogWarning("Failed to regenerate after edit: {Message}", ex.Message); - if (ex.Message.Contains("API key", StringComparison.OrdinalIgnoreCase)) + if (IsGeminiSetupError(ex)) { - await ShowApiKeyRequiredDialogAsync(); + await ShowApiKeyRequiredDialogAsync(ex); CanRetryLastMessage = true; } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -2277,7 +2345,6 @@ private void UpdateProviderStatus() if (!IsComfyRunning) { ProviderStatusMessage = "⚠️ ComfyUI is not running. Click Launch to start."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2286,7 +2353,6 @@ private void UpdateProviderStatus() if (IsWaitingForConnection) { ProviderStatusMessage = "🔄 Connecting to ComfyUI..."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2295,7 +2361,6 @@ private void UpdateProviderStatus() if (!ClientManager.IsConnected) { ProviderStatusMessage = "⚠️ Not connected to ComfyUI. Click Connect."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2307,32 +2372,27 @@ private void UpdateProviderStatus() if (IsDownloadingModels) { ProviderStatusMessage = DownloadProgressText ?? "⬇️ Downloading models..."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } // Check if required models are available - if (!modelManager.AreModelsAvailable(ClientManager)) + if (!AreProviderModelsAvailable(modelManager)) { - var missingModelNames = modelManager.GetMissingModelNames(ClientManager).ToList(); - var modelsList = string.Join(", ", missingModelNames); + var modelsList = string.Join(", ", GetProviderMissingModelNames(modelManager)); ProviderStatusMessage = $"⚠️ Missing: {modelsList}"; - IsFluxKontextAvailable = false; HasMissingModels = true; return; } // All good ProviderStatusMessage = $"✅ {modelManager.ProviderDisplayName} is ready"; - IsFluxKontextAvailable = true; HasMissingModels = false; } else { // Cloud providers or providers without model requirements ProviderStatusMessage = null; - IsFluxKontextAvailable = false; HasMissingModels = false; } } diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivArchiveDetailsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivArchiveDetailsPageViewModel.cs index d84c50186..7e6feedaf 100644 --- a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivArchiveDetailsPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivArchiveDetailsPageViewModel.cs @@ -5,6 +5,7 @@ using System.Reactive.Linq; using System.Threading; using AsyncAwaitBestPractices; +using Avalonia.Platform.Storage; using Avalonia.Threading; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; @@ -13,6 +14,7 @@ using FluentAvalonia.UI.Controls; using Injectio.Attributes; using StabilityMatrix.Avalonia.Extensions; +using StabilityMatrix.Avalonia.Languages; using StabilityMatrix.Avalonia.Models; using StabilityMatrix.Avalonia.Models.Inference; using StabilityMatrix.Avalonia.Services; @@ -108,6 +110,7 @@ IModelIndexService modelIndexService public ObservableCollection Images { get; } = []; public ObservableCollection Files { get; } = []; public ObservableCollection Mirrors { get; } = []; + public ObservableCollection InstallLocations { get; } = []; private static readonly Dictionary< string, @@ -327,6 +330,7 @@ private void PopulateVersionData(CivArchiveModelVersion? version) Mirrors.Add(mirror); } + RefreshInstallLocations(); UpdateInstalledStatus(version); } @@ -524,6 +528,48 @@ private async Task DownloadModel() await ExecuteDownloadAsync(version, primaryFile, GetDownloadUris(version), sourceLabel: null); } + [RelayCommand] + private async Task DownloadModelToLocation(InstallLocationOption? location) + { + var version = Model?.Version; + if (version is null || location is null) + return; + + DirectoryPath destinationDir; + if (location.Directory is { } existingDirectory) + { + destinationDir = existingDirectory; + } + else + { + // No directory means the "Custom..." entry — prompt for a folder. + var folders = await App.StorageProvider.OpenFolderPickerAsync( + new FolderPickerOpenOptions + { + Title = "Select Download Folder", + AllowMultiple = false, + SuggestedStartLocation = await App.StorageProvider.TryGetFolderFromPathAsync( + GetDefaultDownloadFolder() + ), + } + ); + + if (folders.FirstOrDefault()?.TryGetLocalPath() is not { } customPath) + return; + + destinationDir = new DirectoryPath(customPath); + } + + var primaryFile = GetPrimaryFile(version); + await ExecuteDownloadAsync( + version, + primaryFile, + GetDownloadUris(version), + sourceLabel: null, + destinationDir + ); + } + [RelayCommand] private async Task DeleteModel() { @@ -663,7 +709,8 @@ private async Task ExecuteDownloadAsync( CivArchiveModelVersion version, CivArchiveModelFile? file, IReadOnlyList downloadUris, - string? sourceLabel + string? sourceLabel, + DirectoryPath? destinationOverride = null ) { if (Model is null) @@ -684,7 +731,7 @@ private async Task ExecuteDownloadAsync( return; } - var destinationDir = GetDefaultDownloadFolder(); + var destinationDir = destinationOverride ?? GetDefaultDownloadFolder(); var fileName = BuildDownloadFileName(version, file); Uri? previewImageUri = null; @@ -852,6 +899,71 @@ private DirectoryPath GetDefaultDownloadFolder() return new DirectoryPath(settingsManager.ModelsDirectory); } + private void RefreshInstallLocations() + { + var modelsDirectory = new DirectoryPath(settingsManager.ModelsDirectory); + var defaultDirectory = GetDefaultDownloadFolder(); + LoadInstallLocationsAsync(modelsDirectory, defaultDirectory).SafeFireAndForget(); + } + + private async Task LoadInstallLocationsAsync( + DirectoryPath modelsDirectory, + DirectoryPath defaultDirectory + ) + { + // Enumerate off the UI thread — a deeply-nested models folder can hold many subdirectories. + var options = await Task.Run(() => + { + var results = new List(); + + if (defaultDirectory.Exists) + { + results.Add( + new InstallLocationOption( + Path.Combine("Models", Path.GetRelativePath(modelsDirectory, defaultDirectory)), + defaultDirectory + ) + ); + + try + { + foreach ( + var directory in defaultDirectory.EnumerateDirectories( + "*", + EnumerationOptionConstants.AllDirectories + ) + ) + { + results.Add( + new InstallLocationOption( + Path.Combine("Models", Path.GetRelativePath(modelsDirectory, directory)), + directory + ) + ); + } + } + catch (Exception) + { + // Ignore enumeration failures (e.g. an unauthorized subdirectory) so the + // default and custom options remain available. + } + } + + return results; + }); + + // Repopulate atomically on the UI thread; the last completed refresh wins. + await Dispatcher.UIThread.InvokeAsync(() => + { + InstallLocations.Clear(); + foreach (var option in options) + { + InstallLocations.Add(option); + } + InstallLocations.Add(new InstallLocationOption(Resources.Action_CustomFolderEllipsis, null)); + }); + } + private static ConnectedModelInfo BuildConnectedModelInfo( CivArchiveModelDetails model, CivArchiveModelVersion version, diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/InstallLocationOption.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/InstallLocationOption.cs new file mode 100644 index 000000000..39d271fa3 --- /dev/null +++ b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/InstallLocationOption.cs @@ -0,0 +1,10 @@ +using StabilityMatrix.Core.Models.FileInterfaces; + +namespace StabilityMatrix.Avalonia.ViewModels.CheckpointBrowser; + +/// +/// A download target shown in the CivArchive download split-button flyout. +/// is null for the "Custom..." entry, which prompts the user +/// to pick a folder instead of using a known models subdirectory. +/// +public record InstallLocationOption(string DisplayName, DirectoryPath? Directory); diff --git a/StabilityMatrix.Avalonia/ViewModels/Controls/PaintCanvasViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Controls/PaintCanvasViewModel.cs index 99d1444e8..67364ab4a 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Controls/PaintCanvasViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Controls/PaintCanvasViewModel.cs @@ -1161,10 +1161,24 @@ public void RenderToSurface( var needsNewSurface = layer.Surface is null; if (!needsNewSurface) { - // Check if we need to resize - var currentInfo = layer.Surface!.Canvas.DeviceClipBounds; - needsNewSurface = - currentInfo.Width != CanvasSize.Width || currentInfo.Height != CanvasSize.Height; + // Recreate if the existing surface's backing doesn't match the current target. + // On-screen rendering leases a GPU surface, so the persistent layer surfaces are + // GPU-backed and tied to the render thread. Off-screen export (e.g. saving an + // annotation) composites onto a CPU surface from another thread; reusing those + // GPU surfaces there produces a blank image. Forcing a matching CPU surface fixes + // it, and the next on-screen render simply recreates the GPU surface. + var layerIsGpu = layer.Surface!.Context != null; + if (layerIsGpu != useGpu) + { + needsNewSurface = true; + } + else + { + // Check if we need to resize + var currentInfo = layer.Surface!.Canvas.DeviceClipBounds; + needsNewSurface = + currentInfo.Width != CanvasSize.Width || currentInfo.Height != CanvasSize.Height; + } } if (needsNewSurface) diff --git a/StabilityMatrix.Avalonia/ViewModels/Dialogs/MaskEditorViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Dialogs/MaskEditorViewModel.cs index e9fa86cb5..7cd799555 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Dialogs/MaskEditorViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Dialogs/MaskEditorViewModel.cs @@ -244,6 +244,7 @@ public record MaskEditorModel }*/ public void Dispose() { + _cachedMaskRenderImage?.Dispose(); _cachedMaskRenderInverseAlphaImage?.Dispose(); GC.SuppressFinalize(this); } diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs index d081f3717..4f43c77d9 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/InferenceTextToImageViewModel.cs @@ -249,16 +249,6 @@ protected override void BuildPrompt(BuildPromptEventArgs args) protected InferenceWorkflowProfile ResolvedWorkflowProfile => ModelCardViewModel.ResolvedWorkflowProfile; - protected bool IsAnimaUnet => - IsUnetLoader - && ( - ResolvedWorkflowProfile is InferenceWorkflowProfile.Anima - || ( - ResolvedWorkflowProfile is InferenceWorkflowProfile.Custom - && ModelCardViewModel.SelectedClipType is "stable_diffusion" - ) - ); - protected bool IsFlux2Unet => IsUnetLoader && ( diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs index 845ad33b6..0e5429058 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs @@ -1,6 +1,9 @@ using System.Collections.ObjectModel; using System.ComponentModel.DataAnnotations; using System.Text.Json.Nodes; +using System.Text.RegularExpressions; +using Avalonia.Controls.Notifications; +using Avalonia.Threading; using CommunityToolkit.Mvvm.ComponentModel; using CommunityToolkit.Mvvm.Input; using FluentAvalonia.UI.Controls; @@ -14,11 +17,14 @@ using StabilityMatrix.Avalonia.ViewModels.Dialogs; using StabilityMatrix.Avalonia.ViewModels.Inference.Modules; using StabilityMatrix.Core.Attributes; +using StabilityMatrix.Core.Exceptions; using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Api; using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes; using StabilityMatrix.Core.Models.Inference; +using StabilityMatrix.Core.Services; namespace StabilityMatrix.Avalonia.ViewModels.Inference; @@ -28,7 +34,11 @@ namespace StabilityMatrix.Avalonia.ViewModels.Inference; public partial class ModelCardViewModel( IInferenceClientManager clientManager, IServiceManager vmFactory, - TabContext tabContext + TabContext tabContext, + ISettingsManager settingsManager, + IModelIndexService modelIndexService, + INotificationService notificationService, + ModelOrganizationService modelOrganizationService ) : LoadableViewModelBase, IParametersLoadableState, IComfyStep { [ObservableProperty] @@ -36,6 +46,9 @@ TabContext tabContext [NotifyPropertyChangedFor(nameof(WorkflowProfileStatusText))] [NotifyPropertyChangedFor(nameof(ShowWorkflowProfileStatus))] [NotifyPropertyChangedFor(nameof(RecommendedDefaultsToolTip))] + [NotifyPropertyChangedFor(nameof(WorkflowProfileWarningText))] + [NotifyPropertyChangedFor(nameof(ShowWorkflowProfileWarning))] + [NotifyPropertyChangedFor(nameof(MoveModelToRecommendedFolderText))] private HybridModelFile? selectedModel; [ObservableProperty] @@ -219,7 +232,23 @@ public HybridModelFile? SelectedUnifiedModel public List WeightDTypes { get; set; } = ["default", "fp8_e4m3fn", "fp8_e5m2"]; public List ClipTypes { get; set; } = - ["flux", "flux2", "lumina2", "stable_diffusion", "sd3", "HiDream"]; + [ + "sdxl", + "sd3", + "flux", + "hunyuan_video", + "hidream", + "hunyuan_image", + "hunyuan_video_15", + "kandinsky5", + "kandinsky5_image", + "ltxv", + "newbie", + "ace", + "flux2", + "lumina2", + "stable_diffusion", + ]; public List WorkflowProfiles { get; set; } = Enum.GetValues().ToList(); @@ -240,10 +269,10 @@ public HybridModelFile? SelectedUnifiedModel public bool ShowEncoderSection => IsClipModelSelectionEnabled && IsStandaloneModelLoader; public bool IsSd3Clip => SelectedClipType == "sd3"; - public bool IsHiDreamClip => SelectedClipType == "HiDream"; + public bool IsHiDreamClip => SelectedClipType == "hidream"; public bool IsHiDreamWorkflow => ResolvedWorkflowProfile is InferenceWorkflowProfile.HiDream - || (ResolvedWorkflowProfile is InferenceWorkflowProfile.Custom && SelectedClipType == "HiDream"); + || (ResolvedWorkflowProfile is InferenceWorkflowProfile.Custom && SelectedClipType == "hidream"); public bool IsZImageWorkflow => ResolvedWorkflowProfile is InferenceWorkflowProfile.ZImageBase or InferenceWorkflowProfile.ZImageTurbo || (ResolvedWorkflowProfile is InferenceWorkflowProfile.Custom && SelectedClipType == "lumina2"); @@ -318,6 +347,216 @@ SelectedWorkflowProfile is InferenceWorkflowProfile.Auto _ => "No recommended sampler defaults for this workflow", }; + /// + /// Describes a mismatch between the selected model's physical folder and the selected + /// workflow. ComfyUI resolves each loader's paths relative to its own folder, so a + /// checkpoint-folder file can never load through UNETLoader (and vice versa) — without + /// this hint the user only finds out via a cryptic node-validation error at generation. + /// TargetFolder is where the file should live if the profile is right. + /// + private (string Message, SharedFolderType TargetFolder)? GetProfileFolderMismatch() + { + if (SelectedUnifiedModel?.Local is not { } local) + return null; + + var isStandaloneFile = local.SharedFolderType is SharedFolderType.DiffusionModels; + + // Explicit profile that disagrees with the file's folder. Auto follows the + // folder and Custom is the user's business. + if (SelectedWorkflowProfile is not (InferenceWorkflowProfile.Auto or InferenceWorkflowProfile.Custom)) + { + // Flux.1 is packaging-flexible: users commonly have either a split UNet in + // DiffusionModels or an all-in-one checkpoint in StableDiffusion. + var profileWantsStandalone = SelectedWorkflowProfile switch + { + InferenceWorkflowProfile.DefaultCheckpoint => false, + InferenceWorkflowProfile.Flux => (bool?)null, + _ => true, + }; + + if (profileWantsStandalone is true && !isStandaloneFile) + { + return ( + string.Format( + Resources.TextTemplate_WorkflowNeedsDiffusionModelsFile, + SelectedWorkflowProfile.GetStringValue(), + local.SharedFolderType.GetStringValue() + ), + SharedFolderType.DiffusionModels + ); + } + + if (profileWantsStandalone is false && isStandaloneFile) + { + return (Resources.Text_FileInDiffusionModelsNotCheckpoint, SharedFolderType.StableDiffusion); + } + } + else if (SelectedWorkflowProfile is InferenceWorkflowProfile.Auto && !isStandaloneFile) + { + // Respect explicit metadata: an SDXL-tagged checkpoint that merely contains + // "anima" or "flux2" in its filename is not a misplaced standalone model. + var taggedBase = local.ConnectedModelInfo?.BaseModel; + if (!string.IsNullOrWhiteSpace(taggedBase) && IsKnownCheckpointBaseModelTag(taggedBase)) + return null; + + // Checkpoint-folder file whose name/metadata implies a UNet-only architecture + // (the classic "z-image in the StableDiffusion folder" support thread). Anima is + // UNet-only too — it always pairs with a separate qwen_3_06b encoder + Qwen Image + // VAE; the Civitai "Checkpoint" category on Anima models does not mean all-in-one. + // Flux is excluded since Flux.1 all-in-one checkpoints are common and valid. + var impliedProfile = InferWorkflowProfile(SelectedUnifiedModel, isUnetModel: true); + if ( + impliedProfile + is InferenceWorkflowProfile.Flux2 + or InferenceWorkflowProfile.ZImageBase + or InferenceWorkflowProfile.ZImageTurbo + or InferenceWorkflowProfile.HiDream + or InferenceWorkflowProfile.Anima + ) + { + return ( + string.Format( + Resources.TextTemplate_FileLooksLikeStandaloneModel, + impliedProfile.GetStringValue(), + local.SharedFolderType.GetStringValue() + ), + SharedFolderType.DiffusionModels + ); + } + } + + return null; + } + + public string? WorkflowProfileWarningText => GetProfileFolderMismatch()?.Message; + + /// + /// RelativePath of the model whose folder-mismatch warning the user dismissed. + /// Selecting a different mismatched model shows the warning again; re-selecting the + /// dismissed one stays quiet. + /// + private string? dismissedWarningModelPath; + + public bool ShowWorkflowProfileWarning => + WorkflowProfileWarningText is not null + && !string.Equals( + dismissedWarningModelPath, + SelectedUnifiedModel?.Local?.RelativePath, + StringComparison.OrdinalIgnoreCase + ); + + [RelayCommand] + private void DismissWorkflowProfileWarning() + { + if (SelectedUnifiedModel?.Local?.RelativePath is { } path) + { + dismissedWarningModelPath = path; + } + + OnPropertyChanged(nameof(ShowWorkflowProfileWarning)); + } + + /// + /// Label for the one-click fix button next to the mismatch warning, e.g. + /// "Move to DiffusionModels". + /// + public string? MoveModelToRecommendedFolderText => + GetProfileFolderMismatch() is { } mismatch + ? string.Format(Resources.TextTemplate_MoveToFolder, mismatch.TargetFolder.GetStringValue()) + : null; + + /// + /// Moves the selected model (and its metadata sidecars) into the folder the current + /// workflow expects, refreshes the model index, and re-selects the moved file. + /// + [RelayCommand] + private async Task MoveModelToRecommendedFolderAsync() + { + if (GetProfileFolderMismatch() is not { } mismatch || SelectedUnifiedModel?.Local is not { } local) + return; + + var fileName = local.FileName; + var destinationDir = Path.Combine( + settingsManager.ModelsDirectory, + mismatch.TargetFolder.GetStringValue() + ); + + try + { + await modelOrganizationService.MoveModelFileAsync( + local, + settingsManager.ModelsDirectory, + destinationDir + ); + } + catch (FileTransferExistsException) + { + notificationService.Show( + Resources.Label_CouldNotMoveModel, + string.Format( + Resources.TextTemplate_FileAlreadyExistsInFolder, + fileName, + mismatch.TargetFolder.GetStringValue() + ), + NotificationType.Error + ); + return; + } + catch (Exception e) + { + notificationService.Show( + Resources.Label_CouldNotMoveModel, + $"[{e.GetType().Name}] {e.Message}", + NotificationType.Error + ); + return; + } + + await modelIndexService.RefreshIndex(); + + // The index-changed handler posts the refreshed model collections to the UI thread, + // so re-select at Background priority to run after those updates have landed — + // selecting before the item exists in the ComboBox items would reset it to null. + Dispatcher.UIThread.Post( + () => + { + var movedModel = FindMovedModel( + mismatch.TargetFolder is SharedFolderType.DiffusionModels + ? ClientManager.UnetModels + : ClientManager.Models, + fileName + ); + + if (movedModel is not null) + { + SelectedUnifiedModel = movedModel; + } + + NotifyWorkflowProfileStateChanged(); + }, + DispatcherPriority.Background + ); + + notificationService.Show( + Resources.Label_ModelMoved, + string.Format( + Resources.TextTemplate_MovedFileToFolder, + fileName, + mismatch.TargetFolder.GetStringValue() + ), + NotificationType.Success + ); + } + + internal static HybridModelFile? FindMovedModel( + IEnumerable models, + string pathRelativeToSharedFolder + ) => + models.FirstOrDefault(m => + m.Local != null + && string.Equals(m.RelativePath, pathRelativeToSharedFolder, StringComparison.OrdinalIgnoreCase) + ); + public event Action? RecommendedDefaultsRequested; protected override void OnInitialLoaded() @@ -607,7 +846,7 @@ public override void LoadStateFromJsonObject(JsonObject state) : ClientManager.Models.FirstOrDefault(x => x.RelativePath == model.SelectedRefinerName); // Load encoder type first (needed for default encoder count) - SelectedClipType = model.SelectedClipType; + SelectedClipType = model.SelectedClipType == "HiDream" ? "hidream" : model.SelectedClipType; SelectedWorkflowProfile = model.SelectedWorkflowProfile; // Load text encoders from saved state @@ -699,7 +938,7 @@ private static InferenceWorkflowProfile InferWorkflowProfile(HybridModelFile? mo if (name.Contains("flux", StringComparison.OrdinalIgnoreCase)) return InferenceWorkflowProfile.Flux; - if (name.Contains("anima", StringComparison.OrdinalIgnoreCase)) + if (IsAnimaName(name)) return InferenceWorkflowProfile.Anima; if (name.Contains("hidream", StringComparison.OrdinalIgnoreCase)) @@ -708,6 +947,27 @@ private static InferenceWorkflowProfile InferWorkflowProfile(HybridModelFile? mo return InferenceWorkflowProfile.DefaultCheckpoint; } + /// + /// Matches "anima" only as its own token — "animagine" / "animaPencil" are SDXL models + /// that must not be detected as the Anima architecture. + /// + private static bool IsAnimaName(string name) => AnimaNameRegex().IsMatch(name); + + /// + /// Whether a CivitAI BaseModel tag is specific enough to override filename heuristics for + /// a checkpoint-folder file. "Other" remains inconclusive, while known UNet-only tags are + /// allowed to proceed to the misplaced-model check. + /// + private static bool IsKnownCheckpointBaseModelTag(string baseModel) => + Enum.GetValues() + .Any(type => + type is not (CivitBaseModelType.Other or CivitBaseModelType.HiDream) + && type.GetStringValue().Equals(baseModel, StringComparison.OrdinalIgnoreCase) + ); + + [GeneratedRegex(@"(? /// Loads text encoders from the saved model state, supporting both new and legacy formats. /// @@ -751,7 +1011,7 @@ private void LoadTextEncodersFromModel(ModelCardModel model) "flux" => 2, "flux2" or "lumina2" or "stable_diffusion" => 1, "sd3" => 3, - "HiDream" => 4, + "hidream" => 4, _ => 2, }; encoderCount = Math.Max(encoderCount, defaultCount); @@ -870,6 +1130,16 @@ partial void OnSelectedModelLoaderChanged(ModelLoader value) if (TextEncoders.Count == 0) SetDefaultEncoderCount(); } + else if (value is ModelLoader.Default or ModelLoader.Nf4 && !isLoadingState) + { + // The separate text-encoder UI is only reachable for standalone loaders, so a + // stale true here (left over from a previously selected UNet model) makes + // generation demand encoders the user can't even see. Also drop any encoder / + // VAE values WE auto-filled for the previous UNet workflow — a Flux.2 VAE + // override on an all-in-one checkpoint would just corrupt its output. + IsClipModelSelectionEnabled = false; + ClearAutoSelectedComponents(); + } RefreshWorkflowProfileState(); } @@ -914,12 +1184,103 @@ partial void OnSelectedWorkflowProfileChanged(InferenceWorkflowProfile value) { if (!isLoadingState) { + SyncModelLoaderToProfile(value); ApplyDefaultClipTypeForResolvedProfile(preserveUserSelections: true); } RefreshWorkflowProfileState(); } + /// + /// Keeps in sync with the user-facing Workflow dropdown. + /// The encoder / precision / text-encoder fields (and the generated workflow itself) key off + /// the model loader, so without this, picking "Default / Checkpoint" on a model that was + /// auto-detected as a UNet workflow leaves those separate-component fields visible and still + /// required at generation time. All-in-one checkpoints bundle their own VAE + text encoder, + /// so switching to one also drops the VAE / text-encoder selectors the UNet profile turned on. + /// + private void SyncModelLoaderToProfile(InferenceWorkflowProfile profile) + { + // Whether the profile uses separate diffusion-model + encoder + VAE files (true) or an + // all-in-one checkpoint (false). Null = leave the loader alone (Custom, or Auto with + // no model yet to infer from). + var wantStandalone = profile switch + { + InferenceWorkflowProfile.DefaultCheckpoint => false, + InferenceWorkflowProfile.Flux + or InferenceWorkflowProfile.Flux2 + or InferenceWorkflowProfile.Anima + or InferenceWorkflowProfile.ZImageBase + or InferenceWorkflowProfile.ZImageTurbo + or InferenceWorkflowProfile.HiDream => true, + // Auto follows the selected model's folder (diffusion_models -> UNet); Custom is + // left entirely to the user. + InferenceWorkflowProfile.Auto when SelectedUnifiedModel is not null => SelectedUnifiedModel + .Local + ?.SharedFolderType is SharedFolderType.DiffusionModels, + _ => (bool?)null, + }; + + if (wantStandalone is not { } standalone) + return; + + var currentlyStandalone = SelectedModelLoader is ModelLoader.Unet or ModelLoader.Gguf; + var currentModel = SelectedUnifiedModel; + + // A local file can only be loaded by the loader matching its physical folder — ComfyUI + // resolves CheckpointLoader paths against the checkpoints folder and UNETLoader paths + // against diffusion_models, so carrying a model across to the other loader is a + // guaranteed node-validation error. When the profile disagrees with the file's folder, + // keep the current loader and let WorkflowProfileWarningText explain the mismatch. + var canSwitchLoader = + currentModel?.Local is not { } local + || (local.SharedFolderType is SharedFolderType.DiffusionModels) == standalone; + + // The loader this card will actually end on after this call. + var effectiveStandalone = canSwitchLoader ? standalone : currentlyStandalone; + + // Reconcile the encoder-selection flag with that loader IN BOTH DIRECTIONS, even when + // no loader flip happens below. The encoder UI only exists for standalone loaders, so + // a mismatched flag strands the user either way: stuck true on a checkpoint makes + // generation demand encoders the UI doesn't show, and stuck false on a UNet loader + // hides the encoder slots with no profile toggle able to bring them back. + if (effectiveStandalone) + { + if (!IsClipModelSelectionEnabled) + IsClipModelSelectionEnabled = true; + + if (!IsVaeSelectionEnabled) + IsVaeSelectionEnabled = true; + + if (TextEncoders.Count == 0) + SetDefaultEncoderCount(); + } + else if (IsClipModelSelectionEnabled) + { + IsClipModelSelectionEnabled = false; + } + + // Only flip between the checkpoint-style (Default/Nf4) and UNet-style (Unet/Gguf) groups + // when they actually disagree — this preserves Nf4/Gguf nuances when already on the + // correct side. + if (!canSwitchLoader || standalone == currentlyStandalone) + return; + + // The two loaders read different backing fields (SelectedUnetModel vs SelectedModel), so + // carry the current selection across so the model picker doesn't blank out on the switch. + if (standalone) + { + SelectedModelLoader = ModelLoader.Unet; + SelectedUnetModel = currentModel; + } + else + { + IsVaeSelectionEnabled = false; + SelectedModelLoader = ModelLoader.Default; + SelectedModel = currentModel; + } + } + private void NotifyWorkflowProfileStateChanged() { OnPropertyChanged(nameof(ResolvedWorkflowProfile)); @@ -931,6 +1292,9 @@ private void NotifyWorkflowProfileStateChanged() OnPropertyChanged(nameof(WorkflowProfileStatusText)); OnPropertyChanged(nameof(ShowWorkflowProfileStatus)); OnPropertyChanged(nameof(RecommendedDefaultsToolTip)); + OnPropertyChanged(nameof(WorkflowProfileWarningText)); + OnPropertyChanged(nameof(ShowWorkflowProfileWarning)); + OnPropertyChanged(nameof(MoveModelToRecommendedFolderText)); } private void RefreshWorkflowProfileState() @@ -940,6 +1304,7 @@ private void RefreshWorkflowProfileState() if (!isLoadingState) { ApplyDefaultClipTypeForResolvedProfile(preserveUserSelections: true); + AutoSelectComponentsForResolvedProfile(); } } @@ -954,13 +1319,18 @@ private void ApplyDefaultClipTypeForResolvedProfile(bool preserveUserSelections) InferenceWorkflowProfile.Flux2 => "flux2", InferenceWorkflowProfile.ZImageBase or InferenceWorkflowProfile.ZImageTurbo => "lumina2", InferenceWorkflowProfile.Anima => "stable_diffusion", - InferenceWorkflowProfile.HiDream => "HiDream", + InferenceWorkflowProfile.HiDream => "hidream", _ => SelectedClipType, }; if (string.IsNullOrWhiteSpace(clipType) || SelectedClipType == clipType) return; + // The resolved profile genuinely changed. Selections WE made for the previous + // profile are defaults, not user choices — drop them so the slot count can resize + // and the new profile's defaults can fill in. Manual picks are left untouched. + ClearAutoSelectedComponents(); + SelectedClipType = clipType; SetDefaultEncoderCount(preserveUserSelections); } @@ -970,6 +1340,207 @@ private void ApplyDefaultClipTypeForResolvedProfile(bool preserveUserSelections) /// private bool isLoadingState; + /// + /// RelativePaths of encoder / VAE selections this card filled in automatically. When the + /// resolved profile or model changes, selections still in this set are treated as + /// replaceable defaults; anything else is a user choice and is never touched. + /// + private readonly HashSet autoSelectedComponentPaths = new(StringComparer.OrdinalIgnoreCase); + + /// + /// Fills empty text-encoder slots and the default VAE with known-compatible local files + /// for the resolved workflow profile (e.g. qwen_3_4b for Flux.2 Klein 4B, clip_l + t5xxl + /// for Flux.1). Saves users from having to know which encoder file pairs with which + /// model architecture — the top "where has the text encoder gone" support question. + /// + private void AutoSelectComponentsForResolvedProfile() + { + if (isLoadingState || !IsStandaloneModelLoader || SelectedUnetModel is null) + return; + + var profile = ResolvedWorkflowProfile; + + AutoSelectVae(profile); + + var matchers = GetDefaultEncoderMatchersForProfile(profile); + for (var i = 0; i < TextEncoders.Count && i < matchers.Count; i++) + { + AutoSelectEncoderSlot(TextEncoders[i], matchers[i]); + } + } + + private void AutoSelectVae(InferenceWorkflowProfile profile) + { + var current = SelectedVae; + var isEmpty = current is null || current.IsDefault || current.IsNone; + var isAutoSelected = + !isEmpty && current!.RelativePath is { } path && autoSelectedComponentPaths.Contains(path); + + if (!isEmpty && !isAutoSelected) + return; // explicit user pick — never override + + if (isAutoSelected && MatchesDefaultVaeForProfile(profile, current!.FileName)) + return; // our earlier pick is still correct for this profile + + var vae = ClientManager.VaeModels.FirstOrDefault(m => + m.Local != null && MatchesDefaultVaeForProfile(profile, m.FileName) + ); + + if (vae is null) + return; + + if (isAutoSelected) + autoSelectedComponentPaths.Remove(current!.RelativePath); + + SelectedVae = vae; + autoSelectedComponentPaths.Add(vae.RelativePath); + } + + private void AutoSelectEncoderSlot(TextEncoderSlotViewModel slot, Func matcher) + { + var current = slot.SelectedModel; + var isEmpty = current is null || current.IsNone; + var isAutoSelected = + !isEmpty && current!.RelativePath is { } path && autoSelectedComponentPaths.Contains(path); + + if (!isEmpty && !isAutoSelected) + return; // explicit user pick — never override + + if (isAutoSelected && matcher(current!.FileName)) + return; // our earlier pick is still correct (e.g. same encoder across profiles) + + var encoder = ClientManager.ClipModels.FirstOrDefault(m => m.Local != null && matcher(m.FileName)); + + if (encoder is null) + return; + + if (isAutoSelected) + autoSelectedComponentPaths.Remove(current!.RelativePath); + + slot.SelectedModel = encoder; + autoSelectedComponentPaths.Add(encoder.RelativePath); + } + + /// + /// Clears encoder / VAE selections that were auto-filled (leaving user picks alone) so a + /// profile switch can resize the slots and re-derive defaults for the new profile. + /// + private void ClearAutoSelectedComponents() + { + if (autoSelectedComponentPaths.Count == 0) + return; + + foreach (var slot in TextEncoders) + { + if ( + slot.SelectedModel is { IsNone: false, IsDefault: false } model + && autoSelectedComponentPaths.Contains(model.RelativePath) + ) + { + // The slot setter ignores null (transient ComboBox refresh guard), so clear + // with the None sentinel which the empty-slot checks already understand. + slot.SelectedModel = HybridModelFile.None; + } + } + + if ( + SelectedVae is { IsNone: false, IsDefault: false } vae + && autoSelectedComponentPaths.Contains(vae.RelativePath) + ) + { + SelectedVae = HybridModelFile.Default; + } + + autoSelectedComponentPaths.Clear(); + } + + /// + /// Per-slot filename matchers for the encoder files each profile expects, in slot order. + /// Empty for profiles where the encoder is baked into the checkpoint or unknown. + /// + private IReadOnlyList> GetDefaultEncoderMatchersForProfile( + InferenceWorkflowProfile profile + ) + { + switch (profile) + { + case InferenceWorkflowProfile.Flux2: + // Klein 4B pairs with qwen_3_4b, Klein 9B with qwen_3_8b — picking the wrong + // size fails at sampling with a tensor-shape mismatch. + var encoderSize = SelectedUnetModel is { } unet + ? Flux2KleinModelManager.GetExpectedEncoderSize(unet) + : "4b"; + return + [ + name => + Flux2KleinModelManager.IsKleinTextEncoder(name) + && Flux2KleinModelManager.MatchesEncoderSize(name, encoderSize), + ]; + case InferenceWorkflowProfile.Flux: + return + [ + name => name.Contains("clip_l", StringComparison.OrdinalIgnoreCase), + name => + name.Contains("t5xxl", StringComparison.OrdinalIgnoreCase) + || name.Contains("t5-xxl", StringComparison.OrdinalIgnoreCase), + ]; + case InferenceWorkflowProfile.ZImageBase: + case InferenceWorkflowProfile.ZImageTurbo: + // Z-Image uses the Qwen3 4B encoder (the same qwen_3_4b.safetensors file + // Flux.2 Klein 4B uses). + return + [ + name => + Flux2KleinModelManager.IsKleinTextEncoder(name) + && Flux2KleinModelManager.MatchesEncoderSize(name, "4b"), + ]; + case InferenceWorkflowProfile.Anima: + // Official Anima encoder is qwen_3_06b_base.safetensors (Qwen3 0.6B). + return + [ + name => + name.Contains("qwen_3_06b", StringComparison.OrdinalIgnoreCase) + || name.Contains("qwen3_06b", StringComparison.OrdinalIgnoreCase) + || name.Contains("qwen_3_0.6b", StringComparison.OrdinalIgnoreCase), + ]; + case InferenceWorkflowProfile.HiDream: + return + [ + name => name.Contains("clip_l", StringComparison.OrdinalIgnoreCase), + name => name.Contains("clip_g", StringComparison.OrdinalIgnoreCase), + name => + name.Contains("t5xxl", StringComparison.OrdinalIgnoreCase) + || name.Contains("t5-xxl", StringComparison.OrdinalIgnoreCase), + name => name.Contains("llama", StringComparison.OrdinalIgnoreCase), + ]; + default: + return []; + } + } + + /// + /// Whether is the known-default VAE for the profile. + /// Flux.1, Z-Image and HiDream all use the Flux.1 VAE (ae.safetensors); Flux.2 has its + /// own VAE; Anima reuses the Qwen Image VAE. + /// + private static bool MatchesDefaultVaeForProfile(InferenceWorkflowProfile profile, string fileName) => + profile switch + { + InferenceWorkflowProfile.Flux2 => Flux2KleinModelManager.IsFlux2Vae(fileName), + InferenceWorkflowProfile.Flux + or InferenceWorkflowProfile.ZImageBase + or InferenceWorkflowProfile.ZImageTurbo + or InferenceWorkflowProfile.HiDream => fileName.Equals( + "ae.safetensors", + StringComparison.OrdinalIgnoreCase + ), + InferenceWorkflowProfile.Anima => fileName.Contains( + "qwen_image_vae", + StringComparison.OrdinalIgnoreCase + ), + _ => false, + }; + /// /// Sets the default number of encoder slots based on the selected clip type. /// @@ -987,7 +1558,7 @@ private void SetDefaultEncoderCount(bool preserveUserSelections = false) "flux" => 2, "flux2" or "lumina2" or "stable_diffusion" => 1, "sd3" => 3, - "HiDream" => 4, + "hidream" => 4, _ => 2, // Default to 2 for unknown types }; @@ -1137,7 +1708,11 @@ SelectedModelLoader is ModelLoader.Default e.Builder.Connections.Base.Model = baseLoader.Output1; e.Builder.Connections.Base.VAE = baseLoader.Output3; - if (IsClipModelSelectionEnabled) + // Use separate clip loaders only when an encoder is actually configured. The flag can + // be a stale leftover (e.g. an old saved project from a UNet workflow) and checkpoints + // always bundle their own text encoder, so falling back to it beats failing validation + // for encoder slots the checkpoint UI doesn't even show. + if (IsClipModelSelectionEnabled && TextEncoders.Any(slot => slot.SelectedModel is { IsNone: false })) { SetupClipLoaders(e); } diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs index c2682bc77..bc1cd7ec0 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/SelectImageCardViewModel.cs @@ -46,15 +46,15 @@ TabContext tabContext MimeTypes = new[] { "image/jpeg", "image/png" }, }; - private readonly Lazy _lazyMaskEditorViewModel = new( - vmFactory.Get - ); + private MaskEditorViewModel? maskEditorViewModel; /// /// When true, enables a button to open a mask editor for the image. /// This is not saved or loaded from state. /// [ObservableProperty] + [NotifyPropertyChangedFor(nameof(MaskEditorViewModel))] + [NotifyPropertyChangedFor(nameof(MaskOverlayImage))] [property: JsonIgnore] [property: MemberNotNull(nameof(MaskEditorViewModel))] private bool isMaskEditorEnabled; @@ -100,14 +100,37 @@ TabContext tabContext [JsonInclude] public MaskEditorViewModel? MaskEditorViewModel => - IsMaskEditorEnabled ? _lazyMaskEditorViewModel.Value : null; + IsMaskEditorEnabled ? maskEditorViewModel ??= CreateMaskEditorViewModel() : null; + + [JsonIgnore] + public global::Avalonia.Media.IImage? MaskOverlayImage => + MaskEditorViewModel?.CachedOrNewMaskRenderImage?.Bitmap; [JsonIgnore] public ImageSource? LastMaskImage { get; private set; } + private MaskEditorViewModel CreateMaskEditorViewModel() + { + var viewModel = vmFactory.Get(); + viewModel.PropertyChanged += (_, args) => + { + if (args.PropertyName == nameof(MaskEditorViewModel.CachedOrNewMaskRenderImage)) + { + OnPropertyChanged(nameof(MaskOverlayImage)); + } + }; + return viewModel; + } + partial void OnCurrentBitmapSizeChanged(Size value) { PublishCurrentBitmapSizeToTabContext(); + + if (maskEditorViewModel is not null && !value.IsEmpty) + { + maskEditorViewModel.PaintCanvasViewModel.CanvasSize = value; + maskEditorViewModel.InvalidateCachedMaskRenderImage(); + } } partial void OnSyncBitmapSizeToTabContextChanged(bool value) @@ -252,6 +275,7 @@ private async Task OpenEditMaskDialogAsync() if (await MaskEditorViewModel.GetDialog().ShowAsync() == ContentDialogResult.Primary) { MaskEditorViewModel.InvalidateCachedMaskRenderImage(); + OnPropertyChanged(nameof(MaskOverlayImage)); } else { diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/Video/ImgToVidModelCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/Video/ImgToVidModelCardViewModel.cs index 0697e0171..99b5de194 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/Video/ImgToVidModelCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/Video/ImgToVidModelCardViewModel.cs @@ -6,6 +6,7 @@ using StabilityMatrix.Avalonia.ViewModels.Base; using StabilityMatrix.Core.Attributes; using StabilityMatrix.Core.Models.Api.Comfy.Nodes; +using StabilityMatrix.Core.Services; namespace StabilityMatrix.Avalonia.ViewModels.Inference.Video; @@ -17,9 +18,21 @@ public class ImgToVidModelCardViewModel : ModelCardViewModel public ImgToVidModelCardViewModel( IInferenceClientManager clientManager, IServiceManager vmFactory, - TabContext tabContext + TabContext tabContext, + ISettingsManager settingsManager, + IModelIndexService modelIndexService, + INotificationService notificationService, + ModelOrganizationService modelOrganizationService ) - : base(clientManager, vmFactory, tabContext) + : base( + clientManager, + vmFactory, + tabContext, + settingsManager, + modelIndexService, + notificationService, + modelOrganizationService + ) { DisableSettings = true; } diff --git a/StabilityMatrix.Avalonia/ViewModels/Settings/AccountSettingsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Settings/AccountSettingsViewModel.cs index c4dc14cd7..5beaa46a5 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Settings/AccountSettingsViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Settings/AccountSettingsViewModel.cs @@ -643,28 +643,24 @@ private async Task SetGeminiApiKey() var field = new TextBoxField { - Label = "Gemini API Key", + Label = Resources.Label_GeminiApiKey, IsPassword = true, Validator = s => { if (string.IsNullOrWhiteSpace(s)) { - throw new ValidationException("API key is required"); + throw new ValidationException(Resources.Validation_ApiKeyRequired); } }, }; var dialog = DialogHelper.CreateTextEntryDialog( - "Set Gemini API Key", - """ - Get your Gemini API key from [Google AI Studio](https://ai.google.dev/) - - This key will be used for Image Lab image generation. - """, + Resources.Label_SetGeminiApiKey, + Resources.Text_SetGeminiApiKeyDescription, null, [field] ); - dialog.PrimaryButtonText = "Save"; + dialog.PrimaryButtonText = Resources.Action_Save; if (await dialog.ShowAsync() != ContentDialogResult.Primary || field.Text is not { } apiKey) { @@ -672,7 +668,11 @@ Get your Gemini API key from [Google AI Studio](https://ai.google.dev/) } await accountsService.GeminiLoginAsync(apiKey); - notificationService.Show("Success", "Gemini API key saved", NotificationType.Success); + notificationService.Show( + Resources.Label_Success, + Resources.Text_GeminiApiKeySaved, + NotificationType.Success + ); } [RelayCommand] @@ -680,10 +680,10 @@ private async Task RemoveGeminiApiKey() { var dialog = new BetterContentDialog { - Title = "Remove Gemini API Key", - Content = "Are you sure you want to remove your Gemini API key?", - PrimaryButtonText = "Remove", - CloseButtonText = "Cancel", + Title = Resources.Label_RemoveGeminiApiKey, + Content = Resources.Text_ConfirmRemoveGeminiApiKey, + PrimaryButtonText = Resources.Action_Remove, + CloseButtonText = Resources.Action_Cancel, DefaultButton = ContentDialogButton.Close, }; @@ -693,6 +693,10 @@ private async Task RemoveGeminiApiKey() } await accountsService.GeminiLogoutAsync(); - notificationService.Show("Success", "Gemini API key removed", NotificationType.Success); + notificationService.Show( + Resources.Label_Success, + Resources.Text_GeminiApiKeyRemoved, + NotificationType.Success + ); } } diff --git a/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs index ffe5beaa3..40d3aa65f 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Settings/MainSettingsViewModel.cs @@ -65,6 +65,7 @@ using StabilityMatrix.Core.Processes; using StabilityMatrix.Core.Python; using StabilityMatrix.Core.Services; +using StabilityMatrix.Core.Services.ImageGeneration; using Symbol = FluentIcons.Common.Symbol; using SymbolIconSource = FluentIcons.Avalonia.Fluent.SymbolIconSource; @@ -1095,6 +1096,21 @@ private async Task DebugContentDialog() notificationService.Show(new Notification("Content dialog closed", $"Result: {result}")); } + /// + /// Debug: previews an Image Lab / Nano Banana generation error through the real + /// BananaVision handlers. is an + /// name, or any unparseable value (e.g. "Generic") to preview the generic failure path. + /// + [RelayCommand] + private Task DebugShowGeminiErrorDialog(string? errorCode) + { + var code = Enum.TryParse(errorCode, out var parsed) + ? parsed + : (ImageGenerationErrorCode?)null; + + return dialogFactory.Get().DebugSimulateGenerationErrorAsync(code); + } + [RelayCommand] private void DebugThrowException() { diff --git a/StabilityMatrix.Avalonia/Views/BananaVisionPage.axaml.cs b/StabilityMatrix.Avalonia/Views/BananaVisionPage.axaml.cs index 1f8f0cac4..3fcb49ba8 100644 --- a/StabilityMatrix.Avalonia/Views/BananaVisionPage.axaml.cs +++ b/StabilityMatrix.Avalonia/Views/BananaVisionPage.axaml.cs @@ -9,6 +9,7 @@ using Avalonia.Platform.Storage; using Injectio.Attributes; using StabilityMatrix.Avalonia.Controls; +using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.ViewModels; namespace StabilityMatrix.Avalonia.Views; @@ -288,7 +289,7 @@ private void OnDrop(object? sender, DragEventArgs e) private async void OnKeyDown(object? sender, KeyEventArgs e) { // Check for Ctrl+V (paste) - if (e.Key == Key.V && e.KeyModifiers.HasFlag(KeyModifiers.Control)) + if (e.Key == Key.V && e.KeyModifiers.HasFlag(PlatformKeyModifiers.CommandModifier)) { if (DataContext is BananaVisionPageViewModel viewModel) { diff --git a/StabilityMatrix.Avalonia/Views/CivArchiveDetailsPage.axaml b/StabilityMatrix.Avalonia/Views/CivArchiveDetailsPage.axaml index b19f62bf5..2a04df1e9 100644 --- a/StabilityMatrix.Avalonia/Views/CivArchiveDetailsPage.axaml +++ b/StabilityMatrix.Avalonia/Views/CivArchiveDetailsPage.axaml @@ -151,26 +151,43 @@ Foreground="{DynamicResource TextFillColorSecondaryBrush}" Text="{Binding Model.DownloadCount, Converter={StaticResource KiloFormatterConverter}}" /> - + + + + + + + + + + + + + + + + + + + +