From 486853cb309c67c77ec186da35fb0a0938adbbe5 Mon Sep 17 00:00:00 2001 From: jt Date: Tue, 9 Jun 2026 22:36:07 -0700 Subject: [PATCH 01/23] Improve inference model workflow handling --- CHANGELOG.md | 15 + .../Controls/Inference/ModelCard.axaml | 48 ++ .../Inference/InferenceWorkflowProfile.cs | 2 +- .../Services/Flux2KleinModelManager.cs | 101 +++- .../Services/Flux2KleinProvider.cs | 30 +- .../Services/ModelOrganizationService.cs | 22 + .../Services/QwenImageEditModelManager.cs | 49 +- .../BananaVisionPageViewModel.Downloads.cs | 2 +- .../BananaVisionPageViewModel.Models.cs | 119 +++- .../ViewModels/BananaVisionPageViewModel.cs | 5 +- .../InferenceTextToImageViewModel.cs | 10 - .../Inference/ModelCardViewModel.cs | 555 +++++++++++++++++- .../Video/ImgToVidModelCardViewModel.cs | 17 +- StabilityMatrix.Core/Helper/RemoteModels.cs | 38 ++ .../BananaVisionModelCategorizationTests.cs | 43 ++ .../Avalonia/Flux2KleinModelManagerTests.cs | 156 +++++ .../Avalonia/ModelCardWorkflowProfileTests.cs | 467 +++++++++++++++ .../Avalonia/ModelOrganizationServiceTests.cs | 66 +++ .../QwenImageEditModelManagerTests.cs | 109 ++++ 19 files changed, 1762 insertions(+), 92 deletions(-) create mode 100644 StabilityMatrix.Tests/Avalonia/BananaVisionModelCategorizationTests.cs create mode 100644 StabilityMatrix.Tests/Avalonia/Flux2KleinModelManagerTests.cs create mode 100644 StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs create mode 100644 StabilityMatrix.Tests/Avalonia/QwenImageEditModelManagerTests.cs diff --git a/CHANGELOG.md b/CHANGELOG.md index bf6f950a0..e7e4781a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,21 @@ 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 +### 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 +### Fixed +- 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 + ## v2.16.0 ### Added #### New Feature: 🧪 Image Lab - Conversational Image Generation for ComfyUI diff --git a/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml b/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml index bb3ba422c..945835aad 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}" /> + + + + + + + + + "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..11af97971 100644 --- a/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs +++ b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs @@ -41,10 +41,26 @@ public async Task GenerateAsync( }; } + // Resolve the user's UNET selection first — the availability check below 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. + HybridModelFile? customUnetModel = null; + if ( + request.ProviderOptions?.TryGetValue("CustomUnetModel", out var modelObj) == true + && modelObj is HybridModelFile model + ) + { + customUnetModel = model; + logger.LogInformation("Using custom UNet model: {ModelPath}", model.RelativePath); + } + var modelManager = new Flux2KleinModelManager(); - if (!modelManager.AreModelsAvailable(clientManager)) + if (!modelManager.AreModelsAvailable(clientManager, customUnetModel)) { - var modelsList = string.Join(", ", modelManager.GetMissingModelNames(clientManager)); + var modelsList = string.Join( + ", ", + modelManager.GetMissingModelNames(clientManager, customUnetModel) + ); logger.LogWarning("Required models not found: {Models}", modelsList); return new ImageGenerationResponse @@ -65,7 +81,6 @@ await ComfyImageUploadHelper.UploadImagesAsync( cancellationToken ); - HybridModelFile? customUnetModel = null; IEnumerable? loras = null; int? width = null; int? height = null; @@ -75,15 +90,6 @@ await ComfyImageUploadHelper.UploadImagesAsync( 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 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/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..1dcc2f7d5 100644 --- a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs @@ -2313,10 +2313,9 @@ private void UpdateProviderStatus() } // 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; 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..1b3b5fd33 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] @@ -318,6 +331,206 @@ 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 ( + $"The {SelectedWorkflowProfile.GetStringValue()} workflow expects a model from " + + $"DiffusionModels, but this file is in {local.SharedFolderType}. It may not load correctly.", + SharedFolderType.DiffusionModels + ); + } + + if (profileWantsStandalone is false && isStandaloneFile) + { + return ( + "This file is in DiffusionModels and can't load as an all-in-one checkpoint. " + + "It may not load correctly.", + 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 ( + $"This looks like {impliedProfile.GetStringValue()}, which loads from DiffusionModels, " + + $"but the file is in {local.SharedFolderType}. It may not load correctly.", + 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 + ? $"Move to {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( + "Could not move model", + $"A file named '{fileName}' already exists in the {mismatch.TargetFolder.GetStringValue()} folder.", + NotificationType.Error + ); + return; + } + catch (Exception e) + { + notificationService.Show( + "Could not move model", + $"[{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( + "Model moved", + $"Moved '{fileName}' to the {mismatch.TargetFolder.GetStringValue()} folder.", + NotificationType.Success + ); + } + + internal static HybridModelFile? FindMovedModel( + IEnumerable models, + string destinationRelativePath + ) => + models.FirstOrDefault(m => + m.Local != null + && string.Equals(m.RelativePath, destinationRelativePath, StringComparison.OrdinalIgnoreCase) + ); + public event Action? RecommendedDefaultsRequested; protected override void OnInitialLoaded() @@ -699,7 +912,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 +921,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. /// @@ -870,6 +1104,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 +1158,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 +1266,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 +1278,7 @@ private void RefreshWorkflowProfileState() if (!isLoadingState) { ApplyDefaultClipTypeForResolvedProfile(preserveUserSelections: true); + AutoSelectComponentsForResolvedProfile(); } } @@ -961,6 +1300,11 @@ private void ApplyDefaultClipTypeForResolvedProfile(bool preserveUserSelections) 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 +1314,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. /// @@ -1137,7 +1682,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/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.Core/Helper/RemoteModels.cs b/StabilityMatrix.Core/Helper/RemoteModels.cs index 4fd3298ac..659617646 100644 --- a/StabilityMatrix.Core/Helper/RemoteModels.cs +++ b/StabilityMatrix.Core/Helper/RemoteModels.cs @@ -562,4 +562,42 @@ private static RemoteResource ControlNetCommon(string path, string sha256) ]; #endregion + + #region Flux.2 Klein 9B Models + + /// + /// Freely downloadable models for Flux.2 Klein 9B image editing. + /// The 9B UNET itself is gated under the FLUX.2 license and can't be redistributed for + /// one-click download, so 9B users supply that file manually. This list covers the + /// pieces that are downloadable (the qwen_3_8b text encoder + shared Flux.2 VAE) + /// so a 9B install missing only its encoder is offered the matching 8B encoder rather + /// than the 4B one from . + /// + public static IReadOnlyList Flux2Klein9BModels { get; } = + [ + // Flux.2 VAE (~336 MB, shared between all Flux.2 variants) + new() + { + Url = new Uri( + "https://huggingface.co/Comfy-Org/flux2-klein-9B/resolve/main/split_files/vae/flux2-vae.safetensors" + ), + InfoUrl = new Uri("https://huggingface.co/Comfy-Org/flux2-klein-9B"), + Author = "Black Forest Labs", + LicenseType = "Apache 2.0", + ContextType = SharedFolderType.VAE, + }, + // Qwen3 8B text encoder (~16.4 GB, paired with Klein 9B) + new() + { + Url = new Uri( + "https://huggingface.co/Comfy-Org/flux2-klein-9B/resolve/main/split_files/text_encoders/qwen_3_8b.safetensors" + ), + InfoUrl = new Uri("https://huggingface.co/Comfy-Org/flux2-klein-9B"), + Author = "Alibaba Qwen", + LicenseType = "Apache 2.0", + ContextType = SharedFolderType.TextEncoders, + }, + ]; + + #endregion } diff --git a/StabilityMatrix.Tests/Avalonia/BananaVisionModelCategorizationTests.cs b/StabilityMatrix.Tests/Avalonia/BananaVisionModelCategorizationTests.cs new file mode 100644 index 000000000..20e517d50 --- /dev/null +++ b/StabilityMatrix.Tests/Avalonia/BananaVisionModelCategorizationTests.cs @@ -0,0 +1,43 @@ +using StabilityMatrix.Avalonia.ViewModels; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Database; + +namespace StabilityMatrix.Tests.Avalonia; + +[TestClass] +public class BananaVisionModelCategorizationTests +{ + [TestMethod] + public void GetModelTermMatch_OtherMetadataWithSecondaryFilename_ReturnsSecondary() + { + var model = HybridModelFile.FromLocal( + new LocalModelFile + { + RelativePath = Path.Combine("Lora", "my_flux_style.safetensors"), + SharedFolderType = SharedFolderType.Lora, + ConnectedModelInfo = new ConnectedModelInfo { BaseModel = "Other" }, + } + ); + + var match = BananaVisionPageViewModel.GetModelTermMatch(model, ["Kontext"], ["Flux"]); + + Assert.AreEqual(BananaVisionModelTermMatch.Secondary, match); + } + + [TestMethod] + public void GetModelTermMatch_KnownDifferentMetadataWithoutFilenameMatch_IsExcluded() + { + var model = HybridModelFile.FromLocal( + new LocalModelFile + { + RelativePath = Path.Combine("Lora", "watercolor_style.safetensors"), + SharedFolderType = SharedFolderType.Lora, + ConnectedModelInfo = new ConnectedModelInfo { BaseModel = "SDXL 1.0" }, + } + ); + + var match = BananaVisionPageViewModel.GetModelTermMatch(model, ["Klein", "Flux.2"], ["Flux"]); + + Assert.AreEqual(BananaVisionModelTermMatch.Excluded, match); + } +} diff --git a/StabilityMatrix.Tests/Avalonia/Flux2KleinModelManagerTests.cs b/StabilityMatrix.Tests/Avalonia/Flux2KleinModelManagerTests.cs new file mode 100644 index 000000000..3cfd76668 --- /dev/null +++ b/StabilityMatrix.Tests/Avalonia/Flux2KleinModelManagerTests.cs @@ -0,0 +1,156 @@ +using DynamicData.Binding; +using NSubstitute; +using StabilityMatrix.Avalonia.Services; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Database; + +namespace StabilityMatrix.Tests.Avalonia; + +[TestClass] +public class Flux2KleinModelManagerTests +{ + private readonly Flux2KleinModelManager manager = new(); + + private static HybridModelFile LocalModel(string fileName, SharedFolderType type) => + HybridModelFile.FromLocal(new LocalModelFile { RelativePath = fileName, SharedFolderType = type }); + + private static IInferenceClientManager CreateClientManager( + IEnumerable? unet = null, + IEnumerable? vae = null, + IEnumerable? clip = null + ) + { + var clientManager = Substitute.For(); + clientManager.UnetModels.Returns(new ObservableCollectionExtended(unet ?? [])); + clientManager.VaeModels.Returns(new ObservableCollectionExtended(vae ?? [])); + clientManager.ClipModels.Returns(new ObservableCollectionExtended(clip ?? [])); + return clientManager; + } + + private static readonly HybridModelFile Klein9BUnet = LocalModel( + "flux-2-klein-9b.safetensors", + SharedFolderType.DiffusionModels + ); + private static readonly HybridModelFile Klein4BUnet = LocalModel( + "flux-2-klein-4b.safetensors", + SharedFolderType.DiffusionModels + ); + private static readonly HybridModelFile Flux2Vae = LocalModel( + "flux2-vae.safetensors", + SharedFolderType.VAE + ); + private static readonly HybridModelFile Qwen38BEncoder = LocalModel( + "qwen_3_8b.safetensors", + SharedFolderType.TextEncoders + ); + private static readonly HybridModelFile Qwen34BEncoder = LocalModel( + "qwen_3_4b.safetensors", + SharedFolderType.TextEncoders + ); + + [TestMethod] + public void GetMissingModels_PartialNineB_OffersEightBEncoderNotFourB() + { + // User has the 9B UNET + VAE but hasn't downloaded the text encoder yet. + var clientManager = CreateClientManager(unet: [Klein9BUnet], vae: [Flux2Vae], clip: []); + + var missing = manager.GetMissingModels(clientManager).ToList(); + + // Only the encoder should be missing, and it must be the 8B one to pair with the 9B UNET. + Assert.AreEqual(1, missing.Count); + var encoder = missing.Single(); + Assert.AreEqual(SharedFolderType.TextEncoders, encoder.ContextType); + StringAssert.Contains(encoder.FileName, "qwen_3_8b"); + Assert.IsFalse( + encoder.FileName.Contains("qwen_3_4b"), + "Should not offer the 4B encoder for a 9B UNET" + ); + } + + [TestMethod] + public void GetMissingModelNames_PartialNineB_LabelsEncoderAsEightB() + { + var clientManager = CreateClientManager(unet: [Klein9BUnet], vae: [Flux2Vae], clip: []); + + var names = manager.GetMissingModelNames(clientManager).ToList(); + + CollectionAssert.AreEquivalent(new[] { "Qwen3 8B text encoder" }, names); + } + + [TestMethod] + public void GetMissingModels_CompleteNineB_OffersNothing() + { + var clientManager = CreateClientManager(unet: [Klein9BUnet], vae: [Flux2Vae], clip: [Qwen38BEncoder]); + + Assert.IsTrue(manager.AreModelsAvailable(clientManager)); + Assert.AreEqual(0, manager.GetMissingModels(clientManager).Count()); + } + + [TestMethod] + public void GetMissingModels_PartialFourB_OffersFourBEncoder() + { + var clientManager = CreateClientManager(unet: [Klein4BUnet], vae: [Flux2Vae], clip: []); + + var encoder = manager.GetMissingModels(clientManager).Single(); + + Assert.AreEqual(SharedFolderType.TextEncoders, encoder.ContextType); + StringAssert.Contains(encoder.FileName, "qwen_3_4b"); + } + + [TestMethod] + public void GetMissingModels_NineBUnetWithOnlyFourBEncoder_OffersEightBEncoder() + { + // User ran Klein 4B before (has the 4B encoder), then dropped in a 9B UNET. The 4B + // encoder doesn't pair with the 9B UNET, so the 8B encoder must still be offered. + var clientManager = CreateClientManager(unet: [Klein9BUnet], vae: [Flux2Vae], clip: [Qwen34BEncoder]); + + Assert.IsFalse( + manager.AreModelsAvailable(clientManager), + "A 4B encoder must not satisfy a 9B UNET's requirements" + ); + + var encoder = manager.GetMissingModels(clientManager).Single(); + Assert.AreEqual(SharedFolderType.TextEncoders, encoder.ContextType); + StringAssert.Contains(encoder.FileName, "qwen_3_8b"); + } + + [TestMethod] + public void AreModelsAvailable_PreferredUnetOverridesInstalledVariant() + { + // Both UNETs installed but only the 4B encoder present: availability depends on + // which UNET the user has actually selected in the dropdown. + var clientManager = CreateClientManager( + unet: [Klein4BUnet, Klein9BUnet], + vae: [Flux2Vae], + clip: [Qwen34BEncoder] + ); + + Assert.IsTrue(manager.AreModelsAvailable(clientManager, Klein4BUnet)); + Assert.IsFalse(manager.AreModelsAvailable(clientManager, Klein9BUnet)); + + var encoder = manager.GetMissingModels(clientManager, Klein9BUnet).Single(); + StringAssert.Contains(encoder.FileName, "qwen_3_8b"); + } + + [TestMethod] + public void GetMissingModels_FreshInstall_OffersFourBSet() + { + // Nothing installed - default to the freely downloadable Apache 2.0 4B set. + var clientManager = CreateClientManager(); + + var missing = manager.GetMissingModels(clientManager).ToList(); + var names = manager.GetMissingModelNames(clientManager).ToList(); + + Assert.AreEqual(3, missing.Count); + Assert.IsTrue( + missing.Any(m => m.ContextType is SharedFolderType.DiffusionModels && m.FileName.Contains("4b")) + ); + Assert.IsTrue( + missing.Any(m => + m.ContextType is SharedFolderType.TextEncoders && m.FileName.Contains("qwen_3_4b") + ) + ); + CollectionAssert.Contains(names, "Flux.2 Klein 4B UNET"); + CollectionAssert.Contains(names, "Qwen3 4B text encoder"); + } +} diff --git a/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs b/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs new file mode 100644 index 000000000..0df77aa9b --- /dev/null +++ b/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs @@ -0,0 +1,467 @@ +using DynamicData.Binding; +using NSubstitute; +using StabilityMatrix.Avalonia.Models.Inference; +using StabilityMatrix.Avalonia.Services; +using StabilityMatrix.Avalonia.ViewModels.Base; +using StabilityMatrix.Avalonia.ViewModels.Inference; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Database; +using StabilityMatrix.Core.Models.Inference; +using StabilityMatrix.Core.Services; + +namespace StabilityMatrix.Tests.Avalonia; + +[TestClass] +public class ModelCardWorkflowProfileTests +{ + private static HybridModelFile LocalModel(string fileName, SharedFolderType type) => + HybridModelFile.FromLocal(new LocalModelFile { RelativePath = fileName, SharedFolderType = type }); + + private static ModelCardViewModel CreateViewModel( + IEnumerable? clipModels = null, + IEnumerable? vaeModels = null + ) + { + var clientManager = Substitute.For(); + clientManager.ClipModels.Returns(new ObservableCollectionExtended(clipModels ?? [])); + clientManager.VaeModels.Returns(new ObservableCollectionExtended(vaeModels ?? [])); + var vmFactory = Substitute.For>(); + return new ModelCardViewModel( + clientManager, + vmFactory, + new TabContext(), + Substitute.For(), + Substitute.For(), + Substitute.For(), + new ModelOrganizationService() + ); + } + + [TestMethod] + public void SwitchingToDefaultCheckpoint_WithUnetFolderModel_KeepsLoaderAndWarns() + { + var vm = CreateViewModel(); + var anima = LocalModel("anima-base-v1.0.safetensors", SharedFolderType.DiffusionModels); + + vm.SelectedModelLoader = ModelLoader.Unet; + vm.SelectedUnetModel = anima; + + Assert.IsTrue(vm.ShowEncoderSection, "Encoder section should show for a UNet-folder model"); + + // User picks "Default / Checkpoint" — but the file physically lives in DiffusionModels, + // so CheckpointLoader could never load it. The loader must NOT flip (that would build a + // guaranteed-invalid workflow); instead the mismatch warning explains the situation. + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.DefaultCheckpoint; + + Assert.AreEqual(ModelLoader.Unet, vm.SelectedModelLoader); + Assert.AreEqual(anima, vm.SelectedUnetModel); + Assert.IsTrue(vm.ShowWorkflowProfileWarning, "Profile/folder mismatch should surface a warning"); + StringAssert.Contains(vm.WorkflowProfileWarningText, "DiffusionModels"); + } + + [TestMethod] + public void FluxProfile_WithAllInOneCheckpoint_KeepsLoaderWithoutWarning() + { + var vm = CreateViewModel(); + var fluxAllInOne = LocalModel("flux1-dev-fp8.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = fluxAllInOne; + + // Flux supports both split UNETs and all-in-one checkpoints. Choosing the Flux + // workflow should apply its sampling defaults without claiming this file is misplaced. + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Flux; + + Assert.AreEqual(ModelLoader.Default, vm.SelectedModelLoader); + Assert.AreEqual(fluxAllInOne, vm.SelectedModel); + Assert.IsFalse(vm.ShowWorkflowProfileWarning); + Assert.IsNull(vm.MoveModelToRecommendedFolderText); + } + + [TestMethod] + public void SwitchingProfiles_WithNoModelSelected_FlipsLoaderFreely() + { + var vm = CreateViewModel(); + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Flux; + Assert.AreEqual(ModelLoader.Unet, vm.SelectedModelLoader); + Assert.IsTrue(vm.ShowEncoderSection); + Assert.IsTrue(vm.ShowPrecisionSelection); + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.DefaultCheckpoint; + Assert.AreEqual(ModelLoader.Default, vm.SelectedModelLoader); + Assert.IsFalse(vm.ShowEncoderSection, "Text encoders should be hidden for a checkpoint"); + Assert.IsFalse(vm.ShowPrecisionSelection, "Precision should be hidden for a checkpoint"); + Assert.IsFalse(vm.IsVaeSelectionEnabled, "Separate VAE selector should be cleared for a checkpoint"); + Assert.IsFalse(vm.ShowWorkflowProfileWarning, "No model selected - nothing to warn about"); + } + + [TestMethod] + public void CustomProfile_DoesNotChangeModelLoader() + { + var vm = CreateViewModel(); + var anima = LocalModel("anima-base-v1.0.safetensors", SharedFolderType.DiffusionModels); + + vm.SelectedModelLoader = ModelLoader.Unet; + vm.SelectedUnetModel = anima; + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Custom; + + // Custom leaves the loader entirely to the user. + Assert.AreEqual(ModelLoader.Unet, vm.SelectedModelLoader); + Assert.IsFalse(vm.ShowWorkflowProfileWarning); + } + + [TestMethod] + public void AnimaProfile_WithCheckpointFolderModel_KeepsLoaderAndWarns() + { + // Anima is UNet-only (separate qwen_3_06b encoder + Qwen Image VAE) — there is no + // all-in-one packaging, despite Civitai filing Anima models under "Checkpoint". An + // Anima file in StableDiffusion must warn and offer the move, like Z-Image. + var vm = CreateViewModel(); + var animaCheckpoint = LocalModel("anima-base-v1.0.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = animaCheckpoint; + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Anima; + + Assert.AreEqual( + ModelLoader.Default, + vm.SelectedModelLoader, + "Loader must not flip to UNet for a checkpoint-folder file" + ); + Assert.IsTrue(vm.ShowWorkflowProfileWarning); + Assert.AreEqual("Move to DiffusionModels", vm.MoveModelToRecommendedFolderText); + } + + [TestMethod] + public void AutoProfile_AnimaInCheckpointFolder_Warns() + { + var vm = CreateViewModel(); + var animaCheckpoint = LocalModel("anima-base-v1.0.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = animaCheckpoint; + + // Auto can't run a UNet-only file as a checkpoint, so the resolved profile stays + // DefaultCheckpoint but the misplaced-file warning fires with the move offer. + Assert.AreEqual(InferenceWorkflowProfile.DefaultCheckpoint, vm.ResolvedWorkflowProfile); + Assert.IsTrue(vm.ShowWorkflowProfileWarning); + StringAssert.Contains(vm.WorkflowProfileWarningText, "DiffusionModels"); + } + + [TestMethod] + public void AutoProfile_AnimagineCheckpoint_NoAnimaDetectionOrWarning() + { + var vm = CreateViewModel(); + // Animagine XL is an SDXL model - "anima" must only match as its own token. + var animagine = LocalModel("animagine-xl-v3.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = animagine; + + Assert.AreEqual(InferenceWorkflowProfile.DefaultCheckpoint, vm.ResolvedWorkflowProfile); + Assert.IsFalse(vm.ShowWorkflowProfileWarning); + } + + [TestMethod] + public void AutoProfile_SdxlTaggedModelWithAnimaInName_DoesNotWarn() + { + var vm = CreateViewModel(); + // Explicit metadata naming a checkpoint-style base must suppress the name-based + // misplaced-file warning entirely. + var sdxlTagged = HybridModelFile.FromLocal( + new LocalModelFile + { + RelativePath = "anima_style_mix.safetensors", + SharedFolderType = SharedFolderType.StableDiffusion, + ConnectedModelInfo = new ConnectedModelInfo { BaseModel = "SDXL 1.0" }, + } + ); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = sdxlTagged; + + Assert.IsFalse(vm.ShowWorkflowProfileWarning); + } + + [TestMethod] + public void AutoProfile_UnknownMetadataDoesNotSuppressFilenameWarning() + { + var vm = CreateViewModel(); + var mistagged = HybridModelFile.FromLocal( + new LocalModelFile + { + RelativePath = "z_image_turbo_bf16.safetensors", + SharedFolderType = SharedFolderType.StableDiffusion, + ConnectedModelInfo = new ConnectedModelInfo { BaseModel = "Other" }, + } + ); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = mistagged; + + Assert.IsTrue(vm.ShowWorkflowProfileWarning); + Assert.AreEqual("Move to DiffusionModels", vm.MoveModelToRecommendedFolderText); + } + + [TestMethod] + public void AutoProfile_UnetOnlyModelInCheckpointFolder_Warns() + { + var vm = CreateViewModel(); + // The classic support thread: a Z-Image (UNet-only) file dropped into the + // StableDiffusion folder, silently loading as a checkpoint and failing. + var zImage = LocalModel("z_image_turbo_bf16.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = zImage; + + Assert.IsTrue(vm.ShowWorkflowProfileWarning); + StringAssert.Contains(vm.WorkflowProfileWarningText, "DiffusionModels"); + } + + [TestMethod] + public void SwitchingFromUnetModelToCheckpoint_ClearsEncoderRequirement() + { + // Repro of the "No text encoders configured" report: selecting a UNet model turns + // IsClipModelSelectionEnabled on, then picking an all-in-one checkpoint (e.g. Anima + // AIO in StableDiffusion) must turn it back off - the encoder UI is hidden for + // checkpoints, so a stale true makes generation demand encoders the user can't set. + var vm = CreateViewModel(); + + vm.SelectedUnifiedModel = LocalModel( + "z_image_turbo_bf16.safetensors", + SharedFolderType.DiffusionModels + ); + Assert.IsTrue(vm.IsClipModelSelectionEnabled); + + vm.SelectedUnifiedModel = LocalModel("anima-aio-v1.0.safetensors", SharedFolderType.StableDiffusion); + + Assert.AreEqual(ModelLoader.Default, vm.SelectedModelLoader); + Assert.IsFalse(vm.IsClipModelSelectionEnabled); + } + + [TestMethod] + public void SwitchingFromUnetModelToCheckpoint_ClearsAutoFilledVae() + { + var ae = LocalModel("ae.safetensors", SharedFolderType.VAE); + var qwen4B = LocalModel("qwen_3_4b.safetensors", SharedFolderType.TextEncoders); + var vm = CreateViewModel(clipModels: [qwen4B], vaeModels: [ae]); + + // Z-Image UNet auto-fills the Flux.1 VAE... + vm.SelectedUnifiedModel = LocalModel( + "z_image_turbo_bf16.safetensors", + SharedFolderType.DiffusionModels + ); + Assert.AreEqual(ae, vm.SelectedVae); + + // ...which must not stick around as a VAE override on an all-in-one checkpoint. + vm.SelectedUnifiedModel = LocalModel("anima-aio-v1.0.safetensors", SharedFolderType.StableDiffusion); + + Assert.IsTrue(vm.SelectedVae?.IsDefault ?? false, "Auto-filled VAE should reset to Default"); + } + + [TestMethod] + public void SelectingDefaultCheckpointProfile_WithLoaderAlreadyDefault_ClearsEncoderRequirement() + { + // Second half of the report: with the loader already on Default but the encoder flag + // stuck on, picking "Default / Checkpoint" must clear it even though no loader flip + // is needed. + var vm = CreateViewModel(); + + vm.SelectedUnifiedModel = LocalModel( + "z_image_turbo_bf16.safetensors", + SharedFolderType.DiffusionModels + ); + vm.SelectedModelLoader = ModelLoader.Default; + vm.IsClipModelSelectionEnabled = true; + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.DefaultCheckpoint; + + Assert.IsFalse(vm.IsClipModelSelectionEnabled); + } + + [TestMethod] + public void TogglingProfiles_OnUnetFolderModel_EncoderSectionAlwaysComesBack() + { + // Regression: picking "Default / Checkpoint" on a UNet-folder model cleared the + // encoder flag while the folder guard kept the loader on UNet — and since the loader + // never changed again, no amount of toggling brought the encoder slots back. + var vm = CreateViewModel(); + var anima = LocalModel("anima-base-v1.0.safetensors", SharedFolderType.DiffusionModels); + + vm.SelectedUnifiedModel = anima; + Assert.IsTrue(vm.ShowEncoderSection); + + // Toggle through checkpoint and back a few times, like a confused user would. + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.DefaultCheckpoint; + Assert.AreEqual(ModelLoader.Unet, vm.SelectedModelLoader, "Folder guard must keep the loader"); + Assert.IsTrue( + vm.ShowEncoderSection, + "Encoder slots must stay visible - the UNet loader still needs them at generation" + ); + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Anima; + Assert.IsTrue(vm.ShowEncoderSection); + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.DefaultCheckpoint; + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Auto; + Assert.IsTrue(vm.ShowEncoderSection, "Auto on a UNet-folder model must restore the encoder UI"); + } + + [TestMethod] + public void StandaloneProfile_WithStuckDisabledEncoderFlag_ReenablesEncoderSection() + { + // The re-enable direction: loader already UNet but the encoder flag was cleared + // (e.g. by older builds with the one-way clear). Picking any standalone profile + // must bring the encoder UI back. + var vm = CreateViewModel(); + var anima = LocalModel("anima-base-v1.0.safetensors", SharedFolderType.DiffusionModels); + + vm.SelectedUnifiedModel = anima; + vm.IsClipModelSelectionEnabled = false; + Assert.IsFalse(vm.ShowEncoderSection); + + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.Anima; + + Assert.IsTrue(vm.IsClipModelSelectionEnabled); + Assert.IsTrue(vm.ShowEncoderSection); + } + + [TestMethod] + public void DismissedWarning_StaysDismissedForSameModel_ReappearsForOthers() + { + var vm = CreateViewModel(); + var zImage = LocalModel("z_image_turbo_bf16.safetensors", SharedFolderType.StableDiffusion); + var klein = LocalModel("flux2-klein-4b.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = zImage; + Assert.IsTrue(vm.ShowWorkflowProfileWarning); + + // User clicks the dismiss button on the warning strip + vm.DismissWorkflowProfileWarningCommand.Execute(null); + Assert.IsFalse(vm.ShowWorkflowProfileWarning, "Dismissal should stick for the same model"); + + // A different mismatched model must warn again + vm.SelectedModel = klein; + Assert.IsTrue(vm.ShowWorkflowProfileWarning, "A different model should re-show the warning"); + + // Going back to the dismissed model stays quiet + vm.SelectedModel = zImage; + Assert.IsFalse(vm.ShowWorkflowProfileWarning); + } + + [TestMethod] + public void FolderMismatch_OffersMoveToDiffusionModels() + { + var vm = CreateViewModel(); + var zImage = LocalModel("z_image_turbo_bf16.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = zImage; + + Assert.AreEqual("Move to DiffusionModels", vm.MoveModelToRecommendedFolderText); + } + + [TestMethod] + public void FolderMismatch_OffersMoveToStableDiffusion() + { + var vm = CreateViewModel(); + var unetFile = LocalModel("some-model.safetensors", SharedFolderType.DiffusionModels); + + vm.SelectedModelLoader = ModelLoader.Unet; + vm.SelectedUnetModel = unetFile; + vm.SelectedWorkflowProfile = InferenceWorkflowProfile.DefaultCheckpoint; + + Assert.AreEqual("Move to StableDiffusion", vm.MoveModelToRecommendedFolderText); + } + + [TestMethod] + public void NoMismatch_NoMoveButtonText() + { + var vm = CreateViewModel(); + var checkpoint = LocalModel("dreamshaper.safetensors", SharedFolderType.StableDiffusion); + + vm.SelectedModelLoader = ModelLoader.Default; + vm.SelectedModel = checkpoint; + + Assert.IsNull(vm.MoveModelToRecommendedFolderText); + } + + [TestMethod] + public void FindMovedModel_DuplicateFilename_SelectsExactDestinationPath() + { + var nested = HybridModelFile.FromLocal( + new LocalModelFile + { + RelativePath = Path.Combine("DiffusionModels", "archive", "model.safetensors"), + SharedFolderType = SharedFolderType.DiffusionModels, + } + ); + var moved = HybridModelFile.FromLocal( + new LocalModelFile + { + RelativePath = Path.Combine("DiffusionModels", "model.safetensors"), + SharedFolderType = SharedFolderType.DiffusionModels, + } + ); + + var selected = ModelCardViewModel.FindMovedModel([nested, moved], "model.safetensors"); + + Assert.AreSame(moved, selected); + } + + [TestMethod] + public void AutoSelect_KleinUnet_FillsMatchingEncoderAndVae() + { + var qwen4B = LocalModel("qwen_3_4b.safetensors", SharedFolderType.TextEncoders); + var qwen8B = LocalModel("qwen_3_8b.safetensors", SharedFolderType.TextEncoders); + var flux2Vae = LocalModel("flux2-vae.safetensors", SharedFolderType.VAE); + + var vm = CreateViewModel(clipModels: [qwen8B, qwen4B], vaeModels: [flux2Vae]); + + vm.SelectedUnifiedModel = LocalModel("flux-2-klein-4b.safetensors", SharedFolderType.DiffusionModels); + + Assert.AreEqual(InferenceWorkflowProfile.Flux2, vm.ResolvedWorkflowProfile); + Assert.AreEqual(1, vm.TextEncoders.Count); + Assert.AreEqual(qwen4B, vm.TextEncoders[0].SelectedModel, "4B UNET must pair with qwen_3_4b"); + Assert.AreEqual(flux2Vae, vm.SelectedVae); + } + + [TestMethod] + public void AutoSelect_SwitchingKleinVariant_ReplacesAutoFilledEncoder() + { + var qwen4B = LocalModel("qwen_3_4b.safetensors", SharedFolderType.TextEncoders); + var qwen8B = LocalModel("qwen_3_8b.safetensors", SharedFolderType.TextEncoders); + + var vm = CreateViewModel(clipModels: [qwen4B, qwen8B]); + + vm.SelectedUnifiedModel = LocalModel("flux-2-klein-4b.safetensors", SharedFolderType.DiffusionModels); + Assert.AreEqual(qwen4B, vm.TextEncoders[0].SelectedModel); + + // Swapping to the 9B UNET must replace OUR earlier 4B pick with the 8B encoder - + // the mismatched pairing fails at sampling with a tensor-shape error. + vm.SelectedUnifiedModel = LocalModel("flux-2-klein-9b.safetensors", SharedFolderType.DiffusionModels); + Assert.AreEqual(qwen8B, vm.TextEncoders[0].SelectedModel); + } + + [TestMethod] + public void AutoSelect_DoesNotOverrideUserEncoderPick() + { + var qwen4B = LocalModel("qwen_3_4b.safetensors", SharedFolderType.TextEncoders); + var customEncoder = LocalModel("my_custom_encoder.safetensors", SharedFolderType.TextEncoders); + + var vm = CreateViewModel(clipModels: [qwen4B, customEncoder]); + + // User picks the model, auto-fill runs, then the user overrides the encoder manually. + vm.SelectedUnifiedModel = LocalModel("flux-2-klein-4b.safetensors", SharedFolderType.DiffusionModels); + vm.TextEncoders[0].SelectedModel = customEncoder; + + // A later model change must not stomp the user's explicit choice. + vm.SelectedUnifiedModel = LocalModel("flux-2-klein-9b.safetensors", SharedFolderType.DiffusionModels); + + Assert.AreEqual(customEncoder, vm.TextEncoders[0].SelectedModel); + } +} diff --git a/StabilityMatrix.Tests/Avalonia/ModelOrganizationServiceTests.cs b/StabilityMatrix.Tests/Avalonia/ModelOrganizationServiceTests.cs index 46bfccf14..41f5f59d7 100644 --- a/StabilityMatrix.Tests/Avalonia/ModelOrganizationServiceTests.cs +++ b/StabilityMatrix.Tests/Avalonia/ModelOrganizationServiceTests.cs @@ -1,4 +1,5 @@ using StabilityMatrix.Avalonia.Services; +using StabilityMatrix.Core.Exceptions; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Models.Api; using StabilityMatrix.Core.Models.Database; @@ -333,6 +334,71 @@ public async Task ApplyPlan_RollsBackCompletedMovesWhenLaterMoveFails() } } + [TestMethod] + public async Task MoveModelFileAsync_MovesModelAndSidecars() + { + var tempRoot = CreateTempDirectory(); + + try + { + var model = CreateModelFile( + tempRoot, + Path.Combine("StableDiffusion", "z_image_turbo.safetensors"), + "z_image_turbo.safetensors" + ); + + // Sidecars next to the model + var sourceDir = Path.Combine(tempRoot, "StableDiffusion"); + File.WriteAllText(Path.Combine(sourceDir, "z_image_turbo.cm-info.json"), "{}"); + File.WriteAllText(Path.Combine(sourceDir, "z_image_turbo.preview.jpeg"), "img"); + + var destinationDir = Path.Combine(tempRoot, "DiffusionModels"); + + await service.MoveModelFileAsync(model, tempRoot, destinationDir); + + Assert.IsTrue(File.Exists(Path.Combine(destinationDir, "z_image_turbo.safetensors"))); + Assert.IsTrue(File.Exists(Path.Combine(destinationDir, "z_image_turbo.cm-info.json"))); + Assert.IsTrue(File.Exists(Path.Combine(destinationDir, "z_image_turbo.preview.jpeg"))); + Assert.IsFalse(File.Exists(Path.Combine(sourceDir, "z_image_turbo.safetensors"))); + Assert.IsFalse(File.Exists(Path.Combine(sourceDir, "z_image_turbo.cm-info.json"))); + Assert.IsFalse(File.Exists(Path.Combine(sourceDir, "z_image_turbo.preview.jpeg"))); + } + finally + { + Directory.Delete(tempRoot, true); + } + } + + [TestMethod] + public async Task MoveModelFileAsync_DestinationExists_ThrowsWithoutMoving() + { + var tempRoot = CreateTempDirectory(); + + try + { + var model = CreateModelFile( + tempRoot, + Path.Combine("StableDiffusion", "model.safetensors"), + "model.safetensors" + ); + + var destinationDir = Path.Combine(tempRoot, "DiffusionModels"); + Directory.CreateDirectory(destinationDir); + File.WriteAllText(Path.Combine(destinationDir, "model.safetensors"), "existing"); + + await Assert.ThrowsExceptionAsync(() => + service.MoveModelFileAsync(model, tempRoot, destinationDir) + ); + + // Source must be untouched + Assert.IsTrue(File.Exists(Path.Combine(tempRoot, "StableDiffusion", "model.safetensors"))); + } + finally + { + Directory.Delete(tempRoot, true); + } + } + private static LocalModelFile CreateModelFile( string root, string relativePath, diff --git a/StabilityMatrix.Tests/Avalonia/QwenImageEditModelManagerTests.cs b/StabilityMatrix.Tests/Avalonia/QwenImageEditModelManagerTests.cs new file mode 100644 index 000000000..41ac4a852 --- /dev/null +++ b/StabilityMatrix.Tests/Avalonia/QwenImageEditModelManagerTests.cs @@ -0,0 +1,109 @@ +using DynamicData.Binding; +using NSubstitute; +using StabilityMatrix.Avalonia.Services; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Database; + +namespace StabilityMatrix.Tests.Avalonia; + +[TestClass] +public class QwenImageEditModelManagerTests +{ + private readonly QwenImageEditModelManager manager = new(); + + private static HybridModelFile LocalModel(string fileName, SharedFolderType type) => + HybridModelFile.FromLocal(new LocalModelFile { RelativePath = fileName, SharedFolderType = type }); + + private static IInferenceClientManager CreateClientManager( + IEnumerable? unet = null, + IEnumerable? vae = null, + IEnumerable? clip = null + ) + { + var clientManager = Substitute.For(); + clientManager.UnetModels.Returns(new ObservableCollectionExtended(unet ?? [])); + clientManager.VaeModels.Returns(new ObservableCollectionExtended(vae ?? [])); + clientManager.ClipModels.Returns(new ObservableCollectionExtended(clip ?? [])); + return clientManager; + } + + private static readonly HybridModelFile QwenUnet = LocalModel( + "qwen_image_edit_2511_fp8mixed.safetensors", + SharedFolderType.DiffusionModels + ); + private static readonly HybridModelFile QwenVae = LocalModel( + "qwen_image_vae.safetensors", + SharedFolderType.VAE + ); + private static readonly HybridModelFile Vl7BEncoder = LocalModel( + "qwen_2.5_vl_7b_fp8_scaled.safetensors", + SharedFolderType.TextEncoders + ); + private static readonly HybridModelFile Vl3BEncoder = LocalModel( + "qwen_2.5_vl_3b_instruct.safetensors", + SharedFolderType.TextEncoders + ); + + [TestMethod] + public void AreModelsAvailable_WrongSizeVlEncoder_ReportsMissing() + { + // A 3B VL encoder (hidden size 2048) loads but dies mid-sampling with + // "expected input with shape [*, 3584]" - it must not count as available. + var clientManager = CreateClientManager(unet: [QwenUnet], vae: [QwenVae], clip: [Vl3BEncoder]); + + Assert.IsFalse(manager.AreModelsAvailable(clientManager)); + + var encoder = manager.GetMissingModels(clientManager).Single(); + Assert.AreEqual(SharedFolderType.TextEncoders, encoder.ContextType); + StringAssert.Contains(encoder.FileName, "7b"); + } + + [TestMethod] + public void AreModelsAvailable_SevenBEncoder_IsAvailable() + { + var clientManager = CreateClientManager(unet: [QwenUnet], vae: [QwenVae], clip: [Vl7BEncoder]); + + Assert.IsTrue(manager.AreModelsAvailable(clientManager)); + Assert.AreEqual(0, manager.GetMissingModels(clientManager).Count()); + } + + [TestMethod] + public void SelectModels_PrefersSevenBOverOtherSizes() + { + // Collection order intentionally puts the 3B first - selection must not be + // "first VL encoder wins". + var clientManager = CreateClientManager( + unet: [QwenUnet], + vae: [QwenVae], + clip: [Vl3BEncoder, Vl7BEncoder] + ); + + var selected = manager.SelectModels(clientManager); + + Assert.AreEqual(Vl7BEncoder, selected.ClipModel); + } + + [TestMethod] + public void SelectModels_OnlyWrongSizeEncoder_ThrowsWithActionableMessage() + { + var clientManager = CreateClientManager(unet: [QwenUnet], vae: [QwenVae], clip: [Vl3BEncoder]); + + var ex = Assert.ThrowsException(() => manager.SelectModels(clientManager)); + + StringAssert.Contains(ex.Message, "7B"); + StringAssert.Contains(ex.Message, "qwen_2.5_vl_7b_fp8_scaled.safetensors"); + } + + [TestMethod] + public void SelectModels_UnsizedVlEncoder_IsAcceptedAsFallback() + { + // A VL encoder with no size hint is likely a renamed 7B - give it a shot rather + // than refusing to run. + var renamed = LocalModel("qwen_2.5_vl_fp8_scaled.safetensors", SharedFolderType.TextEncoders); + var clientManager = CreateClientManager(unet: [QwenUnet], vae: [QwenVae], clip: [renamed]); + + var selected = manager.SelectModels(clientManager); + + Assert.AreEqual(renamed, selected.ClipModel); + } +} From cdf9b769666dec41f8cfc7fd492f3013e7bf9d26 Mon Sep 17 00:00:00 2001 From: jt Date: Tue, 9 Jun 2026 22:40:10 -0700 Subject: [PATCH 02/23] Reference fixed inference issue --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e7e4781a3..4e56c33c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2 - 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 ### 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 **"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 From c95962be2c380e9a2d8bd85a9def68e8bf744fdd Mon Sep 17 00:00:00 2001 From: jt Date: Tue, 9 Jun 2026 22:44:05 -0700 Subject: [PATCH 03/23] Clarify moved model path matching --- .../ViewModels/Inference/ModelCardViewModel.cs | 4 ++-- .../Avalonia/ModelCardWorkflowProfileTests.cs | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs index 1b3b5fd33..a2abb8a53 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs @@ -524,11 +524,11 @@ mismatch.TargetFolder is SharedFolderType.DiffusionModels internal static HybridModelFile? FindMovedModel( IEnumerable models, - string destinationRelativePath + string pathRelativeToSharedFolder ) => models.FirstOrDefault(m => m.Local != null - && string.Equals(m.RelativePath, destinationRelativePath, StringComparison.OrdinalIgnoreCase) + && string.Equals(m.RelativePath, pathRelativeToSharedFolder, StringComparison.OrdinalIgnoreCase) ); public event Action? RecommendedDefaultsRequested; diff --git a/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs b/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs index 0df77aa9b..5d5248184 100644 --- a/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs +++ b/StabilityMatrix.Tests/Avalonia/ModelCardWorkflowProfileTests.cs @@ -410,6 +410,9 @@ public void FindMovedModel_DuplicateFilename_SelectsExactDestinationPath() var selected = ModelCardViewModel.FindMovedModel([nested, moved], "model.safetensors"); + // HybridModelFile.RelativePath is relative to DiffusionModels, unlike + // LocalModelFile.RelativePath which includes the shared-folder prefix. + Assert.AreEqual("model.safetensors", moved.RelativePath); Assert.AreSame(moved, selected); } From 495be81b4ae1c34948ecda3ac783bf6e97486695 Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 00:52:23 -0700 Subject: [PATCH 04/23] Surface ComfyUI rejection details in Kontext and Qwen providers Only Flux2KleinProvider caught Refit ApiException and showed ComfyUI's node-validation message; Flux Kontext and Qwen Image Edit collapsed the same failures into a generic "Generation failed", hiding the actual problem (wrong encoder, missing file, etc.) from the user. Bring both providers up to parity with the Klein handler. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 1 + .../Services/FluxKontextProvider.cs | 22 +++++++++++++++++++ .../Services/QwenImageEditProvider.cs | 22 +++++++++++++++++++ 3 files changed, 45 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e56c33c2..3eef0786e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2 - 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) ### 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 **"No text encoders configured"** errors when generating with an all-in-one checkpoint after a UNet model had been selected in the same tab diff --git a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs index 00aaf3672..e9138f8a0 100644 --- a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs +++ b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs @@ -1,5 +1,6 @@ using AsyncAwaitBestPractices; using Microsoft.Extensions.Logging; +using Refit; using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; @@ -291,6 +292,24 @@ void OnRunningNodeChanged(object? sender, string? node) logger.LogInformation("Image generation was cancelled"); throw; // Propagate cancellation to ViewModel for proper handling } + 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 Kontext 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 Kontext"); @@ -301,4 +320,7 @@ void OnRunningNodeChanged(object? sender, string? node) }; } } + + private static string Truncate(string value, int maxLength) => + value.Length <= maxLength ? value : value[..maxLength] + "..."; } diff --git a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs index b23f32fc0..8452ecafc 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs @@ -1,5 +1,6 @@ using AsyncAwaitBestPractices; using Microsoft.Extensions.Logging; +using Refit; using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; @@ -293,6 +294,24 @@ void OnRunningNodeChanged(object? sender, string? node) logger.LogInformation("Image generation was cancelled"); throw; // Propagate cancellation to ViewModel for proper handling } + 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 Qwen Image Edit 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 Qwen Image Edit"); @@ -303,4 +322,7 @@ void OnRunningNodeChanged(object? sender, string? node) }; } } + + private static string Truncate(string value, int maxLength) => + value.Length <= maxLength ? value : value[..maxLength] + "..."; } From e37913d55db9b572489648358910eab8c5f15190 Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 11:15:31 -0700 Subject: [PATCH 05/23] Consolidate duplicated ComfyUI plumbing across Image Lab providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three Image Lab local providers (Flux Kontext, Qwen Image Edit, Flux.2 Klein) and their workflow builders carried heavily duplicated ComfyUI plumbing — identical in all three copies. Extract the genuinely shared logic into three helpers in StabilityMatrix.Avalonia/Helpers, removing ~480 lines of duplication while leaving provider-specific logic (model detection, node graphs, per-provider settings) untouched. - ComfyProgressReporter: a disposable that reports the initial "Queued" stage, subscribes to the ComfyTask progress events, dedups updates by (percent, running node), and unsubscribes on dispose. Replaces the per-provider local-function block. - ComfyGenerationHelper: SelectOutputImages (prefer "SaveImage", else first non-empty output), GetMimeTypeForFileName, and Truncate (the three copied provider helpers now share one home). - ComfyWorkflowHelper: GetReferenceImageNames parameterized over the only things that differed (max image count + filename prefix), and ApplyLoras for the LoRA-chaining loop. Also remove the vestigial IsFluxKontextAvailable observable property in BananaVisionPageViewModel — it was set seven times in UpdateProviderStatus but never read (no binding, no consumer). And log a warning when the fire-and-forget InterruptPromptAsync faults instead of swallowing it. Co-Authored-By: Claude Opus 4.8 --- .../Helpers/ComfyGenerationHelper.cs | 52 ++++++++ .../Helpers/ComfyProgressReporter.cs | 79 +++++++++++ .../Helpers/ComfyWorkflowHelper.cs | 88 +++++++++++++ .../Services/Flux2KleinProvider.cs | 120 ++++------------- .../Services/Flux2KleinWorkflowBuilder.cs | 71 ++-------- .../Services/FluxKontextProvider.cs | 123 ++++-------------- .../Services/FluxKontextWorkflowBuilder.cs | 74 ++--------- .../Services/QwenImageEditProvider.cs | 123 ++++-------------- .../Services/QwenImageEditWorkflowBuilder.cs | 67 +--------- .../ViewModels/BananaVisionPageViewModel.cs | 10 -- 10 files changed, 320 insertions(+), 487 deletions(-) create mode 100644 StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs create mode 100644 StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs create mode 100644 StabilityMatrix.Avalonia/Helpers/ComfyWorkflowHelper.cs diff --git a/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs b/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs new file mode 100644 index 000000000..b51d87340 --- /dev/null +++ b/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs @@ -0,0 +1,52 @@ +using StabilityMatrix.Core.Models.Api.Comfy; + +namespace StabilityMatrix.Avalonia.Helpers; + +/// +/// Shared logic for Image Lab providers that generate via a local ComfyUI backend +/// +public static class ComfyGenerationHelper +{ + /// + /// 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..c84b1760a --- /dev/null +++ b/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs @@ -0,0 +1,79 @@ +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 + ) + { + 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/Services/Flux2KleinProvider.cs b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs index 11af97971..f4e91d43a 100644 --- a/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs +++ b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs @@ -4,7 +4,6 @@ using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.Api.Comfy; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -159,60 +158,8 @@ await ComfyImageUploadHelper.UploadImagesAsync( 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" - ); - } - - task.ProgressUpdate += OnProgressUpdate; - task.RunningNodeChanged += OnRunningNodeChanged; + // Reports "Queued" now, then deduplicated progress updates until disposed + using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); await using var promptInterrupt = cancellationToken.Register(() => { @@ -225,47 +172,33 @@ void OnRunningNodeChanged(object? sender, string? node) var interruptCts = new CancellationTokenSource(5000); clientManager .Client.InterruptPromptAsync(interruptCts.Token) - .ContinueWith(_ => interruptCts.Dispose(), TaskScheduler.Default) + .ContinueWith( + t => + { + if (t.IsFaulted) + { + logger.LogWarning( + t.Exception, + "Failed to interrupt ComfyUI prompt {PromptId}", + task.Id + ); + } + interruptCts.Dispose(); + }, + TaskScheduler.Default + ) .SafeFireAndForget(); }); - 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; - } + 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 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; - } + var (selectedOutputKey, candidateImages) = ComfyGenerationHelper.SelectOutputImages(outputImages); if (candidateImages is null || candidateImages.Count == 0) { @@ -292,13 +225,7 @@ void OnRunningNodeChanged(object? sender, string? node) 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"; + var mimeType = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName); generatedImages.Add(new GeneratedImage { Base64Data = base64Image, MimeType = mimeType }); } @@ -340,7 +267,7 @@ void OnRunningNodeChanged(object? sender, string? node) { IsSuccess = false, ErrorMessage = - $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {Truncate(detail, 800)}", + $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {ComfyGenerationHelper.Truncate(detail, 800)}", }; } catch (Exception ex) @@ -353,7 +280,4 @@ void OnRunningNodeChanged(object? sender, string? node) }; } } - - private static string Truncate(string value, int maxLength) => - value.Length <= maxLength ? value : value[..maxLength] + "..."; } 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 e9138f8a0..f05b8fda3 100644 --- a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs +++ b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs @@ -4,7 +4,6 @@ using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.Api.Comfy; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -122,60 +121,8 @@ await ComfyImageUploadHelper.UploadImagesAsync( 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" - ); - } - - task.ProgressUpdate += OnProgressUpdate; - task.RunningNodeChanged += OnRunningNodeChanged; + // Reports "Queued" now, then deduplicated progress updates until disposed + using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); // Register cancellation to interrupt ComfyUI if the user cancels await using var promptInterrupt = cancellationToken.Register(() => @@ -189,21 +136,27 @@ void OnRunningNodeChanged(object? sender, string? node) var interruptCts = new CancellationTokenSource(5000); clientManager .Client.InterruptPromptAsync(interruptCts.Token) - .ContinueWith(_ => interruptCts.Dispose(), TaskScheduler.Default) + .ContinueWith( + t => + { + if (t.IsFaulted) + { + logger.LogWarning( + t.Exception, + "Failed to interrupt ComfyUI prompt {PromptId}", + task.Id + ); + } + 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; - } + // Wait for completion + logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); + await task.Task.WaitAsync(cancellationToken); // Get the output images var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( @@ -211,28 +164,7 @@ void OnRunningNodeChanged(object? sender, string? node) 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; - } + var (selectedOutputKey, candidateImages) = ComfyGenerationHelper.SelectOutputImages(outputImages); if (candidateImages is null || candidateImages.Count == 0) { @@ -259,13 +191,7 @@ void OnRunningNodeChanged(object? sender, string? node) 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"; + var mimeType = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName); generatedImages.Add(new GeneratedImage { Base64Data = base64Image, MimeType = mimeType }); } @@ -307,7 +233,7 @@ void OnRunningNodeChanged(object? sender, string? node) { IsSuccess = false, ErrorMessage = - $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {Truncate(detail, 800)}", + $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {ComfyGenerationHelper.Truncate(detail, 800)}", }; } catch (Exception ex) @@ -320,7 +246,4 @@ void OnRunningNodeChanged(object? sender, string? node) }; } } - - private static string Truncate(string value, int maxLength) => - value.Length <= maxLength ? value : value[..maxLength] + "..."; } 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/QwenImageEditProvider.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs index 8452ecafc..5dc2d7752 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs @@ -4,7 +4,6 @@ using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; using StabilityMatrix.Core.Models; -using StabilityMatrix.Core.Models.Api.Comfy; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -124,60 +123,8 @@ await ComfyImageUploadHelper.UploadImagesAsync( 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" - ); - } - - task.ProgressUpdate += OnProgressUpdate; - task.RunningNodeChanged += OnRunningNodeChanged; + // Reports "Queued" now, then deduplicated progress updates until disposed + using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); // Register cancellation to interrupt ComfyUI if the user cancels await using var promptInterrupt = cancellationToken.Register(() => @@ -191,21 +138,27 @@ void OnRunningNodeChanged(object? sender, string? node) var interruptCts = new CancellationTokenSource(5000); clientManager .Client.InterruptPromptAsync(interruptCts.Token) - .ContinueWith(_ => interruptCts.Dispose(), TaskScheduler.Default) + .ContinueWith( + t => + { + if (t.IsFaulted) + { + logger.LogWarning( + t.Exception, + "Failed to interrupt ComfyUI prompt {PromptId}", + task.Id + ); + } + 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; - } + // Wait for completion + logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); + await task.Task.WaitAsync(cancellationToken); // Get the output images var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( @@ -213,28 +166,7 @@ void OnRunningNodeChanged(object? sender, string? node) 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; - } + var (selectedOutputKey, candidateImages) = ComfyGenerationHelper.SelectOutputImages(outputImages); if (candidateImages is null || candidateImages.Count == 0) { @@ -261,13 +193,7 @@ void OnRunningNodeChanged(object? sender, string? node) 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"; + var mimeType = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName); generatedImages.Add(new GeneratedImage { Base64Data = base64Image, MimeType = mimeType }); } @@ -309,7 +235,7 @@ void OnRunningNodeChanged(object? sender, string? node) { IsSuccess = false, ErrorMessage = - $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {Truncate(detail, 800)}", + $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {ComfyGenerationHelper.Truncate(detail, 800)}", }; } catch (Exception ex) @@ -322,7 +248,4 @@ void OnRunningNodeChanged(object? sender, string? node) }; } } - - private static string Truncate(string value, int maxLength) => - value.Length <= maxLength ? value : value[..maxLength] + "..."; } 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.cs b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs index 1dcc2f7d5..d83eceebc 100644 --- a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs @@ -260,9 +260,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; } @@ -2277,7 +2274,6 @@ private void UpdateProviderStatus() if (!IsComfyRunning) { ProviderStatusMessage = "⚠️ ComfyUI is not running. Click Launch to start."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2286,7 +2282,6 @@ private void UpdateProviderStatus() if (IsWaitingForConnection) { ProviderStatusMessage = "🔄 Connecting to ComfyUI..."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2295,7 +2290,6 @@ private void UpdateProviderStatus() if (!ClientManager.IsConnected) { ProviderStatusMessage = "⚠️ Not connected to ComfyUI. Click Connect."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2307,7 +2301,6 @@ private void UpdateProviderStatus() if (IsDownloadingModels) { ProviderStatusMessage = DownloadProgressText ?? "⬇️ Downloading models..."; - IsFluxKontextAvailable = false; HasMissingModels = false; return; } @@ -2317,21 +2310,18 @@ private void UpdateProviderStatus() { 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; } } From c8762d42bcdd577d67d9dba27545c2ddca260e9c Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 11:22:40 -0700 Subject: [PATCH 06/23] Apply Gemini review: fail-fast null guards in ComfyProgressReporter ctor Add ArgumentNullException.ThrowIfNull for the required task/providerId constructor args of the reusable ComfyProgressReporter, matching the existing pattern used elsewhere in the codebase. The other review suggestions (null checks in SelectOutputImages/GetMimeTypeForFileName/ Truncate) were declined: their inputs are statically non-null under nullable reference types at the only call sites (ComfyImage.FileName is `required`, GetImagesForExecutedPromptAsync returns a non-null dictionary, Exception.Message is never null), so the guards would be redundant. Co-Authored-By: Claude Opus 4.8 --- StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs b/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs index c84b1760a..09b5e9493 100644 --- a/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs +++ b/StabilityMatrix.Avalonia/Helpers/ComfyProgressReporter.cs @@ -24,6 +24,9 @@ public ComfyProgressReporter( IProgress? progress ) { + ArgumentNullException.ThrowIfNull(task); + ArgumentNullException.ThrowIfNull(providerId); + this.task = task; this.providerId = providerId; this.progress = progress; From 35a28f5028601a4f6ef39b5e321d4332e451aa75 Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 11:46:02 -0700 Subject: [PATCH 07/23] Show ComfyUI error detail dialog for Image Lab failures Image Lab failures only surfaced a truncated toast/bubble message while Inference opened a JSON detail dialog for the same errors. Carry the raw ComfyUI error body (queue-time ApiException content or execution-time ComfyNodeException JSON) through ImageGenerationResponse and ImageGenerationException so the chat page can open the same DialogHelper.CreateJsonDialog Inference uses. Node-execution errors were previously caught by the generic handler and lost their details entirely; they now also get a short node-name summary in the chat. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 1 + .../Helpers/ComfyGenerationHelper.cs | 66 +++++++++++++++++++ .../Services/Flux2KleinProvider.cs | 24 +++---- .../Services/FluxKontextProvider.cs | 24 +++---- .../Services/QwenImageEditProvider.cs | 24 +++---- .../ViewModels/BananaVisionPageViewModel.cs | 38 ++++++++--- .../ImageGenerationChatService.cs | 10 ++- .../ImageGenerationException.cs | 7 ++ .../ImageGenerationResponse.cs | 7 ++ 9 files changed, 146 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3eef0786e..39e7d672f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2 - 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 ### 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 **"No text encoders configured"** errors when generating with an all-in-one checkpoint after a UNet model had been selected in the same tab diff --git a/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs b/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs index b51d87340..4c033395b 100644 --- a/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs +++ b/StabilityMatrix.Avalonia/Helpers/ComfyGenerationHelper.cs @@ -1,4 +1,8 @@ +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; @@ -7,6 +11,68 @@ namespace StabilityMatrix.Avalonia.Helpers; /// 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 diff --git a/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs index f4e91d43a..0e3864512 100644 --- a/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs +++ b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs @@ -3,6 +3,7 @@ using Refit; using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; +using StabilityMatrix.Core.Exceptions; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Services.ImageGeneration; @@ -252,23 +253,16 @@ await ComfyImageUploadHelper.UploadImagesAsync( logger.LogInformation("Image generation was cancelled"); throw; } + 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, "Flux.2 Klein", logger); + } 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}): {ComfyGenerationHelper.Truncate(detail, 800)}", - }; + // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. + return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, "Flux.2 Klein", logger); } catch (Exception ex) { diff --git a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs index f05b8fda3..889a4ce7d 100644 --- a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs +++ b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs @@ -3,6 +3,7 @@ using Refit; using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; +using StabilityMatrix.Core.Exceptions; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Services.ImageGeneration; @@ -218,23 +219,16 @@ await ComfyImageUploadHelper.UploadImagesAsync( logger.LogInformation("Image generation was cancelled"); throw; // Propagate cancellation to 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, "Flux Kontext", logger); + } 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 Kontext workflow ({StatusCode}): {Detail}", - apiEx.StatusCode, - detail - ); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {ComfyGenerationHelper.Truncate(detail, 800)}", - }; + // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. + return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, "Flux Kontext", logger); } catch (Exception ex) { diff --git a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs index 5dc2d7752..8208356c5 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs @@ -3,6 +3,7 @@ using Refit; using StabilityMatrix.Avalonia.Helpers; using StabilityMatrix.Avalonia.Models.BananaVision; +using StabilityMatrix.Core.Exceptions; using StabilityMatrix.Core.Models; using StabilityMatrix.Core.Services.ImageGeneration; @@ -220,23 +221,16 @@ await ComfyImageUploadHelper.UploadImagesAsync( logger.LogInformation("Image generation was cancelled"); throw; // Propagate cancellation to 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, "Qwen Image Edit", logger); + } 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 Qwen Image Edit workflow ({StatusCode}): {Detail}", - apiEx.StatusCode, - detail - ); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = - $"ComfyUI rejected the workflow ({(int)apiEx.StatusCode}): {ComfyGenerationHelper.Truncate(detail, 800)}", - }; + // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. + return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, "Qwen Image Edit", logger); } catch (Exception ex) { diff --git a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs index d83eceebc..7f5515c8f 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; @@ -1113,8 +1114,7 @@ private async Task SendMessageAsync(CancellationToken cancellationToken) } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -1140,6 +1140,31 @@ 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) + { + ErrorMessage = ex.Message; + + if (!string.IsNullOrWhiteSpace(ex.DetailJson)) + { + await DialogHelper + .CreateJsonDialog( + ex.DetailJson, + Resources.Label_ComfyError, + Resources.Text_ComfyReportedGenerationError + ) + .ShowAsync(); + } + else + { + notificationService.Show(Resources.Label_GenerationFailed, ex.Message, NotificationType.Warning); + } + } + /// /// Shows a dialog prompting the user to add their Gemini API key in settings /// @@ -1264,8 +1289,7 @@ private async Task RetryLastMessageAsync(CancellationToken cancellationToken) } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -1456,8 +1480,7 @@ private async Task RegenerateLastResponseAsync(CancellationToken cancellationTok } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } @@ -1698,8 +1721,7 @@ string editedText } else { - ErrorMessage = ex.Message; - notificationService.Show("Generation Failed", ex.Message, NotificationType.Warning); + await ShowGenerationFailedAsync(ex); CanRetryLastMessage = true; } } diff --git a/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationChatService.cs b/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationChatService.cs index 6fe90a94f..73d04ea19 100644 --- a/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationChatService.cs +++ b/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationChatService.cs @@ -708,7 +708,10 @@ public async Task RemoveImageFromMessageAsync(Guid messageId, string image await database.Conversations.UpdateAsync(errorUpdatedConversation).ConfigureAwait(false); // Throw exception so caller can handle it appropriately (show notification, etc.) - throw new ImageGenerationException(response.ErrorMessage ?? "Image generation failed"); + throw new ImageGenerationException(response.ErrorMessage ?? "Image generation failed") + { + DetailJson = response.ErrorDetailJson, + }; } // Save generated images @@ -998,7 +1001,10 @@ public async Task RetryGenerationAsync( conversation.UpdatedAt = DateTime.UtcNow; await database.Conversations.UpdateAsync(conversation).ConfigureAwait(false); - throw new ImageGenerationException(response.ErrorMessage ?? "Image generation failed"); + throw new ImageGenerationException(response.ErrorMessage ?? "Image generation failed") + { + DetailJson = response.ErrorDetailJson, + }; } // Save generated images diff --git a/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationException.cs b/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationException.cs index b160e8dd1..41bb7f92f 100644 --- a/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationException.cs +++ b/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationException.cs @@ -10,4 +10,11 @@ public ImageGenerationException(string message) public ImageGenerationException(string message, Exception innerException) : base(message, innerException) { } + + /// + /// Raw machine-readable error detail (e.g. ComfyUI's JSON body for a rejected workflow + /// or node execution error), carried through from the provider so the UI can offer a + /// detail dialog alongside the short message. + /// + public string? DetailJson { get; init; } } diff --git a/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationResponse.cs b/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationResponse.cs index 0f203f3ff..04cc69db4 100644 --- a/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationResponse.cs +++ b/StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationResponse.cs @@ -37,6 +37,13 @@ public record ImageGenerationResponse /// public string? ErrorMessage { get; init; } + /// + /// Raw machine-readable error detail if generation failed (e.g. ComfyUI's JSON body + /// for a rejected workflow or node execution error). Shown to the user in a detail + /// dialog; remains the short human-readable summary. + /// + public string? ErrorDetailJson { get; init; } + /// /// Provider-specific metadata /// From 922d7322e5ca00171d41cfdae80fbb4997927d48 Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 11:46:22 -0700 Subject: [PATCH 08/23] Localize folder-mismatch warning and move action strings Move the user-facing strings added for the Inference folder-mismatch warning, the one-click model move, and the Image Lab Comfy error dialog out of code and into Languages/Resources.resx so they can be translated like the rest of the UI. Co-Authored-By: Claude Fable 5 --- .../Controls/Inference/ModelCard.axaml | 4 +- .../Languages/Resources.Designer.cs | 108 ++++++++++++++++++ .../Languages/Resources.resx | 36 ++++++ .../Inference/ModelCardViewModel.cs | 40 ++++--- 4 files changed, 171 insertions(+), 17 deletions(-) diff --git a/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml b/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml index 945835aad..a227f40af 100644 --- a/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml +++ b/StabilityMatrix.Avalonia/Controls/Inference/ModelCard.axaml @@ -215,7 +215,7 @@ Padding="8,3" VerticalAlignment="Center" Command="{Binding MoveModelToRecommendedFolderCommand}" - Content="Move" + Content="{x:Static lang:Resources.Action_Move}" FontSize="11" ToolTip.Tip="{Binding MoveModelToRecommendedFolderText}" /> diff --git a/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs index 853d33248..9afa340e8 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs +++ b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs @@ -4834,5 +4834,113 @@ 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); + } + } } } diff --git a/StabilityMatrix.Avalonia/Languages/Resources.resx b/StabilityMatrix.Avalonia/Languages/Resources.resx index 8048d9ea9..bafe2b4f2 100644 --- a/StabilityMatrix.Avalonia/Languages/Resources.resx +++ b/StabilityMatrix.Avalonia/Languages/Resources.resx @@ -1744,4 +1744,40 @@ 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 + diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs index a2abb8a53..3fa57c490 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs @@ -361,19 +361,18 @@ SelectedWorkflowProfile is InferenceWorkflowProfile.Auto if (profileWantsStandalone is true && !isStandaloneFile) { return ( - $"The {SelectedWorkflowProfile.GetStringValue()} workflow expects a model from " - + $"DiffusionModels, but this file is in {local.SharedFolderType}. It may not load correctly.", + string.Format( + Resources.TextTemplate_WorkflowNeedsDiffusionModelsFile, + SelectedWorkflowProfile.GetStringValue(), + local.SharedFolderType + ), SharedFolderType.DiffusionModels ); } if (profileWantsStandalone is false && isStandaloneFile) { - return ( - "This file is in DiffusionModels and can't load as an all-in-one checkpoint. " - + "It may not load correctly.", - SharedFolderType.StableDiffusion - ); + return (Resources.Text_FileInDiffusionModelsNotCheckpoint, SharedFolderType.StableDiffusion); } } else if (SelectedWorkflowProfile is InferenceWorkflowProfile.Auto && !isStandaloneFile) @@ -400,8 +399,11 @@ or InferenceWorkflowProfile.Anima ) { return ( - $"This looks like {impliedProfile.GetStringValue()}, which loads from DiffusionModels, " - + $"but the file is in {local.SharedFolderType}. It may not load correctly.", + string.Format( + Resources.TextTemplate_FileLooksLikeStandaloneModel, + impliedProfile.GetStringValue(), + local.SharedFolderType + ), SharedFolderType.DiffusionModels ); } @@ -444,7 +446,7 @@ private void DismissWorkflowProfileWarning() /// public string? MoveModelToRecommendedFolderText => GetProfileFolderMismatch() is { } mismatch - ? $"Move to {mismatch.TargetFolder.GetStringValue()}" + ? string.Format(Resources.TextTemplate_MoveToFolder, mismatch.TargetFolder.GetStringValue()) : null; /// @@ -474,8 +476,12 @@ await modelOrganizationService.MoveModelFileAsync( catch (FileTransferExistsException) { notificationService.Show( - "Could not move model", - $"A file named '{fileName}' already exists in the {mismatch.TargetFolder.GetStringValue()} folder.", + Resources.Label_CouldNotMoveModel, + string.Format( + Resources.TextTemplate_FileAlreadyExistsInFolder, + fileName, + mismatch.TargetFolder.GetStringValue() + ), NotificationType.Error ); return; @@ -483,7 +489,7 @@ await modelOrganizationService.MoveModelFileAsync( catch (Exception e) { notificationService.Show( - "Could not move model", + Resources.Label_CouldNotMoveModel, $"[{e.GetType().Name}] {e.Message}", NotificationType.Error ); @@ -516,8 +522,12 @@ mismatch.TargetFolder is SharedFolderType.DiffusionModels ); notificationService.Show( - "Model moved", - $"Moved '{fileName}' to the {mismatch.TargetFolder.GetStringValue()} folder.", + Resources.Label_ModelMoved, + string.Format( + Resources.TextTemplate_MovedFileToFolder, + fileName, + mismatch.TargetFolder.GetStringValue() + ), NotificationType.Success ); } From b061b67fad4ab00d12e16a6fb5e2424aeb6721b3 Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 17:37:12 -0700 Subject: [PATCH 09/23] Apply Gemini review: use GetStringValue for folder names in warnings Keeps the folder-mismatch warning text consistent with every other SharedFolderType display in this file and with the canonical folder naming used by RelativePathFromSharedFolder. Co-Authored-By: Claude Fable 5 --- .../ViewModels/Inference/ModelCardViewModel.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs index 3fa57c490..a0311d818 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs @@ -364,7 +364,7 @@ SelectedWorkflowProfile is InferenceWorkflowProfile.Auto string.Format( Resources.TextTemplate_WorkflowNeedsDiffusionModelsFile, SelectedWorkflowProfile.GetStringValue(), - local.SharedFolderType + local.SharedFolderType.GetStringValue() ), SharedFolderType.DiffusionModels ); @@ -402,7 +402,7 @@ or InferenceWorkflowProfile.Anima string.Format( Resources.TextTemplate_FileLooksLikeStandaloneModel, impliedProfile.GetStringValue(), - local.SharedFolderType + local.SharedFolderType.GetStringValue() ), SharedFolderType.DiffusionModels ); From 4bad3668e132f21707ddf8998361c2795dd668d6 Mon Sep 17 00:00:00 2001 From: jt Date: Wed, 10 Jun 2026 20:26:53 -0700 Subject: [PATCH 10/23] Extract ComfyImageGenerationProviderBase template for Image Lab providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback that the three Image Lab providers still shared a near-identical GenerateAsync skeleton after the helper extraction — the connection check, prompt queue, cancellation/interrupt registration, output download loop, and four catch arms were copy-pasted in all three. Pull that whole flow into an abstract ComfyImageGenerationProviderBase (template method), mirroring the existing InferenceGenerationViewModel base pattern. Subclasses now supply only: - ProviderId / ProviderName / LogName / MaxInputImages / ProviderPrefix - GetMissingModels(request) — provider model requirements - BuildWorkflow(request) — option parsing + node graph Shared provider-option parsing (custom UNET, LoRAs, dimensions) moves to protected base helpers used by all three. Each provider drops from ~190 lines to ~45-95; net -250 lines. Provider-specific logic stays in the subclasses: Flux.2 Klein keeps its variant-aware model check and Steps/Cfg/ExplicitDimensions options. Because the providers are DI singletons the base stays stateless per request, so Klein re-resolves its UNET in both GetMissingModels and BuildWorkflow; a logSelection:false flag avoids logging the same selection twice. Co-Authored-By: Claude Opus 4.8 --- .../ComfyImageGenerationProviderBase.cs | 310 ++++++++++++++++++ .../Services/Flux2KleinProvider.cs | 301 ++++------------- .../Services/FluxKontextProvider.cs | 255 ++------------ .../Services/QwenImageEditProvider.cs | 255 ++------------ 4 files changed, 427 insertions(+), 694 deletions(-) create mode 100644 StabilityMatrix.Avalonia/Services/ComfyImageGenerationProviderBase.cs 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/Flux2KleinProvider.cs b/StabilityMatrix.Avalonia/Services/Flux2KleinProvider.cs index 0e3864512..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.Exceptions; -using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -15,263 +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.", - }; - } + protected override string LogName => "Flux.2 Klein"; - // Resolve the user's UNET selection first — the availability check below 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. - HybridModelFile? customUnetModel = null; - if ( - request.ProviderOptions?.TryGetValue("CustomUnetModel", out var modelObj) == true - && modelObj is HybridModelFile model - ) - { - customUnetModel = model; - logger.LogInformation("Using custom UNet model: {ModelPath}", model.RelativePath); - } + // Klein supports multi-reference editing; cap at 4 for predictable VRAM use. + protected override int MaxInputImages => 4; + protected override string ProviderPrefix => "flux2_klein"; - var modelManager = new Flux2KleinModelManager(); - if (!modelManager.AreModelsAvailable(clientManager, customUnetModel)) - { - var modelsList = string.Join( - ", ", - modelManager.GetMissingModelNames(clientManager, customUnetModel) - ); + 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); - 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.", - }; - } + var modelManager = new Flux2KleinModelManager(); + if (modelManager.AreModelsAvailable(ClientManager, customUnetModel)) + { + return []; + } + return modelManager.GetMissingModelNames(ClientManager, customUnetModel).ToList(); + } - // Klein supports multi-reference editing; cap at 4 for predictable VRAM use. - await ComfyImageUploadHelper.UploadImagesAsync( - clientManager, - request, - maxInputImages: 4, - providerPrefix: "flux2_klein", - logger, - cancellationToken - ); + protected override Dictionary BuildWorkflow(ImageGenerationRequest request) + { + var customUnetModel = GetCustomUnetModel(request); + var loras = GetSelectedLoras(request); + var (width, height) = GetDimensions(request); - IEnumerable? loras = null; - int? width = null; - int? height = null; - int? steps = null; - double? cfg = null; - var explicitDimensions = false; + int? steps = null; + double? cfg = null; + var explicitDimensions = false; - if (request.ProviderOptions != null) + if (request.ProviderOptions != null) + { + if ( + request.ProviderOptions.TryGetValue("ExplicitDimensions", out var explicitObj) + && explicitObj is bool eb + ) { - 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); - } + explicitDimensions = eb; } - 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); - - // Reports "Queued" now, then deduplicated progress updates until disposed - using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); - - await using var promptInterrupt = cancellationToken.Register(() => + if (request.ProviderOptions.TryGetValue("Steps", out var stepsObj) && stepsObj is int s) { - 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(); - }); - - 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); + 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 = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName); - - 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 (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, "Flux.2 Klein", logger); - } - catch (ApiException apiEx) - { - // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. - return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, "Flux.2 Klein", logger); - } - catch (Exception ex) - { - logger.LogError(ex, "Failed to generate image with Flux.2 Klein"); - return new ImageGenerationResponse - { - IsSuccess = false, - ErrorMessage = $"Generation failed: {ex.Message}", - }; } + + return Flux2KleinWorkflowBuilder.Build( + request, + ClientManager, + customUnetModel, + loras, + width, + height, + steps, + cfg, + explicitDimensions + ); } } diff --git a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs index 889a4ce7d..0ee87e744 100644 --- a/StabilityMatrix.Avalonia/Services/FluxKontextProvider.cs +++ b/StabilityMatrix.Avalonia/Services/FluxKontextProvider.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.Exceptions; -using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -13,231 +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); - - // Reports "Queued" now, then deduplicated progress updates until disposed - using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); - - // 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( - t => - { - if (t.IsFaulted) - { - logger.LogWarning( - t.Exception, - "Failed to interrupt ComfyUI prompt {PromptId}", - task.Id - ); - } - interruptCts.Dispose(); - }, - TaskScheduler.Default - ) - .SafeFireAndForget(); - }); - - // Wait for completion - logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); - await task.Task.WaitAsync(cancellationToken); - - // Get the output images - var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( - task.Id, - cancellationToken - ); - - var (selectedOutputKey, candidateImages) = ComfyGenerationHelper.SelectOutputImages(outputImages); + protected override string LogName => "Flux Kontext"; + protected override int MaxInputImages => 2; + protected override string ProviderPrefix => "flux_kontext"; - 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 = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName); - - 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 (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, "Flux Kontext", logger); - } - catch (ApiException apiEx) - { - // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. - return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, "Flux Kontext", logger); - } - 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/QwenImageEditProvider.cs b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs index 8208356c5..bc98d85ed 100644 --- a/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.cs +++ b/StabilityMatrix.Avalonia/Services/QwenImageEditProvider.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.Exceptions; -using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.Api.Comfy.Nodes; using StabilityMatrix.Core.Services.ImageGeneration; namespace StabilityMatrix.Avalonia.Services; @@ -15,231 +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); - - // Reports "Queued" now, then deduplicated progress updates until disposed - using var progressReporter = new ComfyProgressReporter(task, ProviderId, request.Progress); - - // 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( - t => - { - if (t.IsFaulted) - { - logger.LogWarning( - t.Exception, - "Failed to interrupt ComfyUI prompt {PromptId}", - task.Id - ); - } - interruptCts.Dispose(); - }, - TaskScheduler.Default - ) - .SafeFireAndForget(); - }); - - // Wait for completion - logger.LogInformation("Waiting for generation to complete (Prompt ID: {PromptId})", task.Id); - await task.Task.WaitAsync(cancellationToken); - - // Get the output images - var outputImages = await clientManager.Client.GetImagesForExecutedPromptAsync( - task.Id, - cancellationToken - ); - - var (selectedOutputKey, candidateImages) = ComfyGenerationHelper.SelectOutputImages(outputImages); + protected override string LogName => "Qwen Image Edit"; + protected override int MaxInputImages => 3; + protected override string ProviderPrefix => "qwen_image_edit"; - 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 = ComfyGenerationHelper.GetMimeTypeForFileName(comfyImage.FileName); - - 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 (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, "Qwen Image Edit", logger); - } - catch (ApiException apiEx) - { - // Queue-time rejection; ComfyUI's JSON body explains which node validation failed. - return ComfyGenerationHelper.CreateWorkflowRejectedResponse(apiEx, "Qwen Image Edit", logger); - } - 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 + ); } } From d3a550e504e52d9acff78a05d94bb5e3a817ccdb Mon Sep 17 00:00:00 2001 From: jt Date: Sat, 13 Jun 2026 16:15:05 -0700 Subject: [PATCH 11/23] Fix Image Lab Gemini errors, I2I masks, CivArchive downloads, custom encoders Addresses several Inference / Image Lab / CivArchive issues: - Masks (#1654, #1658): Skia->Avalonia bitmap conversion now copies pixels into a WriteableBitmap instead of wrapping a pointer owned by a disposed SKImage, so enabled clipping masks no longer disappear, fail to preview, or crash on save. MaskEditor disposes both cached render images; SelectImageCard exposes a MaskOverlayImage that refreshes after editing, restore, or size load. - CivArchive downloads (#1660): the primary Download button is now a SplitButton offering the inferred folder, existing subdirectories, and a custom folder picker. Install locations are modelled as a typed InstallLocationOption (no more stringly-typed "Models\" round-trip or "Custom..." equality), and the directory scan runs off the UI thread. ModelImportService sanitizes each pattern path segment and drops rooted/traversal input so downloads stay inside the models dir. - Custom UNet encoders (#1661): ClipTypes now lists ComfyUI's current DualCLIPLoader values while keeping single-encoder compatibility and migrating legacy "HiDream" casing to "hidream". - Gemini errors (#1664): providers return a machine-readable ImageGenerationErrorCode instead of a baked English string, so Image Lab shows the actual invalid-key, billing, quota, or permission error (localized) and explains that Nano Banana needs a paid-tier key with billing enabled. Account/key dialogs are localized, and the add-key dialog now gives step-by-step setup instructions with the AI Studio and Cloud billing links. Adds a debug-only "Gemini error dialogs" panel in Settings to preview each error through the real handlers, plus unit/UI tests for the path sanitization, encoder list, and pixel-ownership conversion. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 4 + .../Controls/Inference/SelectImageCard.axaml | 3 +- .../Extensions/SkiaExtensions.cs | 71 ++++++------ .../Languages/Resources.Designer.cs | 101 +++++++++++++++- .../Languages/Resources.de.resx | 12 +- .../Languages/Resources.fr-FR.resx | 12 +- .../Languages/Resources.ja-JP.resx | 12 +- .../Languages/Resources.ko-KR.resx | 12 +- .../Languages/Resources.resx | 54 ++++++++- .../Languages/Resources.zh-Hans.resx | 12 +- .../Services/ModelImportService.cs | 41 ++++--- .../ViewModels/BananaVisionPageViewModel.cs | 81 ++++++++++--- .../CivArchiveDetailsPageViewModel.cs | 108 +++++++++++++++++- .../InstallLocationOption.cs | 10 ++ .../ViewModels/Dialogs/MaskEditorViewModel.cs | 1 + .../Inference/ModelCardViewModel.cs | 30 +++-- .../Inference/SelectImageCardViewModel.cs | 32 +++++- .../Settings/AccountSettingsViewModel.cs | 34 +++--- .../Settings/MainSettingsViewModel.cs | 16 +++ .../Views/CivArchiveDetailsPage.axaml | 49 +++++--- .../Views/Settings/MainSettingsPage.axaml | 36 ++++++ .../GeminiBaseImageGenerationProvider.cs | 32 +++--- .../ImageGenerationChatService.cs | 2 + .../ImageGenerationErrorCode.cs | 9 ++ .../ImageGenerationException.cs | 2 + .../ImageGenerationResponse.cs | 5 + .../Avalonia/ModelCardWorkflowProfileTests.cs | 12 ++ .../Avalonia/ModelImportServiceTests.cs | 85 ++++++++++++++ .../SkiaExtensionsTests.cs | 42 +++++++ 29 files changed, 785 insertions(+), 135 deletions(-) create mode 100644 StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/InstallLocationOption.cs create mode 100644 StabilityMatrix.Core/Services/ImageGeneration/ImageGenerationErrorCode.cs create mode 100644 StabilityMatrix.UITests/SkiaExtensionsTests.cs diff --git a/CHANGELOG.md b/CHANGELOG.md index 39e7d672f..5462e2511 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,10 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2 - 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 ### 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 diff --git a/StabilityMatrix.Avalonia/Controls/Inference/SelectImageCard.axaml b/StabilityMatrix.Avalonia/Controls/Inference/SelectImageCard.axaml index bf6de24cc..e37b18b3b 100644 --- a/StabilityMatrix.Avalonia/Controls/Inference/SelectImageCard.axaml +++ b/StabilityMatrix.Avalonia/Controls/Inference/SelectImageCard.axaml @@ -59,9 +59,8 @@ diff --git a/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs b/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs index 1b91c5d28..de1696ad9 100644 --- a/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs +++ b/StabilityMatrix.Avalonia/Extensions/SkiaExtensions.cs @@ -1,4 +1,5 @@ using System; +using System.Runtime.InteropServices; using Avalonia; using Avalonia.Media; using Avalonia.Media.Imaging; @@ -110,7 +111,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 +120,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,34 +150,26 @@ public static Bitmap ToAvaloniaBitmap(this SKImage image, Vector dpi) { ArgumentNullException.ThrowIfNull(image, nameof(image)); - // ReSharper disable once SwitchExpressionHandlesSomeKnownEnumValuesWithExceptionInDefault - var avaloniaColorFormat = image.ColorType switch - { - SKColorType.Rgba8888 => PixelFormat.Rgba8888, - SKColorType.Bgra8888 => PixelFormat.Bgra8888, - _ => throw new NotSupportedException($"Unsupported SKColorType: {image.ColorType}") - }; - - // ReSharper disable once SwitchExpressionHandlesSomeKnownEnumValuesWithExceptionInDefault - var avaloniaAlphaFormat = image.AlphaType switch - { - SKAlphaType.Opaque => AlphaFormat.Opaque, - SKAlphaType.Premul => AlphaFormat.Premul, - SKAlphaType.Unpremul => AlphaFormat.Unpremul, - _ => throw new NotSupportedException($"Unsupported SKAlphaType: {image.AlphaType}") - }; + using var bitmap = SKBitmap.FromImage(image); + return bitmap.ToAvaloniaBitmap(dpi); + } - var pixmap = image.PeekPixels(); - var dataPointer = pixmap.GetPixels(); + private static void CopyPixelRows( + IntPtr source, + int sourceRowBytes, + IntPtr destination, + int destinationRowBytes, + int height + ) + { + var bytesPerRow = Math.Min(sourceRowBytes, destinationRowBytes); + var buffer = new byte[bytesPerRow]; - return new Bitmap( - avaloniaColorFormat, - avaloniaAlphaFormat, - dataPointer, - new PixelSize(image.Width, image.Height), - dpi, - pixmap.RowBytes - ); + for (var row = 0; row < height; row++) + { + Marshal.Copy(IntPtr.Add(source, row * sourceRowBytes), buffer, 0, bytesPerRow); + Marshal.Copy(buffer, 0, IntPtr.Add(destination, row * destinationRowBytes), bytesPerRow); + } } public static PixelFormat ToAvaloniaPixelFormat(this SKColorType colorType) @@ -180,7 +179,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/Languages/Resources.Designer.cs b/StabilityMatrix.Avalonia/Languages/Resources.Designer.cs index 9afa340e8..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. /// @@ -4942,5 +4951,95 @@ public static string Label_GenerationFailed { 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 bafe2b4f2..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 @@ -1780,4 +1780,56 @@ Account tokens will not be viewable after saving, please make a note of them if 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/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/ViewModels/BananaVisionPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs index 7f5515c8f..a3f610c19 100644 --- a/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/BananaVisionPageViewModel.cs @@ -1107,9 +1107,9 @@ 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 @@ -1147,7 +1147,8 @@ private async Task SendMessageAsync(CancellationToken cancellationToken) /// private async Task ShowGenerationFailedAsync(ImageGenerationException ex) { - ErrorMessage = ex.Message; + var errorMessage = GetLocalizedGenerationError(ex); + ErrorMessage = errorMessage; if (!string.IsNullOrWhiteSpace(ex.DetailJson)) { @@ -1161,22 +1162,52 @@ await DialogHelper } else { - notificationService.Show(Resources.Label_GenerationFailed, ex.Message, NotificationType.Warning); + 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, }; @@ -1193,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) { @@ -1282,9 +1331,9 @@ 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 @@ -1473,9 +1522,9 @@ 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 @@ -1714,9 +1763,9 @@ 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 diff --git a/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivArchiveDetailsPageViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/CheckpointBrowser/CivArchiveDetailsPageViewModel.cs index d84c50186..03bc5fec3 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,63 @@ 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 + ) + ); + + foreach ( + var directory in defaultDirectory.EnumerateDirectories( + "*", + EnumerationOptionConstants.AllDirectories + ) + ) + { + results.Add( + new InstallLocationOption( + Path.Combine("Models", Path.GetRelativePath(modelsDirectory, directory)), + directory + ) + ); + } + } + + 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/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/ModelCardViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs index a0311d818..0e5429058 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Inference/ModelCardViewModel.cs @@ -232,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(); @@ -253,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"); @@ -830,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 @@ -995,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); @@ -1303,7 +1319,7 @@ 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, }; @@ -1542,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 }; 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/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/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}}" /> - + + + + + + + + + + + + + + + + + + + +