From 391403cbd07b6a03dfdf017e11f3ec53192134ba Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Thu, 30 Jul 2026 11:36:42 -0500 Subject: [PATCH 1/7] Adds support for Final Group --- docs/Features/Upscaling.md | 9 +++ .../ComfyUIBackend/ComfyUIBackendExtension.cs | 6 +- .../ComfyUIBackend/WorkflowGenerator.cs | 58 +++++++++++++++++++ .../ComfyUIBackend/WorkflowGeneratorSteps.cs | 47 ++++++++++++++- src/Text2Image/T2IParamTypes.cs | 9 ++- 5 files changed, 125 insertions(+), 4 deletions(-) diff --git a/docs/Features/Upscaling.md b/docs/Features/Upscaling.md index 0da35122e..d026dd34e 100644 --- a/docs/Features/Upscaling.md +++ b/docs/Features/Upscaling.md @@ -2,6 +2,15 @@ (TODO) +# Upscale Stages + +There are two places to upscale, differing in what happens after: + +- **Refine / Upscale** group: the refiner model samples over the upscaled result, per **Refiner Control Percentage**. Set that to `0` for an upscale-only stage. +- **Final Stage** group: runs after the base and refiner stages are done, and nothing samples over the result, only whatever the upscaler does itself. + - **Final Upscale** stacks on top of Refiner Upscale, eg `1.5` refiner upscale and `2` final upscale is 3x total. + - **Final Upscale Method** offers the same methods as Refiner Upscale Method, minus the latent upscalers (which need a sampler after them). + # Pixel Decoder (PiD) (TODO) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs index 999ebc564..358f29c22 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs @@ -606,7 +606,7 @@ public static void AssignValuesFromRaw(JObject rawObjectInfo) } } - public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; + public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, FinalUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG; @@ -772,6 +772,10 @@ public override void OnInit() "pixel-lanczos", Group: T2IParamTypes.GroupRefiners, OrderPriority: -1, FeatureFlag: "comfyui", ChangeWeight: 1, GetValues: (session) => [.. UpscalerModels, .. PidUpscaleModels(session)], DependNonDefault: T2IParamTypes.RefinerUpscale.Type.ID )); + FinalUpscaleMethod = T2IParamTypes.Register(new("Final Upscale Method", "How to upscale the image, if upscaling is used.\nNo sampler runs after this, so latent upscalers are not available here.", + "pixel-lanczos", Group: T2IParamTypes.GroupFinalStage, OrderPriority: -1, FeatureFlag: "comfyui", ChangeWeight: 1, + GetValues: (session) => [.. RefinerUpscaleMethod.Type.GetValues(session).Where(u => !u.StartsWith("latent"))] + )); PixelDecoderModel = T2IParamTypes.Register(new("Pixel Decoder Model", "Optionally use a PiD (Pixel Diffusion Decoder) model.", "", Toggleable: true, FeatureFlag: "comfyui", Group: T2IParamTypes.GroupAdvancedModelAddons, IsAdvanced: true, Subtype: "Stable-Diffusion", ChangeWeight: 4, DoNotPreview: true, OrderPriority: 14, GetValues: (session) => T2IParamTypes.CleanModelList(Program.MainSDModels.ListModelsFor(session).Where(m => m.ModelClass?.CompatClass?.ID == "pid").OrderBy(m => m.Name).Select(m => m.Name)) diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index 286f46154..6bdee9f75 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -169,6 +169,9 @@ public JArray FinalImageOut /// If true, the generator is currently working on the pixel-decoder stage. public bool IsPixelDecoderStage = false; + /// If true, the generator is currently working on the final stage upscale. + public bool IsFinalStage = false; + /// If true, the generator is currently working on Image2Video. public bool IsImageToVideo = false; @@ -2670,6 +2673,61 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD return result; } + /// Scales raw media to an exact pixel size, or does nothing if it's already that size. + public WGNodeData ScaleRawMedia(WGNodeData raw, int width, int height, string method = "lanczos", string id = null) + { + if (raw.Width == width && raw.Height == height) + { + return raw; + } + string scaled = CreateNode("ImageScale", new JObject() + { + ["image"] = raw.Path, + ["width"] = width, + ["height"] = height, + ["upscale_method"] = method, + ["crop"] = "disabled" + }, id); + WGNodeData result = raw.WithPath([scaled, 0]); + result.Width = width; + result.Height = height; + return result; + } + + /// Upscales raw media by any of the Final Upscale Methods. + public WGNodeData CreatePixelUpscale(string method, WGNodeData media, WGNodeData vae, double scale, long seed) + { + WGNodeData raw = media.AsRawImage(vae); + int width = (int)Math.Round((raw.Width ?? UserInput.GetImageWidth()) * scale) / 16 * 16; + int height = (int)Math.Round((raw.Height ?? UserInput.GetImageHeight()) * scale) / 16 * 16; + if (method.StartsWith("pidmodel-")) + { + T2IModel pidModel = ComfyUIBackendExtension.GetPidModel(method.After("pidmodel-"), UserInput.SourceSession); + return ScaleRawMedia(CreatePixelDecode(pidModel, raw, vae, seed), width, height); + } + if (method.StartsWith("model-")) + { + string loaderNode = CreateNode("UpscaleModelLoader", new JObject() + { + ["model_name"] = method.After("model-") + }); + string upscaledNode = CreateNode("ImageUpscaleWithModel", new JObject() + { + ["upscale_model"] = NodePath(loaderNode, 0), + ["image"] = raw.Path + }); + WGNodeData upscaled = raw.WithPath([upscaledNode, 0]); + upscaled.Width = null; // the model's own scale factor is unknown here, so always correct after + upscaled.Height = null; + return ScaleRawMedia(upscaled, width, height); + } + if (method.StartsWith("pixel-")) + { + return ScaleRawMedia(raw, width, height, method.After("pixel-")); + } + throw new SwarmUserErrorException($"Upscale method '{method}' needs a sampler after it, so it can't be used as a Final Upscale Method. Use it in the Refine/Upscale group instead."); + } + /// Creates a "CLIPTextEncode" or equivalent node for the given input, applying prompt-given conditioning modifiers as relevant. public JArray CreateConditioning(string prompt, JArray clip, T2IModel model, bool isPositive, string firstId = null, bool isRefiner = false, bool isVideo = false, bool isVideoSwap = false, bool isPixelDecoder = false) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs index c2ed9900e..7eb79b826 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs @@ -65,6 +65,10 @@ public static void Register() }, -15); AddModelGenStep(g => { + if (g.IsFinalStage) + { + return; + } if (g.IsRefinerStage && g.UserInput.TryGet(T2IParamTypes.RefinerVAE, out T2IModel rvae)) { g.LoadingVAE = g.CreateVAELoader(rvae.ToString(g.ModelFolderFormat), g.HasNode("21") ? null : "21"); @@ -102,7 +106,7 @@ public static void Register() }, -14); AddModelGenStep(g => { - if (g.LoadingModelType == "negative" && !g.UserInput.Get(T2IParamTypes.NegativeModelIncludeLoras, true)) + if (g.LoadingModelType == "negative" && !g.UserInput.Get(T2IParamTypes.NegativeModelIncludeLoras, true) || g.IsFinalStage) { return; } @@ -143,6 +147,10 @@ public static void Register() }, -9); AddModelGenStep(g => { + if (g.IsFinalStage) + { + return; + } string applyTo = g.UserInput.Get(T2IParamTypes.FreeUApplyTo, null); if (g.Features.Contains("freeu") && applyTo is not null) { @@ -163,6 +171,10 @@ public static void Register() }, -8); AddModelGenStep(g => { + if (g.IsFinalStage) + { + return; + } if (g.UserInput.TryGet(ComfyUIBackendExtension.SelfAttentionGuidanceScale, out double sagScale)) { string patched = g.CreateNode("SelfAttentionGuidance", new JObject() @@ -243,6 +255,10 @@ public static void Register() }, -6); AddModelGenStep(g => { + if (g.IsFinalStage) + { + return; + } if (g.UserInput.TryGet(T2IParamTypes.SeamlessTileable, out string tileable) && tileable != "false") { string mode = "Both"; @@ -264,6 +280,10 @@ public static void Register() }, -5); AddModelGenStep(g => { + if (g.IsFinalStage) + { + return; + } if (g.UserInput.TryGet(ComfyUIBackendExtension.TeaCacheMode, out string teaCacheMode) && teaCacheMode != "disabled") { double teaCacheThreshold = g.UserInput.Get(ComfyUIBackendExtension.TeaCacheThreshold, 0.25); @@ -370,6 +390,10 @@ public static void Register() }, -4); AddModelGenStep(g => { + if (g.IsFinalStage) + { + return; + } if (g.Features.Contains("aitemplate") && g.UserInput.Get(ComfyUIBackendExtension.AITemplateParam)) { string aitLoad = g.CreateNode("AITemplateLoader", new JObject() @@ -1913,6 +1937,27 @@ void RunSegmentationProcessing(WorkflowGenerator g, bool isBeforeRefiner) RunSegmentationProcessing(g, isBeforeRefiner: false); }, 5); #endregion + #region Final Stage + AddStep(g => + { + if (!g.UserInput.TryGet(ComfyUIBackendExtension.FinalUpscaleMethod, out string method)) + { + return; + } + double scale = g.UserInput.Get(T2IParamTypes.FinalUpscale, 2); + if (scale == 1 && method.StartsWith("pixel-")) + { + return; + } + if (g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) + { + g.CurrentMedia.SaveOutput(g.CurrentVae, g.CurrentAudioVae, g.GetStableDynamicID(50000, 0)); + } + g.IsFinalStage = true; + g.CurrentMedia = g.CreatePixelUpscale(method, g.CurrentMedia, g.CurrentVae, scale, g.UserInput.Get(T2IParamTypes.Seed) + 500); + g.IsFinalStage = false; + }, 6); + #endregion #region SaveImage AddStep(g => { diff --git a/src/Text2Image/T2IParamTypes.cs b/src/Text2Image/T2IParamTypes.cs index 9b95f10a8..062639ca9 100644 --- a/src/Text2Image/T2IParamTypes.cs +++ b/src/Text2Image/T2IParamTypes.cs @@ -326,7 +326,7 @@ public static string ApplyStringEdit(string prior, string update) public static T2IRegisteredParam Prompt, NegativePrompt, AspectRatio, BackendType, RefinerMethod, FreeUApplyTo, FreeUVersion, PersonalNote, VideoFormat, VideoResolution, UnsamplerPrompt, ImageFormat, MaskBehavior, ColorCorrectionBehavior, RawResolution, SeamlessTileable, SD3TextEncs, BitDepth, Webhooks, WildcardSeedBehavior, SegmentSortOrder, SegmentTargetResolution, SegmentApplyAfter, TorchCompile, VideoExtendFormat, ExactBackendID, OverridePredictionType, OverrideOutpathFormat, Text2AudioTimeSignature, Text2AudioLanguage, Text2AudioKeyScale, Text2AudioStyle; public static T2IRegisteredParam Images, Steps, Width, Height, SideLength, BatchSize, VAETileSize, VAETileOverlap, VAETemporalTileSize, VAETemporalTileOverlap, ClipStopAtLayer, VideoFrames, VideoMotionBucket, VideoFPS, VideoSteps, RefinerSteps, CascadeLatentCompression, MaskShrinkGrow, MaskBlur, MaskGrow, SegmentMaskBlur, SegmentMaskGrow, SegmentMaskOversize, SegmentSteps, Text2VideoFrames, TrimVideoStartFrames, TrimVideoEndFrames, VideoExtendFrameOverlap; public static T2IRegisteredParam Seed, VariationSeed, WildcardSeed, Text2AudioBPM; - public static T2IRegisteredParam CFGScale, VariationSeedStrength, InitImageCreativity, InitImageResetToNorm, InitImageNoise, RefinerControl, RefinerUpscale, RefinerCFGScale, ReVisionStrength, AltResolutionHeightMult, + public static T2IRegisteredParam CFGScale, VariationSeedStrength, InitImageCreativity, InitImageResetToNorm, InitImageNoise, RefinerControl, RefinerUpscale, RefinerCFGScale, FinalUpscale, ReVisionStrength, AltResolutionHeightMult, FreeUBlock1, FreeUBlock2, FreeUSkip1, FreeUSkip2, GlobalRegionFactor, EndStepsEarly, SamplerSigmaMin, SamplerSigmaMax, SamplerRho, VideoAugmentationLevel, VideoCFG, VideoMinCFG, Video2VideoCreativity, VideoSwapPercent, VideoExtendSwapPercent, IP2PCFG2, RegionalObjectCleanupFactor, SigmaShift, SegmentThresholdMax, SegmentCFGScale, FluxGuidanceScale, Text2AudioDuration; public static T2IRegisteredParam InitImage, MaskImage, VideoEndFrame; public static T2IRegisteredParam VideoAudioInput, VideoAudioReference; @@ -336,7 +336,7 @@ public static string ApplyStringEdit(string prior, string update) public static T2IRegisteredParam OutputIntermediateImages, DoNotSave, DoNotSaveIntermediates, ControlNetPreviewOnly, RevisionZeroPrompt, RemoveBackground, NoSeedIncrement, NoPreviews, VideoBoomerang, ModelSpecificEnhancements, UseInpaintingEncode, MaskCompositeUnthresholded, SaveSegmentMask, InitImageRecompositeMask, UseReferenceOnly, RefinerDoTiling, AutomaticVAE, ZeroNegative, FluxDisableGuidance, SmartImagePromptResizing, NoLoadModels, NoInternalSpecialHandling, ForwardRawBackendData, ForwardSwarmData, NegativeModelIncludeLoras, ContinueAfterErrors, PlaceholderParamGroupStarred, PlaceholderParamGroupUser1, PlaceholderParamGroupUser2, PlaceholderParamGroupUser3; - public static T2IParamGroup GroupImagePrompting, GroupCore, GroupVariation, GroupResolution, GroupSampling, GroupInitImage, GroupRefiners, GroupRefinerOverrides, + public static T2IParamGroup GroupImagePrompting, GroupCore, GroupVariation, GroupResolution, GroupSampling, GroupInitImage, GroupRefiners, GroupRefinerOverrides, GroupFinalStage, GroupAdvancedModelAddons, GroupSwarmInternal, GroupFreeU, GroupRegionalPrompting, GroupSegmentRefining, GroupSegmentOverrides, GroupAdvancedSampling, GroupAlternateGuidance, GroupVideo, GroupText2Video, GroupAdvancedVideo, GroupAdvancedVideoObscure, GroupVideoExtend, GroupText2Audio, GroupStarred, GroupUser1, GroupUser2, GroupUser3; @@ -552,6 +552,11 @@ static List listVaes(Session s) + "Too-high values can cause corrupted/burnt images, too-low can cause nonsensical images.\n7 is a good baseline. Normal usages vary between 4 and 9.\nSome model types, such as Turbo, expect CFG around 1.", "7", Min: 0, Max: 100, ViewMax: 20, Step: 0.5, Examples: ["5", "6", "7", "8", "9"], OrderPriority: -4, ViewType: ParamViewType.SLIDER, Group: GroupRefinerOverrides, ChangeWeight: -3, Toggleable: true, IsAdvanced: true )); + // ================================================ Final Stage ================================================ + GroupFinalStage = new("Final Stage", Toggles: true, Open: false, OrderPriority: -2.5, Description: "An upscale applied after the base and refiner stages are done.\nUnlike the Refine/Upscale group, nothing samples over the result after, only whatever the upscaler does itself."); + FinalUpscale = Register(new("Final Upscale", "How much to upscale by in the final stage.\nThis stacks on top of Refiner Upscale, so 1.5 refiner upscale and 2 final upscale is 3x total.", + "2", Min: 0.25, Max: 8, ViewMax: 4, Step: 0.25, OrderPriority: -2, ViewType: ParamViewType.SLIDER, Group: GroupFinalStage, DoNotPreview: true, Examples: ["1.5", "2", "4"] + )); // ================================================ ControlNet ================================================ for (int i = 1; i <= 3; i++) { From c8254f0e26eb8104d3cea7bcfa1ba76a38ac068f Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Thu, 30 Jul 2026 13:31:45 -0500 Subject: [PATCH 2/7] Adds support for *2v workflows in Final Stage --- .../ComfyUIBackend/WorkflowGenerator.cs | 21 +++++++++++++ .../ComfyUIBackend/WorkflowGeneratorSteps.cs | 30 ++++++------------- 2 files changed, 30 insertions(+), 21 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index 6bdee9f75..96bb60416 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -2694,6 +2694,27 @@ public WGNodeData ScaleRawMedia(WGNodeData raw, int width, int height, string me return result; } + /// Runs the Final Stage upscale over the current media, if the user configured one and it hasn't run yet. + public void RunFinalStage() + { + if (!UserInput.TryGet(ComfyUIBackendExtension.FinalUpscaleMethod, out string method)) + { + return; + } + double scale = UserInput.Get(T2IParamTypes.FinalUpscale, 2); + if (scale == 1 && method.StartsWith("pixel-")) + { + return; + } + if (UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) + { + CurrentMedia.SaveOutput(CurrentVae, CurrentAudioVae, GetStableDynamicID(50000, 0)); + } + IsFinalStage = true; + CurrentMedia = CreatePixelUpscale(method, CurrentMedia, CurrentVae, scale, UserInput.Get(T2IParamTypes.Seed) + 500); + IsFinalStage = false; + } + /// Upscales raw media by any of the Final Upscale Methods. public WGNodeData CreatePixelUpscale(string method, WGNodeData media, WGNodeData vae, double scale, long seed) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs index 7eb79b826..55a1c756e 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs @@ -1937,27 +1937,6 @@ void RunSegmentationProcessing(WorkflowGenerator g, bool isBeforeRefiner) RunSegmentationProcessing(g, isBeforeRefiner: false); }, 5); #endregion - #region Final Stage - AddStep(g => - { - if (!g.UserInput.TryGet(ComfyUIBackendExtension.FinalUpscaleMethod, out string method)) - { - return; - } - double scale = g.UserInput.Get(T2IParamTypes.FinalUpscale, 2); - if (scale == 1 && method.StartsWith("pixel-")) - { - return; - } - if (g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) - { - g.CurrentMedia.SaveOutput(g.CurrentVae, g.CurrentAudioVae, g.GetStableDynamicID(50000, 0)); - } - g.IsFinalStage = true; - g.CurrentMedia = g.CreatePixelUpscale(method, g.CurrentMedia, g.CurrentVae, scale, g.UserInput.Get(T2IParamTypes.Seed) + 500); - g.IsFinalStage = false; - }, 6); - #endregion #region SaveImage AddStep(g => { @@ -2022,6 +2001,10 @@ void RunSegmentationProcessing(WorkflowGenerator g, bool isBeforeRefiner) bool willHaveFollowupVideo = g.UserInput.TryGet(T2IParamTypes.VideoModel, out _) || g.UserInput.Get(T2IParamTypes.Prompt, "").Contains(" 1) { if (g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) @@ -2292,6 +2279,7 @@ WGNodeData appendAudio(WGNodeData combinedAudio, WGNodeData nextAudio) } g.CurrentMedia = conjoinedLast; g.CurrentMedia.FPS = videoFps ?? g.CurrentMedia.FPS; + g.RunFinalStage(); if (g.UserInput.TryGet(ComfyUIBackendExtension.VideoFrameInterpolationMethod, out string method) && g.UserInput.TryGet(ComfyUIBackendExtension.VideoFrameInterpolationMultiplier, out int mult) && mult > 1) { if (saveIntermediate) From f10a30de681d06714258e30deba731acc0387af5 Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Thu, 30 Jul 2026 14:30:00 -0500 Subject: [PATCH 3/7] Adds native SeedVR2 support --- docs/Features/Upscaling.md | 8 ++ .../ComfyUIBackend/ComfyUIBackendExtension.cs | 43 +++++++++-- .../ComfyUIBackend/WorkflowGenerator.cs | 77 +++++++++++++++++++ .../WorkflowGeneratorModelSupport.cs | 15 +++- .../ComfyUIBackend/WorkflowGeneratorSteps.cs | 40 +++++++--- src/Text2Image/CommonModels.cs | 1 + src/Text2Image/T2IModelClassSorter.cs | 12 +++ 7 files changed, 175 insertions(+), 21 deletions(-) diff --git a/docs/Features/Upscaling.md b/docs/Features/Upscaling.md index d026dd34e..b76fdb0a4 100644 --- a/docs/Features/Upscaling.md +++ b/docs/Features/Upscaling.md @@ -16,3 +16,11 @@ There are two places to upscale, differing in what happens after: (TODO) Downloads here: + +# SeedVR2 + +- [SeedVR2]() is a one-step restoration model for upscaling images or videos. + - Models can be downloaded here: [Comfy-Org/SeedVR2]() + - Save in `diffusion_models` + - The VAE will be automatically downloaded +- Select it as your `Refiner Upscale Method` or your `Final Upscale Method`, see [Upscale Stages](#upscale-stages) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs index 358f29c22..c8d037da8 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs @@ -35,10 +35,10 @@ public record class ComfyCustomWorkflow(string Name, string Workflow, string Pro public static ConcurrentDictionary CustomWorkflows = new(); /// Set of all feature-ids supported by ComfyUI backends. - public static HashSet FeaturesSupported = ["comfyui", "refiners", "controlnet", "endstepsearly", "seamless", "video", "variation_seed", "freeu", "yolov8"]; + public static HashSet FeaturesSupported = ["comfyui", "refiners", "controlnet", "endstepsearly", "seamless", "video", "variation_seed", "freeu", "yolov8", "seedvr2"]; /// Set of feature-ids that were added presumptively during loading and should be removed if the backend turns out to be missing them. - public static HashSet FeaturesDiscardIfNotFound = ["variation_seed", "freeu", "yolov8"]; + public static HashSet FeaturesDiscardIfNotFound = ["variation_seed", "freeu", "yolov8", "seedvr2"]; /// Extensible map of ComfyUI Node IDs to supported feature IDs. public static Dictionary NodeToFeatureMap = new() @@ -68,7 +68,8 @@ public record class ComfyCustomWorkflow(string Name, string Workflow, string Pro ["TeaCache"] = "teacache", ["TeaCacheForVidGen"] = "teacache", ["TeaCacheForImgGen"] = "teacache_oldvers", - ["OverrideCLIPDevice"] = "set_clip_device" + ["OverrideCLIPDevice"] = "set_clip_device", + ["SeedVR2Conditioning"] = "seedvr2", }; /// @@ -606,13 +607,13 @@ public static void AssignValuesFromRaw(JObject rawObjectInfo) } } - public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, FinalUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; + public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, FinalUpscaleMethod, SeedVR2ColorCorrectionBehavior, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; - public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG; + public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG, SeedVR2SplitLatent; public static T2IRegisteredParam IPAdapterWeight, IPAdapterStart, IPAdapterEnd, SelfAttentionGuidanceScale, SelfAttentionGuidanceSigmaBlur, PerturbedAttentionGuidanceScale, StyleModelMergeStrength, StyleModelApplyStart, StyleModelMultiplyStrength, RescaleCFGMultiplier, TeaCacheThreshold, TeaCacheStart, NunchakuCacheThreshold, EasyCacheThreshold, EasyCacheStart, EasyCacheEnd, RenormCFG, NormalizedAttentionGuidanceScale, NormalizedAttentionGuidanceAlpha, NormalizedAttentionGuidanceTau; - public static T2IRegisteredParam RefinerHyperTile, VideoFrameInterpolationMultiplier; + public static T2IRegisteredParam RefinerHyperTile, VideoFrameInterpolationMultiplier, SeedVR2TemporalVideoOverlap; public static T2IRegisteredParam PixelDecoderModel; @@ -656,6 +657,25 @@ public static T2IModel GetPidModel(string name, Session session) return model; } + /// Lists SeedVR2 upscaler models. + public static List SeedVR2UpscaleModels(Session session) => [.. Program.MainSDModels.ListModelsFor(session).Where(m => m.ModelClass?.CompatClass?.ID == "seedvr2").OrderBy(m => m.Name).Select(m => $"seedvr2model-{m.Name}///SeedVR2 Model: {m.Name}")]; + + /// Resolves a SeedVR2 model from a model name. + public static T2IModel GetSeedVR2Model(string name, Session session) + { + string matched = T2IParamTypes.GetBestModelInList(name, Program.MainSDModels.ListModelNamesFor(session)); + if (matched is not null && matched.EndsWith(".safetensors")) + { + matched = matched.BeforeLast('.'); + } + T2IModel model = matched is null ? null : Program.MainSDModels.GetModel(matched); + if (model is null || model.ModelClass?.CompatClass?.ID != "seedvr2") + { + throw new SwarmUserErrorException($"SeedVR2 model '{name}' could not be found, or is not a valid SeedVR2 model."); + } + return model; + } + public static List IPAdapterModels = ["None"], IPAdapterWeightTypes = ["standard", "prompt is more important", "style transfer"]; public static List GligenModels = ["None"], YoloModels = [], StyleModels = ["None"], SetClipDevices = ["cpu"]; @@ -770,7 +790,7 @@ public override void OnInit() )); RefinerUpscaleMethod = T2IParamTypes.Register(new("Refiner Upscale Method", "How to upscale the image, if upscaling is used.", "pixel-lanczos", Group: T2IParamTypes.GroupRefiners, OrderPriority: -1, FeatureFlag: "comfyui", ChangeWeight: 1, - GetValues: (session) => [.. UpscalerModels, .. PidUpscaleModels(session)], DependNonDefault: T2IParamTypes.RefinerUpscale.Type.ID + GetValues: (session) => [.. UpscalerModels, .. PidUpscaleModels(session), .. SeedVR2UpscaleModels(session)], DependNonDefault: T2IParamTypes.RefinerUpscale.Type.ID )); FinalUpscaleMethod = T2IParamTypes.Register(new("Final Upscale Method", "How to upscale the image, if upscaling is used.\nNo sampler runs after this, so latent upscalers are not available here.", "pixel-lanczos", Group: T2IParamTypes.GroupFinalStage, OrderPriority: -1, FeatureFlag: "comfyui", ChangeWeight: 1, @@ -803,6 +823,15 @@ public override void OnInit() RefinerHyperTile = T2IParamTypes.Register(new("Refiner HyperTile", "The size of hypertiles to use for the refining stage.\nHyperTile is a technique to speed up sampling of large images by tiling the image and batching the tiles.\nThis is useful when using SDv1 models as the refiner. SDXL-Base models do not benefit as much.", "256", Min: 64, Max: 2048, Step: 32, Toggleable: true, IsAdvanced: true, FeatureFlag: "comfyui", ViewType: ParamViewType.POT_SLIDER, Group: T2IParamTypes.GroupAdvancedSampling, OrderPriority: 20 )); + SeedVR2ColorCorrectionBehavior = T2IParamTypes.Register(new("SeedVR2 Color Correction Behavior", "How to match the colors of a SeedVR2 upscale back to the image it was given.\n'None' = Do not attempt color correction, only align the geometry.\n'CIELAB' = Transfer the color in CIELAB space, preserving detail.\n'Wavelet' = Transfer the low-frequency color, keeping the upscaled high-frequency detail.\n'AdaIN' = Match the per-channel mean and standard deviation.", + "none", FeatureFlag: "seedvr2", Group: T2IParamTypes.GroupAdvancedSampling, IsAdvanced: true, OrderPriority: 21, GetValues: (_) => ["none///None", "lab///CIELAB", "wavelet///Wavelet", "adain///AdaIN"] + )); + SeedVR2SplitLatent = T2IParamTypes.Register(new("SeedVR2 Split Latent", "If enabled, samples a SeedVR2 video upscale as chunks of frames instead of all at once, sized to fit in free VRAM.\nChunking reduces VRAM consumption.\nDoes nothing to a single image, or to a video that already fits.", + "false", IgnoreIf: "false", FeatureFlag: "seedvr2", Group: T2IParamTypes.GroupAdvancedSampling, IsAdvanced: true, OrderPriority: 22 + )); + SeedVR2TemporalVideoOverlap = T2IParamTypes.Register(new("SeedVR2 Temporal Video Overlap", "Overrides 'VAE Temporal Tile Overlap' for 'SeedVR2 Split Latent' chunks.\nHigher overlap hides the chunk seams better but takes longer.", + "0", Min: 0, Max: 4096, Step: 1, Toggleable: true, VisibleNormally: false, IsAdvanced: true, FeatureFlag: "seedvr2", Group: T2IParamTypes.GroupAdvancedSampling, OrderPriority: 23 + )); List interpolators = ["RIFE", "FILM", "GIMM-VFI"]; VideoPreviewType = T2IParamTypes.Register(new("Video Preview Type", "How to display previews for generating videos.\n'Animate' shows a low-res animated video preview.\n'iterate' shows one frame at a time while it goes.\n'one' displays just the first frame.\n'none' disables previews.", "animate", IgnoreIf: "animate", FeatureFlag: "comfyui", Group: T2IParamTypes.GroupAdvancedVideo, Permission: Permissions.ParamVideo, IsAdvanced: true, GetValues: (_) => ["animate", "iterate", "one", "none"] diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index 96bb60416..c9095e6e8 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -2673,6 +2673,78 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD return result; } + /// Creates a SeedVR2 restoration stage. + public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, WGNodeData decodeVae, long seed) + { + if (!Features.Contains("seedvr2")) + { + throw new SwarmUserErrorException($"Cannot use SeedVR2 model '{seedVrModel.Name}', the backend is missing SeedVR2 support. Update ComfyUI."); + } + WGNodeData raw = media.AsRawImage(decodeVae); + JArray resized = raw.Path; + string preprocessed = CreateNode("SeedVR2Preprocess", new JObject() + { + ["resized_images"] = resized + }); + T2IModel priorFinalModel = FinalLoadedModel; + List priorFinalModelList = FinalLoadedModelList; + WGNodeData priorModel = CurrentModel, priorTextEnc = CurrentTextEnc, priorVae = CurrentVae; + bool priorNoVae = NoVAEOverride; + FinalLoadedModel = seedVrModel; + FinalLoadedModelList = [seedVrModel]; + NoVAEOverride = true; + (FinalLoadedModel, CurrentModel, CurrentTextEnc, CurrentVae) = CreateModelLoader(seedVrModel, "SeedVR2"); + NoVAEOverride = priorNoVae; + WGNodeData encoded = raw.WithPath([preprocessed, 0]).EncodeToLatent(CurrentVae); + JArray latent = encoded.Path; + JArray chunkOverlap = null; + if (UserInput.Get(ComfyUIBackendExtension.SeedVR2SplitLatent, false)) + { + int overlap = UserInput.TryGet(ComfyUIBackendExtension.SeedVR2TemporalVideoOverlap, out int seedVrOverlap) ? seedVrOverlap : UserInput.Get(T2IParamTypes.VAETemporalTileOverlap, 0); + string chunked = CreateNode("SeedVR2TemporalChunk", new JObject() + { + ["latent"] = latent, + ["temporal_overlap"] = overlap, + ["chunking_mode"] = "auto" + }); + latent = [chunked, 0]; + chunkOverlap = [chunked, 1]; + } + string cond = CreateNode("SeedVR2Conditioning", new JObject() + { + ["model"] = CurrentModel.Path, + ["vae_conditioning"] = latent + }); + string sampled = CreateKSampler(CurrentModel.Path, [cond, 0], [cond, 1], latent, 1, 1, 0, 10000, seed, false, true, + explicitSampler: "euler", explicitScheduler: "simple"); + JArray sampledLatent = [sampled, 0]; + if (chunkOverlap is not null) + { + string merged = CreateNode("SeedVR2TemporalMerge", new JObject() + { + ["latents"] = sampledLatent, + ["temporal_overlap"] = chunkOverlap + }); + sampledLatent = [merged, 0]; + } + WGNodeData decoded = encoded.WithPath(sampledLatent).DecodeLatents(CurrentVae, false); + string post = CreateNode("SeedVR2PostProcessing", new JObject() + { + ["images"] = decoded.Path, + ["original_resized_images"] = resized, + ["color_correction_method"] = UserInput.Get(ComfyUIBackendExtension.SeedVR2ColorCorrectionBehavior, "none") + }); + FinalLoadedModel = priorFinalModel; + FinalLoadedModelList = priorFinalModelList; + CurrentModel = priorModel; + CurrentTextEnc = priorTextEnc; + CurrentVae = priorVae; + WGNodeData result = raw.WithPath([post, 0]); + result.Width = raw.Width; + result.Height = raw.Height; + return result; + } + /// Scales raw media to an exact pixel size, or does nothing if it's already that size. public WGNodeData ScaleRawMedia(WGNodeData raw, int width, int height, string method = "lanczos", string id = null) { @@ -2721,6 +2793,11 @@ public WGNodeData CreatePixelUpscale(string method, WGNodeData media, WGNodeData WGNodeData raw = media.AsRawImage(vae); int width = (int)Math.Round((raw.Width ?? UserInput.GetImageWidth()) * scale) / 16 * 16; int height = (int)Math.Round((raw.Height ?? UserInput.GetImageHeight()) * scale) / 16 * 16; + if (method.StartsWith("seedvr2model-")) + { + T2IModel seedVrModel = ComfyUIBackendExtension.GetSeedVR2Model(method.After("seedvr2model-"), UserInput.SourceSession); + return CreateSeedVR2Restore(seedVrModel, ScaleRawMedia(raw, width, height), vae, seed); + } if (method.StartsWith("pidmodel-")) { T2IModel pidModel = ComfyUIBackendExtension.GetPidModel(method.After("pidmodel-"), UserInput.SourceSession); diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs index 54be505c4..de8f92a1d 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs @@ -94,6 +94,9 @@ public bool IsKontext() /// Returns true if the current model is PiD. public bool IsPiD() => IsModelCompatClass(T2IModelClassSorter.CompatPiD); + /// Returns true if the current model is SeedVR2. + public bool IsSeedVR2() => IsModelCompatClass(T2IModelClassSorter.CompatSeedVR2); + /// Returns true if the current model is HiDream-i1. public bool IsHiDream() => IsModelCompatClass(T2IModelClassSorter.CompatHiDreamI1); @@ -815,6 +818,10 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) } IsDifferentialDiffusion = false; LoadingModelType = type; + if (type != "SeedVR2" && model.ModelClass?.CompatClass?.ID == "seedvr2") + { + throw new SwarmUserErrorException($"Model '{model.Name}' is a SeedVR2 model, which can only be used as an upscale method."); + } if (!noCascadeFix && model.ModelClass?.ID == "stable-cascade-v1-stage-b" && model.Name.Contains("stage_b") && Program.MainSDModels.Models.TryGetValue(model.Name.Replace("stage_b", "stage_c"), out T2IModel altCascadeModel)) { model = altCascadeModel; @@ -966,7 +973,7 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) { dtype = "default"; } - else if (IsZImage() || IsZetaChroma() || IsAnima() || IsLens() || IsPixelDiT() || IsPiD()) // Model is small and dense, so trust user preferred download format + else if (IsZImage() || IsZetaChroma() || IsAnima() || IsLens() || IsPixelDiT() || IsPiD() || IsSeedVR2()) // Model is small and dense, so trust user preferred download format { dtype = "default"; } @@ -1233,6 +1240,10 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) helpers.LoadClip("pixeldit", helpers.GetGemma2_2bElmModel()); LoadingVAE = CreateVAELoader("pixel_space"); } + else if (IsSeedVR2()) + { + helpers.DoVaeLoader(null, T2IModelClassSorter.CompatSeedVR2, "seedvr2-vae"); + } else if (IsHiDream()) { string loaderType = "QuadrupleCLIPLoader"; @@ -1484,7 +1495,7 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) { step.Action(this); } - if (LoadingClip is null) + if (LoadingClip is null && !IsSeedVR2()) { if (string.IsNullOrWhiteSpace(model.Metadata?.ModelClassType)) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs index 55a1c756e..79fb2fb5a 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs @@ -65,7 +65,7 @@ public static void Register() }, -15); AddModelGenStep(g => { - if (g.IsFinalStage) + if (g.IsFinalStage || g.IsSeedVR2()) { return; } @@ -106,7 +106,7 @@ public static void Register() }, -14); AddModelGenStep(g => { - if (g.LoadingModelType == "negative" && !g.UserInput.Get(T2IParamTypes.NegativeModelIncludeLoras, true) || g.IsFinalStage) + if (g.LoadingModelType == "negative" && !g.UserInput.Get(T2IParamTypes.NegativeModelIncludeLoras, true) || g.IsFinalStage || g.IsSeedVR2()) { return; } @@ -135,7 +135,7 @@ public static void Register() }, -10); AddModelGenStep(g => { - if (g.UserInput.TryGet(ComfyUIBackendExtension.SetClipDevice, out string device) && g.Features.Contains("set_clip_device")) + if (g.LoadingClip is not null && g.UserInput.TryGet(ComfyUIBackendExtension.SetClipDevice, out string device) && g.Features.Contains("set_clip_device")) { string clipDeviceNode = g.CreateNode("OverrideCLIPDevice", new JObject() { @@ -147,7 +147,7 @@ public static void Register() }, -9); AddModelGenStep(g => { - if (g.IsFinalStage) + if (g.IsFinalStage || g.IsSeedVR2()) { return; } @@ -171,7 +171,7 @@ public static void Register() }, -8); AddModelGenStep(g => { - if (g.IsFinalStage) + if (g.IsFinalStage || g.IsSeedVR2()) { return; } @@ -243,7 +243,7 @@ public static void Register() }, -7); AddModelGenStep(g => { - if (g.UserInput.TryGet(T2IParamTypes.ClipStopAtLayer, out int layer)) + if (g.LoadingClip is not null && g.UserInput.TryGet(T2IParamTypes.ClipStopAtLayer, out int layer)) { string clipSkip = g.CreateNode("CLIPSetLastLayer", new JObject() { @@ -255,7 +255,7 @@ public static void Register() }, -6); AddModelGenStep(g => { - if (g.IsFinalStage) + if (g.IsFinalStage || g.IsSeedVR2()) { return; } @@ -280,7 +280,7 @@ public static void Register() }, -5); AddModelGenStep(g => { - if (g.IsFinalStage) + if (g.IsFinalStage || g.IsSeedVR2()) { return; } @@ -390,10 +390,6 @@ public static void Register() }, -4); AddModelGenStep(g => { - if (g.IsFinalStage) - { - return; - } if (g.Features.Contains("aitemplate") && g.UserInput.Get(ComfyUIBackendExtension.AITemplateParam)) { string aitLoad = g.CreateNode("AITemplateLoader", new JObject() @@ -1562,6 +1558,7 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) // TODO: Better same-VAE check bool doPixelUpscale = doUpscale && (upscaleMethod.StartsWith("pixel-") || upscaleMethod.StartsWith("model-")); bool doPidUpscale = doUpscale && upscaleMethod.StartsWith("pidmodel-"); + bool doSeedVR2Upscale = doUpscale && upscaleMethod.StartsWith("seedvr2model-"); int width = (int)Math.Round(g.UserInput.GetImageWidth() * refineUpscale); int height = (int)Math.Round(g.UserInput.GetImageHeight() * refineUpscale); width = (width / 16) * 16; // avoid unworkable output sizes @@ -1595,6 +1592,25 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) } g.CurrentMedia = decoded.EncodeToLatent(g.CurrentVae, "25"); } + else if (doSeedVR2Upscale) + { + T2IModel seedVrModel = ComfyUIBackendExtension.GetSeedVR2Model(upscaleMethod.After("seedvr2model-"), g.UserInput.SourceSession); + WGNodeData decoded = g.CurrentMedia.DecodeLatents(origVae, false, "24"); + decoded = decoded.WithPath(doMaskShrinkApply(g, decoded.Path)); + if (doSave) + { + decoded.SaveOutput(null, null, id: "29"); + } + decoded = g.ScaleRawMedia(decoded, width, height, id: "26"); + decoded = g.CreateSeedVR2Restore(seedVrModel, decoded, origVae, g.UserInput.Get(T2IParamTypes.Seed) + 2); + if (refinerControl <= 0) + { + g.CurrentMedia = decoded; + g.IsRefinerStage = false; + return; + } + g.CurrentMedia = decoded.EncodeToLatent(g.CurrentVae, "25"); + } else if (modelMustReencode || doPixelUpscale || doSave || g.MaskShrunkInfo.BoundsNode is not null) { WGNodeData decoded = g.CurrentMedia.DecodeLatents(origVae, false, "24"); diff --git a/src/Text2Image/CommonModels.cs b/src/Text2Image/CommonModels.cs index 1c5c8ce79..09956a379 100644 --- a/src/Text2Image/CommonModels.cs +++ b/src/Text2Image/CommonModels.cs @@ -95,6 +95,7 @@ public static void RegisterCoreSet() Register(new("mage-flow-vae", "Mage-VAE", "The VAE for Mage-Flow", "https://huggingface.co/microsoft/Mage-Flow/resolve/main/vae/diffusion_pytorch_model.safetensors", "34e076dc1e8a15321e1e07be5111d59cf16dd10b804b7c7e20b4de29013427e0", "VAE", "MageFlow/mage_flow_vae.safetensors")); Register(new("hunyuan-image-2_1-vae", "Hunyuan Image 2.1 VAE", "The VAE for Hunyuan Image 2.1 Base", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_2.1_vae_fp16.safetensors", "f2ae19863609206196b5e3a86bfd94f67bd3866f5042004e3994f07e3c93b2f9", "VAE", "HunyuanImage/hunyuan_image_2.1_vae_fp16.safetensors")); Register(new("hunyuan-image-2_1-refiner-vae", "Hunyuan Image 2.1 Refiner VAE", "The VAE for Hunyuan Image 2.1 Refiner", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_refiner_vae_fp16.safetensors", "e1b74e85d61b65e18cc05ca390e387d93cfadf161e737de229ebb800ea3db769", "VAE", "HunyuanImage/hunyuan_image_2.1_refiner_vae_fp16.safetensors")); + Register(new("seedvr2-vae", "SeedVR2 VAE", "The VAE for SeedVR2", "https://huggingface.co/Comfy-Org/SeedVR2/resolve/main/vae/seedvr2_ema_vae_fp16.safetensors", "20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1", "VAE", "SeedVR2/seedvr2_ema_vae_fp16.safetensors")); Register(new("hunyuan-video-1_5-vae", "Hunyuan Video 1.5 VAE", "The VAE for Hunyuan Video 1.5", "https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/resolve/main/split_files/vae/hunyuanvideo15_vae_fp16.safetensors", "e7c3091949c27e2d55ae6d5df917b99dadfebbf308e5a50d0ade0d16c90297ae", "VAE", "HunyuanVideo/hunyuanvideo15_vae_fp16.safetensors")); // Audio VAEs diff --git a/src/Text2Image/T2IModelClassSorter.cs b/src/Text2Image/T2IModelClassSorter.cs index 18d841b8a..8face175f 100644 --- a/src/Text2Image/T2IModelClassSorter.cs +++ b/src/Text2Image/T2IModelClassSorter.cs @@ -92,6 +92,7 @@ public static T2IModelCompatClass CompatLens = RegisterCompat(new() { ID = "lens", ShortCode = "Lens", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatPiD = RegisterCompat(new() { ID = "pid", ShortCode = "PiD", LorasTargetTextEnc = false }), CompatPixelDiT = RegisterCompat(new() { ID = "pixeldit", ShortCode = "PixDiT", LorasTargetTextEnc = false }), + CompatSeedVR2 = RegisterCompat(new() { ID = "seedvr2", ShortCode = "SeedVR2", LorasTargetTextEnc = false }), CompatIdeogram4 = RegisterCompat(new() { ID = "ideogram-4", ShortCode = "Ideo4", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatKrea2 = RegisterCompat(new() { ID = "krea-2", ShortCode = "Krea2", VaeFamily = VaeQwenImage }), CompatBoogu = RegisterCompat(new() { ID = "boogu", ShortCode = "Boogu", LorasTargetTextEnc = false, VaeFamily = VaeFlux1 }), @@ -241,6 +242,8 @@ bool isZImageLora(JObject h) => (hasLoraKey(h, "layers.0.adaLN_modulation.0") && bool isChromaRadiance(JObject h) => hasKey(h, "nerf_image_embedder.embedder.0.bias"); bool isPiD(JObject h) => h.ContainsKey("net.lq_proj.latent_proj.0.weight") && h.ContainsKey("net.pixel_blocks.0.attn.q_norm.weight") && h.ContainsKey("net.pixel_blocks.0.compress_to_attn.weight"); bool isPixelDiT(JObject h) => h.ContainsKey("core.pixel_embedder.proj.weight") && h.ContainsKey("core.pixel_blocks.0.attn.q_norm.weight") && h.ContainsKey("core.pixel_blocks.0.compress_to_attn.weight") && !isPiD(h); + bool isSeedVR2(JObject h) => hasKey(h, "blocks.31.mlp.all.proj_in_gate.weight") || hasKey(h, "blocks.35.mlp.all.proj_in_gate.weight") || hasKey(h, "blocks.35.mlp.vid.proj_out.weight"); + bool isSeedVR2Vae(JObject h) => h.ContainsKey("decoder.up_blocks.2.upsamplers.0.upscale_conv.weight"); bool isOmniGen(JObject h) => h.ContainsKey("time_caption_embed.timestep_embedder.linear_2.weight") && h.ContainsKey("context_refiner.0.attn.norm_k.weight") && !isBoogu(h); bool isBoogu(JObject h) => hasKey(h, "double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") && hasKey(h, "double_stream_layers.0.img_instruct_attn.processor.instruct_to_q.weight"); bool isQwenImage(JObject h) => (h.ContainsKey("time_text_embed.timestep_embedder.linear_1.bias") && h.ContainsKey("img_in.bias") && (h.ContainsKey("transformer_blocks.0.attn.add_k_proj.bias") || h.ContainsKey("transformer_blocks.0.attn.add_qkv_proj.bias"))) @@ -776,6 +779,15 @@ JToken GetEmbeddingKey(JObject h) { return isPixelDiT(h); }}); + // ====================== SeedVR2 ====================== + Register(new() { ID = "seedvr2", CompatClass = CompatSeedVR2, Name = "SeedVR2", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => + { + return isSeedVR2(h); + }}); + Register(new() { ID = "seedvr2/vae", CompatClass = CompatSeedVR2, Name = "SeedVR2 VAE", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => + { + return isSeedVR2Vae(h); + }}); Register(new() { ID = "alt_diffusion_v1_512_placeholder", CompatClass = CompatAltDiffusion, Name = "Alt-Diffusion", StandardWidth = 512, StandardHeight = 512, IsThisModelOfClass = (m, h) => { return IsAlt(h); From 3b1aec454a5250b516f761f31f88dd1dac3056bf Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Fri, 31 Jul 2026 17:24:08 -0500 Subject: [PATCH 4/7] Revert "Adds support for Final Group", "Adds support for *2v workflows in Final Stage", "Adds native SeedVR2 support" --- docs/Features/Upscaling.md | 17 -- .../ComfyUIBackend/ComfyUIBackendExtension.cs | 47 +----- .../ComfyUIBackend/WorkflowGenerator.cs | 156 ------------------ .../WorkflowGeneratorModelSupport.cs | 15 +- .../ComfyUIBackend/WorkflowGeneratorSteps.cs | 55 +----- src/Text2Image/CommonModels.cs | 1 - src/Text2Image/T2IModelClassSorter.cs | 12 -- src/Text2Image/T2IParamTypes.cs | 9 +- 8 files changed, 14 insertions(+), 298 deletions(-) diff --git a/docs/Features/Upscaling.md b/docs/Features/Upscaling.md index b76fdb0a4..0da35122e 100644 --- a/docs/Features/Upscaling.md +++ b/docs/Features/Upscaling.md @@ -2,25 +2,8 @@ (TODO) -# Upscale Stages - -There are two places to upscale, differing in what happens after: - -- **Refine / Upscale** group: the refiner model samples over the upscaled result, per **Refiner Control Percentage**. Set that to `0` for an upscale-only stage. -- **Final Stage** group: runs after the base and refiner stages are done, and nothing samples over the result, only whatever the upscaler does itself. - - **Final Upscale** stacks on top of Refiner Upscale, eg `1.5` refiner upscale and `2` final upscale is 3x total. - - **Final Upscale Method** offers the same methods as Refiner Upscale Method, minus the latent upscalers (which need a sampler after them). - # Pixel Decoder (PiD) (TODO) Downloads here: - -# SeedVR2 - -- [SeedVR2]() is a one-step restoration model for upscaling images or videos. - - Models can be downloaded here: [Comfy-Org/SeedVR2]() - - Save in `diffusion_models` - - The VAE will be automatically downloaded -- Select it as your `Refiner Upscale Method` or your `Final Upscale Method`, see [Upscale Stages](#upscale-stages) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs index c8d037da8..999ebc564 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs @@ -35,10 +35,10 @@ public record class ComfyCustomWorkflow(string Name, string Workflow, string Pro public static ConcurrentDictionary CustomWorkflows = new(); /// Set of all feature-ids supported by ComfyUI backends. - public static HashSet FeaturesSupported = ["comfyui", "refiners", "controlnet", "endstepsearly", "seamless", "video", "variation_seed", "freeu", "yolov8", "seedvr2"]; + public static HashSet FeaturesSupported = ["comfyui", "refiners", "controlnet", "endstepsearly", "seamless", "video", "variation_seed", "freeu", "yolov8"]; /// Set of feature-ids that were added presumptively during loading and should be removed if the backend turns out to be missing them. - public static HashSet FeaturesDiscardIfNotFound = ["variation_seed", "freeu", "yolov8", "seedvr2"]; + public static HashSet FeaturesDiscardIfNotFound = ["variation_seed", "freeu", "yolov8"]; /// Extensible map of ComfyUI Node IDs to supported feature IDs. public static Dictionary NodeToFeatureMap = new() @@ -68,8 +68,7 @@ public record class ComfyCustomWorkflow(string Name, string Workflow, string Pro ["TeaCache"] = "teacache", ["TeaCacheForVidGen"] = "teacache", ["TeaCacheForImgGen"] = "teacache_oldvers", - ["OverrideCLIPDevice"] = "set_clip_device", - ["SeedVR2Conditioning"] = "seedvr2", + ["OverrideCLIPDevice"] = "set_clip_device" }; /// @@ -607,13 +606,13 @@ public static void AssignValuesFromRaw(JObject rawObjectInfo) } } - public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, FinalUpscaleMethod, SeedVR2ColorCorrectionBehavior, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; + public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; - public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG, SeedVR2SplitLatent; + public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG; public static T2IRegisteredParam IPAdapterWeight, IPAdapterStart, IPAdapterEnd, SelfAttentionGuidanceScale, SelfAttentionGuidanceSigmaBlur, PerturbedAttentionGuidanceScale, StyleModelMergeStrength, StyleModelApplyStart, StyleModelMultiplyStrength, RescaleCFGMultiplier, TeaCacheThreshold, TeaCacheStart, NunchakuCacheThreshold, EasyCacheThreshold, EasyCacheStart, EasyCacheEnd, RenormCFG, NormalizedAttentionGuidanceScale, NormalizedAttentionGuidanceAlpha, NormalizedAttentionGuidanceTau; - public static T2IRegisteredParam RefinerHyperTile, VideoFrameInterpolationMultiplier, SeedVR2TemporalVideoOverlap; + public static T2IRegisteredParam RefinerHyperTile, VideoFrameInterpolationMultiplier; public static T2IRegisteredParam PixelDecoderModel; @@ -657,25 +656,6 @@ public static T2IModel GetPidModel(string name, Session session) return model; } - /// Lists SeedVR2 upscaler models. - public static List SeedVR2UpscaleModels(Session session) => [.. Program.MainSDModels.ListModelsFor(session).Where(m => m.ModelClass?.CompatClass?.ID == "seedvr2").OrderBy(m => m.Name).Select(m => $"seedvr2model-{m.Name}///SeedVR2 Model: {m.Name}")]; - - /// Resolves a SeedVR2 model from a model name. - public static T2IModel GetSeedVR2Model(string name, Session session) - { - string matched = T2IParamTypes.GetBestModelInList(name, Program.MainSDModels.ListModelNamesFor(session)); - if (matched is not null && matched.EndsWith(".safetensors")) - { - matched = matched.BeforeLast('.'); - } - T2IModel model = matched is null ? null : Program.MainSDModels.GetModel(matched); - if (model is null || model.ModelClass?.CompatClass?.ID != "seedvr2") - { - throw new SwarmUserErrorException($"SeedVR2 model '{name}' could not be found, or is not a valid SeedVR2 model."); - } - return model; - } - public static List IPAdapterModels = ["None"], IPAdapterWeightTypes = ["standard", "prompt is more important", "style transfer"]; public static List GligenModels = ["None"], YoloModels = [], StyleModels = ["None"], SetClipDevices = ["cpu"]; @@ -790,11 +770,7 @@ public override void OnInit() )); RefinerUpscaleMethod = T2IParamTypes.Register(new("Refiner Upscale Method", "How to upscale the image, if upscaling is used.", "pixel-lanczos", Group: T2IParamTypes.GroupRefiners, OrderPriority: -1, FeatureFlag: "comfyui", ChangeWeight: 1, - GetValues: (session) => [.. UpscalerModels, .. PidUpscaleModels(session), .. SeedVR2UpscaleModels(session)], DependNonDefault: T2IParamTypes.RefinerUpscale.Type.ID - )); - FinalUpscaleMethod = T2IParamTypes.Register(new("Final Upscale Method", "How to upscale the image, if upscaling is used.\nNo sampler runs after this, so latent upscalers are not available here.", - "pixel-lanczos", Group: T2IParamTypes.GroupFinalStage, OrderPriority: -1, FeatureFlag: "comfyui", ChangeWeight: 1, - GetValues: (session) => [.. RefinerUpscaleMethod.Type.GetValues(session).Where(u => !u.StartsWith("latent"))] + GetValues: (session) => [.. UpscalerModels, .. PidUpscaleModels(session)], DependNonDefault: T2IParamTypes.RefinerUpscale.Type.ID )); PixelDecoderModel = T2IParamTypes.Register(new("Pixel Decoder Model", "Optionally use a PiD (Pixel Diffusion Decoder) model.", "", Toggleable: true, FeatureFlag: "comfyui", Group: T2IParamTypes.GroupAdvancedModelAddons, IsAdvanced: true, Subtype: "Stable-Diffusion", ChangeWeight: 4, DoNotPreview: true, OrderPriority: 14, @@ -823,15 +799,6 @@ public override void OnInit() RefinerHyperTile = T2IParamTypes.Register(new("Refiner HyperTile", "The size of hypertiles to use for the refining stage.\nHyperTile is a technique to speed up sampling of large images by tiling the image and batching the tiles.\nThis is useful when using SDv1 models as the refiner. SDXL-Base models do not benefit as much.", "256", Min: 64, Max: 2048, Step: 32, Toggleable: true, IsAdvanced: true, FeatureFlag: "comfyui", ViewType: ParamViewType.POT_SLIDER, Group: T2IParamTypes.GroupAdvancedSampling, OrderPriority: 20 )); - SeedVR2ColorCorrectionBehavior = T2IParamTypes.Register(new("SeedVR2 Color Correction Behavior", "How to match the colors of a SeedVR2 upscale back to the image it was given.\n'None' = Do not attempt color correction, only align the geometry.\n'CIELAB' = Transfer the color in CIELAB space, preserving detail.\n'Wavelet' = Transfer the low-frequency color, keeping the upscaled high-frequency detail.\n'AdaIN' = Match the per-channel mean and standard deviation.", - "none", FeatureFlag: "seedvr2", Group: T2IParamTypes.GroupAdvancedSampling, IsAdvanced: true, OrderPriority: 21, GetValues: (_) => ["none///None", "lab///CIELAB", "wavelet///Wavelet", "adain///AdaIN"] - )); - SeedVR2SplitLatent = T2IParamTypes.Register(new("SeedVR2 Split Latent", "If enabled, samples a SeedVR2 video upscale as chunks of frames instead of all at once, sized to fit in free VRAM.\nChunking reduces VRAM consumption.\nDoes nothing to a single image, or to a video that already fits.", - "false", IgnoreIf: "false", FeatureFlag: "seedvr2", Group: T2IParamTypes.GroupAdvancedSampling, IsAdvanced: true, OrderPriority: 22 - )); - SeedVR2TemporalVideoOverlap = T2IParamTypes.Register(new("SeedVR2 Temporal Video Overlap", "Overrides 'VAE Temporal Tile Overlap' for 'SeedVR2 Split Latent' chunks.\nHigher overlap hides the chunk seams better but takes longer.", - "0", Min: 0, Max: 4096, Step: 1, Toggleable: true, VisibleNormally: false, IsAdvanced: true, FeatureFlag: "seedvr2", Group: T2IParamTypes.GroupAdvancedSampling, OrderPriority: 23 - )); List interpolators = ["RIFE", "FILM", "GIMM-VFI"]; VideoPreviewType = T2IParamTypes.Register(new("Video Preview Type", "How to display previews for generating videos.\n'Animate' shows a low-res animated video preview.\n'iterate' shows one frame at a time while it goes.\n'one' displays just the first frame.\n'none' disables previews.", "animate", IgnoreIf: "animate", FeatureFlag: "comfyui", Group: T2IParamTypes.GroupAdvancedVideo, Permission: Permissions.ParamVideo, IsAdvanced: true, GetValues: (_) => ["animate", "iterate", "one", "none"] diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index c9095e6e8..286f46154 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -169,9 +169,6 @@ public JArray FinalImageOut /// If true, the generator is currently working on the pixel-decoder stage. public bool IsPixelDecoderStage = false; - /// If true, the generator is currently working on the final stage upscale. - public bool IsFinalStage = false; - /// If true, the generator is currently working on Image2Video. public bool IsImageToVideo = false; @@ -2673,159 +2670,6 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD return result; } - /// Creates a SeedVR2 restoration stage. - public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, WGNodeData decodeVae, long seed) - { - if (!Features.Contains("seedvr2")) - { - throw new SwarmUserErrorException($"Cannot use SeedVR2 model '{seedVrModel.Name}', the backend is missing SeedVR2 support. Update ComfyUI."); - } - WGNodeData raw = media.AsRawImage(decodeVae); - JArray resized = raw.Path; - string preprocessed = CreateNode("SeedVR2Preprocess", new JObject() - { - ["resized_images"] = resized - }); - T2IModel priorFinalModel = FinalLoadedModel; - List priorFinalModelList = FinalLoadedModelList; - WGNodeData priorModel = CurrentModel, priorTextEnc = CurrentTextEnc, priorVae = CurrentVae; - bool priorNoVae = NoVAEOverride; - FinalLoadedModel = seedVrModel; - FinalLoadedModelList = [seedVrModel]; - NoVAEOverride = true; - (FinalLoadedModel, CurrentModel, CurrentTextEnc, CurrentVae) = CreateModelLoader(seedVrModel, "SeedVR2"); - NoVAEOverride = priorNoVae; - WGNodeData encoded = raw.WithPath([preprocessed, 0]).EncodeToLatent(CurrentVae); - JArray latent = encoded.Path; - JArray chunkOverlap = null; - if (UserInput.Get(ComfyUIBackendExtension.SeedVR2SplitLatent, false)) - { - int overlap = UserInput.TryGet(ComfyUIBackendExtension.SeedVR2TemporalVideoOverlap, out int seedVrOverlap) ? seedVrOverlap : UserInput.Get(T2IParamTypes.VAETemporalTileOverlap, 0); - string chunked = CreateNode("SeedVR2TemporalChunk", new JObject() - { - ["latent"] = latent, - ["temporal_overlap"] = overlap, - ["chunking_mode"] = "auto" - }); - latent = [chunked, 0]; - chunkOverlap = [chunked, 1]; - } - string cond = CreateNode("SeedVR2Conditioning", new JObject() - { - ["model"] = CurrentModel.Path, - ["vae_conditioning"] = latent - }); - string sampled = CreateKSampler(CurrentModel.Path, [cond, 0], [cond, 1], latent, 1, 1, 0, 10000, seed, false, true, - explicitSampler: "euler", explicitScheduler: "simple"); - JArray sampledLatent = [sampled, 0]; - if (chunkOverlap is not null) - { - string merged = CreateNode("SeedVR2TemporalMerge", new JObject() - { - ["latents"] = sampledLatent, - ["temporal_overlap"] = chunkOverlap - }); - sampledLatent = [merged, 0]; - } - WGNodeData decoded = encoded.WithPath(sampledLatent).DecodeLatents(CurrentVae, false); - string post = CreateNode("SeedVR2PostProcessing", new JObject() - { - ["images"] = decoded.Path, - ["original_resized_images"] = resized, - ["color_correction_method"] = UserInput.Get(ComfyUIBackendExtension.SeedVR2ColorCorrectionBehavior, "none") - }); - FinalLoadedModel = priorFinalModel; - FinalLoadedModelList = priorFinalModelList; - CurrentModel = priorModel; - CurrentTextEnc = priorTextEnc; - CurrentVae = priorVae; - WGNodeData result = raw.WithPath([post, 0]); - result.Width = raw.Width; - result.Height = raw.Height; - return result; - } - - /// Scales raw media to an exact pixel size, or does nothing if it's already that size. - public WGNodeData ScaleRawMedia(WGNodeData raw, int width, int height, string method = "lanczos", string id = null) - { - if (raw.Width == width && raw.Height == height) - { - return raw; - } - string scaled = CreateNode("ImageScale", new JObject() - { - ["image"] = raw.Path, - ["width"] = width, - ["height"] = height, - ["upscale_method"] = method, - ["crop"] = "disabled" - }, id); - WGNodeData result = raw.WithPath([scaled, 0]); - result.Width = width; - result.Height = height; - return result; - } - - /// Runs the Final Stage upscale over the current media, if the user configured one and it hasn't run yet. - public void RunFinalStage() - { - if (!UserInput.TryGet(ComfyUIBackendExtension.FinalUpscaleMethod, out string method)) - { - return; - } - double scale = UserInput.Get(T2IParamTypes.FinalUpscale, 2); - if (scale == 1 && method.StartsWith("pixel-")) - { - return; - } - if (UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) - { - CurrentMedia.SaveOutput(CurrentVae, CurrentAudioVae, GetStableDynamicID(50000, 0)); - } - IsFinalStage = true; - CurrentMedia = CreatePixelUpscale(method, CurrentMedia, CurrentVae, scale, UserInput.Get(T2IParamTypes.Seed) + 500); - IsFinalStage = false; - } - - /// Upscales raw media by any of the Final Upscale Methods. - public WGNodeData CreatePixelUpscale(string method, WGNodeData media, WGNodeData vae, double scale, long seed) - { - WGNodeData raw = media.AsRawImage(vae); - int width = (int)Math.Round((raw.Width ?? UserInput.GetImageWidth()) * scale) / 16 * 16; - int height = (int)Math.Round((raw.Height ?? UserInput.GetImageHeight()) * scale) / 16 * 16; - if (method.StartsWith("seedvr2model-")) - { - T2IModel seedVrModel = ComfyUIBackendExtension.GetSeedVR2Model(method.After("seedvr2model-"), UserInput.SourceSession); - return CreateSeedVR2Restore(seedVrModel, ScaleRawMedia(raw, width, height), vae, seed); - } - if (method.StartsWith("pidmodel-")) - { - T2IModel pidModel = ComfyUIBackendExtension.GetPidModel(method.After("pidmodel-"), UserInput.SourceSession); - return ScaleRawMedia(CreatePixelDecode(pidModel, raw, vae, seed), width, height); - } - if (method.StartsWith("model-")) - { - string loaderNode = CreateNode("UpscaleModelLoader", new JObject() - { - ["model_name"] = method.After("model-") - }); - string upscaledNode = CreateNode("ImageUpscaleWithModel", new JObject() - { - ["upscale_model"] = NodePath(loaderNode, 0), - ["image"] = raw.Path - }); - WGNodeData upscaled = raw.WithPath([upscaledNode, 0]); - upscaled.Width = null; // the model's own scale factor is unknown here, so always correct after - upscaled.Height = null; - return ScaleRawMedia(upscaled, width, height); - } - if (method.StartsWith("pixel-")) - { - return ScaleRawMedia(raw, width, height, method.After("pixel-")); - } - throw new SwarmUserErrorException($"Upscale method '{method}' needs a sampler after it, so it can't be used as a Final Upscale Method. Use it in the Refine/Upscale group instead."); - } - /// Creates a "CLIPTextEncode" or equivalent node for the given input, applying prompt-given conditioning modifiers as relevant. public JArray CreateConditioning(string prompt, JArray clip, T2IModel model, bool isPositive, string firstId = null, bool isRefiner = false, bool isVideo = false, bool isVideoSwap = false, bool isPixelDecoder = false) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs index de8f92a1d..54be505c4 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs @@ -94,9 +94,6 @@ public bool IsKontext() /// Returns true if the current model is PiD. public bool IsPiD() => IsModelCompatClass(T2IModelClassSorter.CompatPiD); - /// Returns true if the current model is SeedVR2. - public bool IsSeedVR2() => IsModelCompatClass(T2IModelClassSorter.CompatSeedVR2); - /// Returns true if the current model is HiDream-i1. public bool IsHiDream() => IsModelCompatClass(T2IModelClassSorter.CompatHiDreamI1); @@ -818,10 +815,6 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) } IsDifferentialDiffusion = false; LoadingModelType = type; - if (type != "SeedVR2" && model.ModelClass?.CompatClass?.ID == "seedvr2") - { - throw new SwarmUserErrorException($"Model '{model.Name}' is a SeedVR2 model, which can only be used as an upscale method."); - } if (!noCascadeFix && model.ModelClass?.ID == "stable-cascade-v1-stage-b" && model.Name.Contains("stage_b") && Program.MainSDModels.Models.TryGetValue(model.Name.Replace("stage_b", "stage_c"), out T2IModel altCascadeModel)) { model = altCascadeModel; @@ -973,7 +966,7 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) { dtype = "default"; } - else if (IsZImage() || IsZetaChroma() || IsAnima() || IsLens() || IsPixelDiT() || IsPiD() || IsSeedVR2()) // Model is small and dense, so trust user preferred download format + else if (IsZImage() || IsZetaChroma() || IsAnima() || IsLens() || IsPixelDiT() || IsPiD()) // Model is small and dense, so trust user preferred download format { dtype = "default"; } @@ -1240,10 +1233,6 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) helpers.LoadClip("pixeldit", helpers.GetGemma2_2bElmModel()); LoadingVAE = CreateVAELoader("pixel_space"); } - else if (IsSeedVR2()) - { - helpers.DoVaeLoader(null, T2IModelClassSorter.CompatSeedVR2, "seedvr2-vae"); - } else if (IsHiDream()) { string loaderType = "QuadrupleCLIPLoader"; @@ -1495,7 +1484,7 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) { step.Action(this); } - if (LoadingClip is null && !IsSeedVR2()) + if (LoadingClip is null) { if (string.IsNullOrWhiteSpace(model.Metadata?.ModelClassType)) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs index 79fb2fb5a..c2ed9900e 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs @@ -65,10 +65,6 @@ public static void Register() }, -15); AddModelGenStep(g => { - if (g.IsFinalStage || g.IsSeedVR2()) - { - return; - } if (g.IsRefinerStage && g.UserInput.TryGet(T2IParamTypes.RefinerVAE, out T2IModel rvae)) { g.LoadingVAE = g.CreateVAELoader(rvae.ToString(g.ModelFolderFormat), g.HasNode("21") ? null : "21"); @@ -106,7 +102,7 @@ public static void Register() }, -14); AddModelGenStep(g => { - if (g.LoadingModelType == "negative" && !g.UserInput.Get(T2IParamTypes.NegativeModelIncludeLoras, true) || g.IsFinalStage || g.IsSeedVR2()) + if (g.LoadingModelType == "negative" && !g.UserInput.Get(T2IParamTypes.NegativeModelIncludeLoras, true)) { return; } @@ -135,7 +131,7 @@ public static void Register() }, -10); AddModelGenStep(g => { - if (g.LoadingClip is not null && g.UserInput.TryGet(ComfyUIBackendExtension.SetClipDevice, out string device) && g.Features.Contains("set_clip_device")) + if (g.UserInput.TryGet(ComfyUIBackendExtension.SetClipDevice, out string device) && g.Features.Contains("set_clip_device")) { string clipDeviceNode = g.CreateNode("OverrideCLIPDevice", new JObject() { @@ -147,10 +143,6 @@ public static void Register() }, -9); AddModelGenStep(g => { - if (g.IsFinalStage || g.IsSeedVR2()) - { - return; - } string applyTo = g.UserInput.Get(T2IParamTypes.FreeUApplyTo, null); if (g.Features.Contains("freeu") && applyTo is not null) { @@ -171,10 +163,6 @@ public static void Register() }, -8); AddModelGenStep(g => { - if (g.IsFinalStage || g.IsSeedVR2()) - { - return; - } if (g.UserInput.TryGet(ComfyUIBackendExtension.SelfAttentionGuidanceScale, out double sagScale)) { string patched = g.CreateNode("SelfAttentionGuidance", new JObject() @@ -243,7 +231,7 @@ public static void Register() }, -7); AddModelGenStep(g => { - if (g.LoadingClip is not null && g.UserInput.TryGet(T2IParamTypes.ClipStopAtLayer, out int layer)) + if (g.UserInput.TryGet(T2IParamTypes.ClipStopAtLayer, out int layer)) { string clipSkip = g.CreateNode("CLIPSetLastLayer", new JObject() { @@ -255,10 +243,6 @@ public static void Register() }, -6); AddModelGenStep(g => { - if (g.IsFinalStage || g.IsSeedVR2()) - { - return; - } if (g.UserInput.TryGet(T2IParamTypes.SeamlessTileable, out string tileable) && tileable != "false") { string mode = "Both"; @@ -280,10 +264,6 @@ public static void Register() }, -5); AddModelGenStep(g => { - if (g.IsFinalStage || g.IsSeedVR2()) - { - return; - } if (g.UserInput.TryGet(ComfyUIBackendExtension.TeaCacheMode, out string teaCacheMode) && teaCacheMode != "disabled") { double teaCacheThreshold = g.UserInput.Get(ComfyUIBackendExtension.TeaCacheThreshold, 0.25); @@ -1558,7 +1538,6 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) // TODO: Better same-VAE check bool doPixelUpscale = doUpscale && (upscaleMethod.StartsWith("pixel-") || upscaleMethod.StartsWith("model-")); bool doPidUpscale = doUpscale && upscaleMethod.StartsWith("pidmodel-"); - bool doSeedVR2Upscale = doUpscale && upscaleMethod.StartsWith("seedvr2model-"); int width = (int)Math.Round(g.UserInput.GetImageWidth() * refineUpscale); int height = (int)Math.Round(g.UserInput.GetImageHeight() * refineUpscale); width = (width / 16) * 16; // avoid unworkable output sizes @@ -1592,25 +1571,6 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) } g.CurrentMedia = decoded.EncodeToLatent(g.CurrentVae, "25"); } - else if (doSeedVR2Upscale) - { - T2IModel seedVrModel = ComfyUIBackendExtension.GetSeedVR2Model(upscaleMethod.After("seedvr2model-"), g.UserInput.SourceSession); - WGNodeData decoded = g.CurrentMedia.DecodeLatents(origVae, false, "24"); - decoded = decoded.WithPath(doMaskShrinkApply(g, decoded.Path)); - if (doSave) - { - decoded.SaveOutput(null, null, id: "29"); - } - decoded = g.ScaleRawMedia(decoded, width, height, id: "26"); - decoded = g.CreateSeedVR2Restore(seedVrModel, decoded, origVae, g.UserInput.Get(T2IParamTypes.Seed) + 2); - if (refinerControl <= 0) - { - g.CurrentMedia = decoded; - g.IsRefinerStage = false; - return; - } - g.CurrentMedia = decoded.EncodeToLatent(g.CurrentVae, "25"); - } else if (modelMustReencode || doPixelUpscale || doSave || g.MaskShrunkInfo.BoundsNode is not null) { WGNodeData decoded = g.CurrentMedia.DecodeLatents(origVae, false, "24"); @@ -2017,10 +1977,6 @@ void RunSegmentationProcessing(WorkflowGenerator g, bool isBeforeRefiner) bool willHaveFollowupVideo = g.UserInput.TryGet(T2IParamTypes.VideoModel, out _) || g.UserInput.Get(T2IParamTypes.Prompt, "").Contains(" 1) { if (g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) @@ -2295,7 +2247,6 @@ WGNodeData appendAudio(WGNodeData combinedAudio, WGNodeData nextAudio) } g.CurrentMedia = conjoinedLast; g.CurrentMedia.FPS = videoFps ?? g.CurrentMedia.FPS; - g.RunFinalStage(); if (g.UserInput.TryGet(ComfyUIBackendExtension.VideoFrameInterpolationMethod, out string method) && g.UserInput.TryGet(ComfyUIBackendExtension.VideoFrameInterpolationMultiplier, out int mult) && mult > 1) { if (saveIntermediate) diff --git a/src/Text2Image/CommonModels.cs b/src/Text2Image/CommonModels.cs index 09956a379..1c5c8ce79 100644 --- a/src/Text2Image/CommonModels.cs +++ b/src/Text2Image/CommonModels.cs @@ -95,7 +95,6 @@ public static void RegisterCoreSet() Register(new("mage-flow-vae", "Mage-VAE", "The VAE for Mage-Flow", "https://huggingface.co/microsoft/Mage-Flow/resolve/main/vae/diffusion_pytorch_model.safetensors", "34e076dc1e8a15321e1e07be5111d59cf16dd10b804b7c7e20b4de29013427e0", "VAE", "MageFlow/mage_flow_vae.safetensors")); Register(new("hunyuan-image-2_1-vae", "Hunyuan Image 2.1 VAE", "The VAE for Hunyuan Image 2.1 Base", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_2.1_vae_fp16.safetensors", "f2ae19863609206196b5e3a86bfd94f67bd3866f5042004e3994f07e3c93b2f9", "VAE", "HunyuanImage/hunyuan_image_2.1_vae_fp16.safetensors")); Register(new("hunyuan-image-2_1-refiner-vae", "Hunyuan Image 2.1 Refiner VAE", "The VAE for Hunyuan Image 2.1 Refiner", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_refiner_vae_fp16.safetensors", "e1b74e85d61b65e18cc05ca390e387d93cfadf161e737de229ebb800ea3db769", "VAE", "HunyuanImage/hunyuan_image_2.1_refiner_vae_fp16.safetensors")); - Register(new("seedvr2-vae", "SeedVR2 VAE", "The VAE for SeedVR2", "https://huggingface.co/Comfy-Org/SeedVR2/resolve/main/vae/seedvr2_ema_vae_fp16.safetensors", "20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1", "VAE", "SeedVR2/seedvr2_ema_vae_fp16.safetensors")); Register(new("hunyuan-video-1_5-vae", "Hunyuan Video 1.5 VAE", "The VAE for Hunyuan Video 1.5", "https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/resolve/main/split_files/vae/hunyuanvideo15_vae_fp16.safetensors", "e7c3091949c27e2d55ae6d5df917b99dadfebbf308e5a50d0ade0d16c90297ae", "VAE", "HunyuanVideo/hunyuanvideo15_vae_fp16.safetensors")); // Audio VAEs diff --git a/src/Text2Image/T2IModelClassSorter.cs b/src/Text2Image/T2IModelClassSorter.cs index 8face175f..18d841b8a 100644 --- a/src/Text2Image/T2IModelClassSorter.cs +++ b/src/Text2Image/T2IModelClassSorter.cs @@ -92,7 +92,6 @@ public static T2IModelCompatClass CompatLens = RegisterCompat(new() { ID = "lens", ShortCode = "Lens", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatPiD = RegisterCompat(new() { ID = "pid", ShortCode = "PiD", LorasTargetTextEnc = false }), CompatPixelDiT = RegisterCompat(new() { ID = "pixeldit", ShortCode = "PixDiT", LorasTargetTextEnc = false }), - CompatSeedVR2 = RegisterCompat(new() { ID = "seedvr2", ShortCode = "SeedVR2", LorasTargetTextEnc = false }), CompatIdeogram4 = RegisterCompat(new() { ID = "ideogram-4", ShortCode = "Ideo4", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatKrea2 = RegisterCompat(new() { ID = "krea-2", ShortCode = "Krea2", VaeFamily = VaeQwenImage }), CompatBoogu = RegisterCompat(new() { ID = "boogu", ShortCode = "Boogu", LorasTargetTextEnc = false, VaeFamily = VaeFlux1 }), @@ -242,8 +241,6 @@ bool isZImageLora(JObject h) => (hasLoraKey(h, "layers.0.adaLN_modulation.0") && bool isChromaRadiance(JObject h) => hasKey(h, "nerf_image_embedder.embedder.0.bias"); bool isPiD(JObject h) => h.ContainsKey("net.lq_proj.latent_proj.0.weight") && h.ContainsKey("net.pixel_blocks.0.attn.q_norm.weight") && h.ContainsKey("net.pixel_blocks.0.compress_to_attn.weight"); bool isPixelDiT(JObject h) => h.ContainsKey("core.pixel_embedder.proj.weight") && h.ContainsKey("core.pixel_blocks.0.attn.q_norm.weight") && h.ContainsKey("core.pixel_blocks.0.compress_to_attn.weight") && !isPiD(h); - bool isSeedVR2(JObject h) => hasKey(h, "blocks.31.mlp.all.proj_in_gate.weight") || hasKey(h, "blocks.35.mlp.all.proj_in_gate.weight") || hasKey(h, "blocks.35.mlp.vid.proj_out.weight"); - bool isSeedVR2Vae(JObject h) => h.ContainsKey("decoder.up_blocks.2.upsamplers.0.upscale_conv.weight"); bool isOmniGen(JObject h) => h.ContainsKey("time_caption_embed.timestep_embedder.linear_2.weight") && h.ContainsKey("context_refiner.0.attn.norm_k.weight") && !isBoogu(h); bool isBoogu(JObject h) => hasKey(h, "double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") && hasKey(h, "double_stream_layers.0.img_instruct_attn.processor.instruct_to_q.weight"); bool isQwenImage(JObject h) => (h.ContainsKey("time_text_embed.timestep_embedder.linear_1.bias") && h.ContainsKey("img_in.bias") && (h.ContainsKey("transformer_blocks.0.attn.add_k_proj.bias") || h.ContainsKey("transformer_blocks.0.attn.add_qkv_proj.bias"))) @@ -779,15 +776,6 @@ JToken GetEmbeddingKey(JObject h) { return isPixelDiT(h); }}); - // ====================== SeedVR2 ====================== - Register(new() { ID = "seedvr2", CompatClass = CompatSeedVR2, Name = "SeedVR2", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => - { - return isSeedVR2(h); - }}); - Register(new() { ID = "seedvr2/vae", CompatClass = CompatSeedVR2, Name = "SeedVR2 VAE", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => - { - return isSeedVR2Vae(h); - }}); Register(new() { ID = "alt_diffusion_v1_512_placeholder", CompatClass = CompatAltDiffusion, Name = "Alt-Diffusion", StandardWidth = 512, StandardHeight = 512, IsThisModelOfClass = (m, h) => { return IsAlt(h); diff --git a/src/Text2Image/T2IParamTypes.cs b/src/Text2Image/T2IParamTypes.cs index 062639ca9..9b95f10a8 100644 --- a/src/Text2Image/T2IParamTypes.cs +++ b/src/Text2Image/T2IParamTypes.cs @@ -326,7 +326,7 @@ public static string ApplyStringEdit(string prior, string update) public static T2IRegisteredParam Prompt, NegativePrompt, AspectRatio, BackendType, RefinerMethod, FreeUApplyTo, FreeUVersion, PersonalNote, VideoFormat, VideoResolution, UnsamplerPrompt, ImageFormat, MaskBehavior, ColorCorrectionBehavior, RawResolution, SeamlessTileable, SD3TextEncs, BitDepth, Webhooks, WildcardSeedBehavior, SegmentSortOrder, SegmentTargetResolution, SegmentApplyAfter, TorchCompile, VideoExtendFormat, ExactBackendID, OverridePredictionType, OverrideOutpathFormat, Text2AudioTimeSignature, Text2AudioLanguage, Text2AudioKeyScale, Text2AudioStyle; public static T2IRegisteredParam Images, Steps, Width, Height, SideLength, BatchSize, VAETileSize, VAETileOverlap, VAETemporalTileSize, VAETemporalTileOverlap, ClipStopAtLayer, VideoFrames, VideoMotionBucket, VideoFPS, VideoSteps, RefinerSteps, CascadeLatentCompression, MaskShrinkGrow, MaskBlur, MaskGrow, SegmentMaskBlur, SegmentMaskGrow, SegmentMaskOversize, SegmentSteps, Text2VideoFrames, TrimVideoStartFrames, TrimVideoEndFrames, VideoExtendFrameOverlap; public static T2IRegisteredParam Seed, VariationSeed, WildcardSeed, Text2AudioBPM; - public static T2IRegisteredParam CFGScale, VariationSeedStrength, InitImageCreativity, InitImageResetToNorm, InitImageNoise, RefinerControl, RefinerUpscale, RefinerCFGScale, FinalUpscale, ReVisionStrength, AltResolutionHeightMult, + public static T2IRegisteredParam CFGScale, VariationSeedStrength, InitImageCreativity, InitImageResetToNorm, InitImageNoise, RefinerControl, RefinerUpscale, RefinerCFGScale, ReVisionStrength, AltResolutionHeightMult, FreeUBlock1, FreeUBlock2, FreeUSkip1, FreeUSkip2, GlobalRegionFactor, EndStepsEarly, SamplerSigmaMin, SamplerSigmaMax, SamplerRho, VideoAugmentationLevel, VideoCFG, VideoMinCFG, Video2VideoCreativity, VideoSwapPercent, VideoExtendSwapPercent, IP2PCFG2, RegionalObjectCleanupFactor, SigmaShift, SegmentThresholdMax, SegmentCFGScale, FluxGuidanceScale, Text2AudioDuration; public static T2IRegisteredParam InitImage, MaskImage, VideoEndFrame; public static T2IRegisteredParam VideoAudioInput, VideoAudioReference; @@ -336,7 +336,7 @@ public static string ApplyStringEdit(string prior, string update) public static T2IRegisteredParam OutputIntermediateImages, DoNotSave, DoNotSaveIntermediates, ControlNetPreviewOnly, RevisionZeroPrompt, RemoveBackground, NoSeedIncrement, NoPreviews, VideoBoomerang, ModelSpecificEnhancements, UseInpaintingEncode, MaskCompositeUnthresholded, SaveSegmentMask, InitImageRecompositeMask, UseReferenceOnly, RefinerDoTiling, AutomaticVAE, ZeroNegative, FluxDisableGuidance, SmartImagePromptResizing, NoLoadModels, NoInternalSpecialHandling, ForwardRawBackendData, ForwardSwarmData, NegativeModelIncludeLoras, ContinueAfterErrors, PlaceholderParamGroupStarred, PlaceholderParamGroupUser1, PlaceholderParamGroupUser2, PlaceholderParamGroupUser3; - public static T2IParamGroup GroupImagePrompting, GroupCore, GroupVariation, GroupResolution, GroupSampling, GroupInitImage, GroupRefiners, GroupRefinerOverrides, GroupFinalStage, + public static T2IParamGroup GroupImagePrompting, GroupCore, GroupVariation, GroupResolution, GroupSampling, GroupInitImage, GroupRefiners, GroupRefinerOverrides, GroupAdvancedModelAddons, GroupSwarmInternal, GroupFreeU, GroupRegionalPrompting, GroupSegmentRefining, GroupSegmentOverrides, GroupAdvancedSampling, GroupAlternateGuidance, GroupVideo, GroupText2Video, GroupAdvancedVideo, GroupAdvancedVideoObscure, GroupVideoExtend, GroupText2Audio, GroupStarred, GroupUser1, GroupUser2, GroupUser3; @@ -552,11 +552,6 @@ static List listVaes(Session s) + "Too-high values can cause corrupted/burnt images, too-low can cause nonsensical images.\n7 is a good baseline. Normal usages vary between 4 and 9.\nSome model types, such as Turbo, expect CFG around 1.", "7", Min: 0, Max: 100, ViewMax: 20, Step: 0.5, Examples: ["5", "6", "7", "8", "9"], OrderPriority: -4, ViewType: ParamViewType.SLIDER, Group: GroupRefinerOverrides, ChangeWeight: -3, Toggleable: true, IsAdvanced: true )); - // ================================================ Final Stage ================================================ - GroupFinalStage = new("Final Stage", Toggles: true, Open: false, OrderPriority: -2.5, Description: "An upscale applied after the base and refiner stages are done.\nUnlike the Refine/Upscale group, nothing samples over the result after, only whatever the upscaler does itself."); - FinalUpscale = Register(new("Final Upscale", "How much to upscale by in the final stage.\nThis stacks on top of Refiner Upscale, so 1.5 refiner upscale and 2 final upscale is 3x total.", - "2", Min: 0.25, Max: 8, ViewMax: 4, Step: 0.25, OrderPriority: -2, ViewType: ParamViewType.SLIDER, Group: GroupFinalStage, DoNotPreview: true, Examples: ["1.5", "2", "4"] - )); // ================================================ ControlNet ================================================ for (int i = 1; i <= 3; i++) { From 5d3bc95c66b0d75d6ca9616439ed94dd435a0a61 Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Fri, 31 Jul 2026 19:34:21 -0500 Subject: [PATCH 5/7] SeedVR support as a group --- docs/Features/Upscaling.md | 8 + .../ComfyUIBackend/ComfyUIBackendExtension.cs | 37 ++++- .../ComfyUIBackend/WorkflowGenerator.cs | 150 +++++++++++++++++- .../WorkflowGeneratorModelSupport.cs | 11 +- .../ComfyUIBackend/WorkflowGeneratorSteps.cs | 39 +++-- src/Text2Image/CommonModels.cs | 1 + src/Text2Image/T2IModelClassSorter.cs | 12 ++ 7 files changed, 240 insertions(+), 18 deletions(-) diff --git a/docs/Features/Upscaling.md b/docs/Features/Upscaling.md index 0da35122e..3dda27e94 100644 --- a/docs/Features/Upscaling.md +++ b/docs/Features/Upscaling.md @@ -7,3 +7,11 @@ (TODO) Downloads here: + +# SeedVR2 + +- [SeedVR2]() is a one-step restoration model for upscaling images or videos. + - Models can be downloaded here: [Comfy-Org/SeedVR2]() + - Save in `diffusion_models` + - The VAE will be automatically downloaded +- Select it as your `Refiner Upscale Method` or your `Final Upscale Method`, see [Upscale Stages](#upscale-stages) \ No newline at end of file diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs index 999ebc564..05ef49575 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs @@ -68,7 +68,8 @@ public record class ComfyCustomWorkflow(string Name, string Workflow, string Pro ["TeaCache"] = "teacache", ["TeaCacheForVidGen"] = "teacache", ["TeaCacheForImgGen"] = "teacache_oldvers", - ["OverrideCLIPDevice"] = "set_clip_device" + ["OverrideCLIPDevice"] = "set_clip_device", + ["SeedVR2Conditioning"] = "seedvr2", }; /// @@ -606,15 +607,17 @@ public static void AssignValuesFromRaw(JObject rawObjectInfo) } } - public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice; + public static T2IRegisteredParam CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice, SeedVRUpscaleMethod, SeedVRColorCorrectionBehavior; - public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG; + public static T2IRegisteredParam AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG, SeedVRSplitLatent; - public static T2IRegisteredParam IPAdapterWeight, IPAdapterStart, IPAdapterEnd, SelfAttentionGuidanceScale, SelfAttentionGuidanceSigmaBlur, PerturbedAttentionGuidanceScale, StyleModelMergeStrength, StyleModelApplyStart, StyleModelMultiplyStrength, RescaleCFGMultiplier, TeaCacheThreshold, TeaCacheStart, NunchakuCacheThreshold, EasyCacheThreshold, EasyCacheStart, EasyCacheEnd, RenormCFG, NormalizedAttentionGuidanceScale, NormalizedAttentionGuidanceAlpha, NormalizedAttentionGuidanceTau; + public static T2IRegisteredParam IPAdapterWeight, IPAdapterStart, IPAdapterEnd, SelfAttentionGuidanceScale, SelfAttentionGuidanceSigmaBlur, PerturbedAttentionGuidanceScale, StyleModelMergeStrength, StyleModelApplyStart, StyleModelMultiplyStrength, RescaleCFGMultiplier, TeaCacheThreshold, TeaCacheStart, NunchakuCacheThreshold, EasyCacheThreshold, EasyCacheStart, EasyCacheEnd, RenormCFG, NormalizedAttentionGuidanceScale, NormalizedAttentionGuidanceAlpha, NormalizedAttentionGuidanceTau, SeedVRUpscale; - public static T2IRegisteredParam RefinerHyperTile, VideoFrameInterpolationMultiplier; + public static T2IRegisteredParam RefinerHyperTile, VideoFrameInterpolationMultiplier, SeedVRTemporalVideoOverlap; - public static T2IRegisteredParam PixelDecoderModel; + public static T2IRegisteredParam PixelDecoderModel, SeedVRModel; + + public static T2IParamGroup GroupSeedVR; public static T2IRegisteredParam[] ControlNetPreprocessorParams = new T2IRegisteredParam[3], ControlNetUnionTypeParams = new T2IRegisteredParam[3]; @@ -799,6 +802,28 @@ public override void OnInit() RefinerHyperTile = T2IParamTypes.Register(new("Refiner HyperTile", "The size of hypertiles to use for the refining stage.\nHyperTile is a technique to speed up sampling of large images by tiling the image and batching the tiles.\nThis is useful when using SDv1 models as the refiner. SDXL-Base models do not benefit as much.", "256", Min: 64, Max: 2048, Step: 32, Toggleable: true, IsAdvanced: true, FeatureFlag: "comfyui", ViewType: ParamViewType.POT_SLIDER, Group: T2IParamTypes.GroupAdvancedSampling, OrderPriority: 20 )); + // ================================================ SeedVR ================================================ + GroupSeedVR = new T2IParamGroup("SeedVR", Toggles: true, Open: false, OrderPriority: -2.5, Description: "SeedVR2 is a one-step restoration model, run over the result of the normal generation."); + SeedVRModel = T2IParamTypes.Register(new("SeedVR Model", "Which SeedVR2 model to restore with.", + "", Toggleable: true, FeatureFlag: "seedvr2", Group: GroupSeedVR, Subtype: "Stable-Diffusion", ChangeWeight: 9, DoNotPreview: true, OrderPriority: -10, + GetValues: (session) => T2IParamTypes.CleanModelList(Program.MainSDModels.ListModelsFor(session).Where(m => m.ModelClass?.CompatClass?.ID == "seedvr2").OrderBy(m => m.Name).Select(m => m.Name)) + )); + SeedVRUpscale = T2IParamTypes.Register(new("SeedVR Upscale", "Optional upscale of the image before SeedVR2 runs over it.\nSetting to '1' disables the upscale, and just restores at the current size.", + "1", Min: 0.25, Max: 8, ViewMax: 4, Step: 0.25, OrderPriority: -9, ViewType: ParamViewType.SLIDER, FeatureFlag: "seedvr2", Group: GroupSeedVR, DoNotPreview: true, Examples: ["1", "1.5", "2"] + )); + SeedVRUpscaleMethod = T2IParamTypes.Register(new("SeedVR Upscale Method", "How to upscale the image before SeedVR2 runs over it, if upscaling is used.", + "pixel-lanczos", OrderPriority: -8, FeatureFlag: "seedvr2", Group: GroupSeedVR, ChangeWeight: 1, + GetValues: (session) => RefinerUpscaleMethod.Type.GetValues(session), DependNonDefault: SeedVRUpscale.Type.ID + )); + SeedVRColorCorrectionBehavior = T2IParamTypes.Register(new("SeedVR Color Correction Behavior", "How to match the colors of a SeedVR2 restore back to the image it was given.\n'None' = Do not attempt color correction, only align the geometry.\n'CIELAB' = Transfer the color in CIELAB space, preserving detail.\n'Wavelet' = Transfer the low-frequency color, keeping the upscaled high-frequency detail.\n'AdaIN' = Match the per-channel mean and standard deviation.", + "none", IgnoreIf: "none", FeatureFlag: "seedvr2", Group: GroupSeedVR, IsAdvanced: true, OrderPriority: 1, GetValues: (_) => ["none///None", "lab///CIELAB", "wavelet///Wavelet", "adain///AdaIN"] + )); + SeedVRSplitLatent = T2IParamTypes.Register(new("SeedVR Split Latent", "If enabled, samples a SeedVR2 video restore as chunks of frames instead of all at once, sized to fit in free VRAM.\nChunking reduces VRAM consumption.\nDoes nothing to a single image, or to a video that already fits.", + "false", IgnoreIf: "false", FeatureFlag: "seedvr2", Group: GroupSeedVR, IsAdvanced: true, OrderPriority: 2 + )); + SeedVRTemporalVideoOverlap = T2IParamTypes.Register(new("SeedVR Temporal Video Overlap", "How many frames of overlap to keep between 'SeedVR Split Latent' chunks.\nHigher overlap hides the chunk seams better but takes longer.", + "0", Min: 0, Max: 4096, Step: 1, IsAdvanced: true, FeatureFlag: "seedvr2", Group: GroupSeedVR, OrderPriority: 3 + )); List interpolators = ["RIFE", "FILM", "GIMM-VFI"]; VideoPreviewType = T2IParamTypes.Register(new("Video Preview Type", "How to display previews for generating videos.\n'Animate' shows a low-res animated video preview.\n'iterate' shows one frame at a time while it goes.\n'one' displays just the first frame.\n'none' disables previews.", "animate", IgnoreIf: "animate", FeatureFlag: "comfyui", Group: T2IParamTypes.GroupAdvancedVideo, Permission: Permissions.ParamVideo, IsAdvanced: true, GetValues: (_) => ["animate", "iterate", "one", "none"] diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index 286f46154..ec090bbd9 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -2066,8 +2066,12 @@ public void CreateImageToVideo(ImageToVideoGenInfo genInfo) ["upscale_method"] = "lanczos", ["crop"] = "disabled" }); - // TODO: Update width/height properly CurrentMedia = CurrentMedia.WithPath([scaled, 0]); + if (genInfo.Width?.Type == JTokenType.Integer && genInfo.Height?.Type == JTokenType.Integer) + { + CurrentMedia.Width = (int)genInfo.Width; + CurrentMedia.Height = (int)genInfo.Height; + } foreach (Action altHandler in AltImageToVideoPreHandlers) { altHandler(genInfo); @@ -2670,6 +2674,150 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD return result; } + /// Creates a SeedVR2 restoration stage: converts to a SeedVR2-space latent and samples a one-step restore from it. + public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, WGNodeData decodeVae, long seed) + { + WGNodeData raw = media.AsRawImage(decodeVae); + JArray resized = raw.Path; + string preprocessed = CreateNode("SeedVR2Preprocess", new JObject() + { + ["resized_images"] = resized + }); + T2IModel priorFinalModel = FinalLoadedModel; + List priorFinalModelList = FinalLoadedModelList; + WGNodeData priorModel = CurrentModel, priorTextEnc = CurrentTextEnc, priorVae = CurrentVae; + bool priorNoVae = NoVAEOverride, priorRefinerStage = IsRefinerStage; + FinalLoadedModel = seedVrModel; + FinalLoadedModelList = [seedVrModel]; + NoVAEOverride = true; + IsRefinerStage = false; + (FinalLoadedModel, CurrentModel, CurrentTextEnc, CurrentVae) = CreateModelLoader(seedVrModel, "SeedVR2"); + IsRefinerStage = priorRefinerStage; + NoVAEOverride = priorNoVae; + WGNodeData encoded = raw.WithPath([preprocessed, 0]).EncodeToLatent(CurrentVae); + JArray latent = encoded.Path; + JArray chunkOverlap = null; + if (UserInput.Get(ComfyUIBackendExtension.SeedVRSplitLatent, false)) + { + int overlap = UserInput.Get(ComfyUIBackendExtension.SeedVRTemporalVideoOverlap, 0); + string chunked = CreateNode("SeedVR2TemporalChunk", new JObject() + { + ["latent"] = latent, + ["temporal_overlap"] = overlap, + ["chunking_mode"] = "auto" + }); + latent = [chunked, 0]; + chunkOverlap = [chunked, 1]; + } + string cond = CreateNode("SeedVR2Conditioning", new JObject() + { + ["model"] = CurrentModel.Path, + ["vae_conditioning"] = latent + }); + string sampled = CreateKSampler(CurrentModel.Path, [cond, 0], [cond, 1], latent, 1, 1, 0, 10000, seed, false, true, explicitSampler: "euler", explicitScheduler: "normal"); + JArray sampledLatent = [sampled, 0]; + if (chunkOverlap is not null) + { + string merged = CreateNode("SeedVR2TemporalMerge", new JObject() + { + ["latents"] = sampledLatent, + ["temporal_overlap"] = chunkOverlap + }); + sampledLatent = [merged, 0]; + } + WGNodeData decoded = encoded.WithPath(sampledLatent).DecodeLatents(CurrentVae, false); + string post = CreateNode("SeedVR2PostProcessing", new JObject() + { + ["images"] = decoded.Path, + ["original_resized_images"] = resized, + ["color_correction_method"] = UserInput.Get(ComfyUIBackendExtension.SeedVRColorCorrectionBehavior, "none") + }); + WGNodeData result = raw.WithPath([post, 0]); + result.Width = raw.Width; + result.Height = raw.Height; + FinalLoadedModel = priorFinalModel; + FinalLoadedModelList = priorFinalModelList; + CurrentModel = priorModel; + CurrentTextEnc = priorTextEnc; + CurrentVae = priorVae; + return result; + } + + /// Runs the SeedVR group's restoration pass over the current media. + public void RunSeedVR2Stage(WGNodeData vaeOverride = null) + { + if (!UserInput.TryGet(ComfyUIBackendExtension.SeedVRModel, out T2IModel seedVrModel) || seedVrModel is null) + { + return; + } + WGNodeData vae = vaeOverride ?? CurrentVae; + if (UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) + { + CurrentMedia.SaveOutput(vae, CurrentAudioVae, GetStableDynamicID(50000, 0)); + } + long seed = UserInput.Get(T2IParamTypes.Seed) + 500; + WGNodeData media = CurrentMedia; + double scale = UserInput.Get(ComfyUIBackendExtension.SeedVRUpscale, 1); + if (scale != 1) + { + string method = UserInput.Get(ComfyUIBackendExtension.SeedVRUpscaleMethod, "pixel-lanczos"); + int width = (((int)Math.Round((media.Width ?? UserInput.GetImageWidth()) * scale)) / 16) * 16; + int height = (((int)Math.Round((media.Height ?? UserInput.GetImageHeight()) * scale)) / 16) * 16; + if (method.StartsWith("latent-")) + { + media = media.AsLatentImage(vae); + string latentUpscaled = CreateNode("LatentUpscaleBy", new JObject() + { + ["samples"] = media.Path, + ["upscale_method"] = method.After("latent-"), + ["scale_by"] = scale + }); + media = media.WithPath([latentUpscaled, 0]); + media.Width = width; + media.Height = height; + } + else + { + media = media.AsRawImage(vae); + if (method.StartsWith("pidmodel-")) + { + T2IModel pidModel = ComfyUIBackendExtension.GetPidModel(method.After("pidmodel-"), UserInput.SourceSession); + media = CreatePixelDecode(pidModel, media, vae, seed); + } + else if (method.StartsWith("model-")) + { + string loaderNode = CreateNode("UpscaleModelLoader", new JObject() + { + ["model_name"] = method.After("model-") + }); + string upscaledNode = CreateNode("ImageUpscaleWithModel", new JObject() + { + ["upscale_model"] = NodePath(loaderNode, 0), + ["image"] = media.Path + }); + media = media.WithPath([upscaledNode, 0]); + media.Width = null; // the model's own scale factor is unknown here, so always correct after + media.Height = null; + } + if (media.Width != width || media.Height != height) + { + string scaled = CreateNode("ImageScale", new JObject() + { + ["image"] = media.Path, + ["width"] = width, + ["height"] = height, + ["upscale_method"] = method.StartsWith("pixel-") ? method.After("pixel-") : "lanczos", + ["crop"] = "disabled" + }); + media = media.WithPath([scaled, 0]); + media.Width = width; + media.Height = height; + } + } + } + CurrentMedia = CreateSeedVR2Restore(seedVrModel, media, vae, seed); + } + /// Creates a "CLIPTextEncode" or equivalent node for the given input, applying prompt-given conditioning modifiers as relevant. public JArray CreateConditioning(string prompt, JArray clip, T2IModel model, bool isPositive, string firstId = null, bool isRefiner = false, bool isVideo = false, bool isVideoSwap = false, bool isPixelDecoder = false) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs index 54be505c4..0d82ba297 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs @@ -94,6 +94,9 @@ public bool IsKontext() /// Returns true if the current model is PiD. public bool IsPiD() => IsModelCompatClass(T2IModelClassSorter.CompatPiD); + /// Returns true if the current model is SeedVR2. + public bool IsSeedVR2() => IsModelCompatClass(T2IModelClassSorter.CompatSeedVR2); + /// Returns true if the current model is HiDream-i1. public bool IsHiDream() => IsModelCompatClass(T2IModelClassSorter.CompatHiDreamI1); @@ -966,7 +969,7 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) { dtype = "default"; } - else if (IsZImage() || IsZetaChroma() || IsAnima() || IsLens() || IsPixelDiT() || IsPiD()) // Model is small and dense, so trust user preferred download format + else if (IsZImage() || IsZetaChroma() || IsAnima() || IsLens() || IsPixelDiT() || IsPiD() || IsSeedVR2()) // Model is small and dense, so trust user preferred download format { dtype = "default"; } @@ -1233,6 +1236,10 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) helpers.LoadClip("pixeldit", helpers.GetGemma2_2bElmModel()); LoadingVAE = CreateVAELoader("pixel_space"); } + else if (IsSeedVR2()) + { + helpers.DoVaeLoader(null, T2IModelClassSorter.CompatSeedVR2, "seedvr2-vae"); + } else if (IsHiDream()) { string loaderType = "QuadrupleCLIPLoader"; @@ -1484,7 +1491,7 @@ public void LoadClip3(string type, string modelA, string modelB, string modelC) { step.Action(this); } - if (LoadingClip is null) + if (LoadingClip is null && type != "SeedVR2") { if (string.IsNullOrWhiteSpace(model.Metadata?.ModelClassType)) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs index c2ed9900e..400e7dffa 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs @@ -131,7 +131,7 @@ public static void Register() }, -10); AddModelGenStep(g => { - if (g.UserInput.TryGet(ComfyUIBackendExtension.SetClipDevice, out string device) && g.Features.Contains("set_clip_device")) + if (g.LoadingClip is not null && g.UserInput.TryGet(ComfyUIBackendExtension.SetClipDevice, out string device) && g.Features.Contains("set_clip_device")) { string clipDeviceNode = g.CreateNode("OverrideCLIPDevice", new JObject() { @@ -231,7 +231,7 @@ public static void Register() }, -7); AddModelGenStep(g => { - if (g.UserInput.TryGet(T2IParamTypes.ClipStopAtLayer, out int layer)) + if (g.LoadingClip is not null && g.UserInput.TryGet(T2IParamTypes.ClipStopAtLayer, out int layer)) { string clipSkip = g.CreateNode("CLIPSetLastLayer", new JObject() { @@ -1525,13 +1525,17 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) { modelMustReencode = true; } - g.NoVAEOverride = refineModel.ModelClass?.CompatClass != baseModel.ModelClass?.CompatClass; - g.FinalLoadedModel = refineModel; - g.FinalLoadedModelList = [refineModel]; - (g.FinalLoadedModel, g.CurrentModel, g.CurrentTextEnc, g.CurrentVae) = g.CreateModelLoader(refineModel, "Refiner", loaderNodeId, sectionId: T2IParamInput.SectionID_Refiner); - g.NoVAEOverride = false; - prompt = g.CreateConditioning(g.UserInput.Get(T2IParamTypes.Prompt), g.CurrentTextEnc.Path, g.FinalLoadedModel, true, isRefiner: true); - negPrompt = g.CreateConditioning(g.UserInput.Get(T2IParamTypes.NegativePrompt), g.CurrentTextEnc.Path, g.FinalLoadedModel, false, isRefiner: true); + bool isSeedVr = refineModel.ModelClass?.CompatClass?.ID == "seedvr2"; + if (!isSeedVr) + { + g.NoVAEOverride = refineModel.ModelClass?.CompatClass != baseModel.ModelClass?.CompatClass; + g.FinalLoadedModel = refineModel; + g.FinalLoadedModelList = [refineModel]; + (g.FinalLoadedModel, g.CurrentModel, g.CurrentTextEnc, g.CurrentVae) = g.CreateModelLoader(refineModel, "Refiner", loaderNodeId, sectionId: T2IParamInput.SectionID_Refiner); + g.NoVAEOverride = false; + prompt = g.CreateConditioning(g.UserInput.Get(T2IParamTypes.Prompt), g.CurrentTextEnc.Path, g.FinalLoadedModel, true, isRefiner: true); + negPrompt = g.CreateConditioning(g.UserInput.Get(T2IParamTypes.NegativePrompt), g.CurrentTextEnc.Path, g.FinalLoadedModel, false, isRefiner: true); + } bool doSave = g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false); bool doUpscale = g.UserInput.TryGet(T2IParamTypes.RefinerUpscale, out double refineUpscale) && refineUpscale != 1; string upscaleMethod = g.UserInput.Get(ComfyUIBackendExtension.RefinerUpscaleMethod, "None"); @@ -1686,6 +1690,12 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) g.CurrentMedia.Width = width; g.CurrentMedia.Height = height; } + if (isSeedVr) + { + g.CurrentMedia = g.CreateSeedVR2Restore(refineModel, g.CurrentMedia, origVae, g.UserInput.Get(T2IParamTypes.Seed) + 1); + g.IsRefinerStage = false; + return; + } WGNodeData model = g.CurrentModel; if (g.UserInput.TryGet(ComfyUIBackendExtension.RefinerHyperTile, out int tileSize)) { @@ -1977,6 +1987,10 @@ void RunSegmentationProcessing(WorkflowGenerator g, bool isBeforeRefiner) bool willHaveFollowupVideo = g.UserInput.TryGet(T2IParamTypes.VideoModel, out _) || g.UserInput.Get(T2IParamTypes.Prompt, "").Contains(" 1) { if (g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false)) @@ -2156,6 +2174,7 @@ WGNodeData appendAudio(WGNodeData combinedAudio, WGNodeData nextAudio) string negPrompt = g.UserInput.Get(T2IParamTypes.NegativePrompt, ""); long seed = g.UserInput.Get(T2IParamTypes.Seed) + 600; int? videoFps = g.UserInput.TryGet(T2IParamTypes.VideoFPS, out int fpsRaw) ? fpsRaw : null; + WGNodeData extendVae = null; string format = g.UserInput.Get(T2IParamTypes.VideoExtendFormat, "mp4").ToLowerFast(); int frameExtendOverlap = g.UserInput.Get(T2IParamTypes.VideoExtendFrameOverlap, 9); bool saveIntermediate = g.UserInput.Get(T2IParamTypes.OutputIntermediateImages, false); @@ -2221,6 +2240,7 @@ WGNodeData appendAudio(WGNodeData combinedAudio, WGNodeData nextAudio) ContextID = part.ContextID }; g.CreateImageToVideo(genInfo); + extendVae = genInfo.Vae; g.CurrentMedia = g.CurrentMedia.AsRawImage(genInfo.Vae); WGNodeData stageWithAudio = ensureAttachedAudio(g.CurrentMedia); videoFps = genInfo.VideoFPS; @@ -2247,6 +2267,7 @@ WGNodeData appendAudio(WGNodeData combinedAudio, WGNodeData nextAudio) } g.CurrentMedia = conjoinedLast; g.CurrentMedia.FPS = videoFps ?? g.CurrentMedia.FPS; + g.RunSeedVR2Stage(extendVae); if (g.UserInput.TryGet(ComfyUIBackendExtension.VideoFrameInterpolationMethod, out string method) && g.UserInput.TryGet(ComfyUIBackendExtension.VideoFrameInterpolationMultiplier, out int mult) && mult > 1) { if (saveIntermediate) diff --git a/src/Text2Image/CommonModels.cs b/src/Text2Image/CommonModels.cs index 1c5c8ce79..09956a379 100644 --- a/src/Text2Image/CommonModels.cs +++ b/src/Text2Image/CommonModels.cs @@ -95,6 +95,7 @@ public static void RegisterCoreSet() Register(new("mage-flow-vae", "Mage-VAE", "The VAE for Mage-Flow", "https://huggingface.co/microsoft/Mage-Flow/resolve/main/vae/diffusion_pytorch_model.safetensors", "34e076dc1e8a15321e1e07be5111d59cf16dd10b804b7c7e20b4de29013427e0", "VAE", "MageFlow/mage_flow_vae.safetensors")); Register(new("hunyuan-image-2_1-vae", "Hunyuan Image 2.1 VAE", "The VAE for Hunyuan Image 2.1 Base", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_2.1_vae_fp16.safetensors", "f2ae19863609206196b5e3a86bfd94f67bd3866f5042004e3994f07e3c93b2f9", "VAE", "HunyuanImage/hunyuan_image_2.1_vae_fp16.safetensors")); Register(new("hunyuan-image-2_1-refiner-vae", "Hunyuan Image 2.1 Refiner VAE", "The VAE for Hunyuan Image 2.1 Refiner", "https://huggingface.co/Comfy-Org/HunyuanImage_2.1_ComfyUI/resolve/main/split_files/vae/hunyuan_image_refiner_vae_fp16.safetensors", "e1b74e85d61b65e18cc05ca390e387d93cfadf161e737de229ebb800ea3db769", "VAE", "HunyuanImage/hunyuan_image_2.1_refiner_vae_fp16.safetensors")); + Register(new("seedvr2-vae", "SeedVR2 VAE", "The VAE for SeedVR2", "https://huggingface.co/Comfy-Org/SeedVR2/resolve/main/vae/seedvr2_ema_vae_fp16.safetensors", "20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1", "VAE", "SeedVR2/seedvr2_ema_vae_fp16.safetensors")); Register(new("hunyuan-video-1_5-vae", "Hunyuan Video 1.5 VAE", "The VAE for Hunyuan Video 1.5", "https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/resolve/main/split_files/vae/hunyuanvideo15_vae_fp16.safetensors", "e7c3091949c27e2d55ae6d5df917b99dadfebbf308e5a50d0ade0d16c90297ae", "VAE", "HunyuanVideo/hunyuanvideo15_vae_fp16.safetensors")); // Audio VAEs diff --git a/src/Text2Image/T2IModelClassSorter.cs b/src/Text2Image/T2IModelClassSorter.cs index 18d841b8a..8face175f 100644 --- a/src/Text2Image/T2IModelClassSorter.cs +++ b/src/Text2Image/T2IModelClassSorter.cs @@ -92,6 +92,7 @@ public static T2IModelCompatClass CompatLens = RegisterCompat(new() { ID = "lens", ShortCode = "Lens", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatPiD = RegisterCompat(new() { ID = "pid", ShortCode = "PiD", LorasTargetTextEnc = false }), CompatPixelDiT = RegisterCompat(new() { ID = "pixeldit", ShortCode = "PixDiT", LorasTargetTextEnc = false }), + CompatSeedVR2 = RegisterCompat(new() { ID = "seedvr2", ShortCode = "SeedVR2", LorasTargetTextEnc = false }), CompatIdeogram4 = RegisterCompat(new() { ID = "ideogram-4", ShortCode = "Ideo4", LorasTargetTextEnc = false, VaeFamily = VaeFlux2 }), CompatKrea2 = RegisterCompat(new() { ID = "krea-2", ShortCode = "Krea2", VaeFamily = VaeQwenImage }), CompatBoogu = RegisterCompat(new() { ID = "boogu", ShortCode = "Boogu", LorasTargetTextEnc = false, VaeFamily = VaeFlux1 }), @@ -241,6 +242,8 @@ bool isZImageLora(JObject h) => (hasLoraKey(h, "layers.0.adaLN_modulation.0") && bool isChromaRadiance(JObject h) => hasKey(h, "nerf_image_embedder.embedder.0.bias"); bool isPiD(JObject h) => h.ContainsKey("net.lq_proj.latent_proj.0.weight") && h.ContainsKey("net.pixel_blocks.0.attn.q_norm.weight") && h.ContainsKey("net.pixel_blocks.0.compress_to_attn.weight"); bool isPixelDiT(JObject h) => h.ContainsKey("core.pixel_embedder.proj.weight") && h.ContainsKey("core.pixel_blocks.0.attn.q_norm.weight") && h.ContainsKey("core.pixel_blocks.0.compress_to_attn.weight") && !isPiD(h); + bool isSeedVR2(JObject h) => hasKey(h, "blocks.31.mlp.all.proj_in_gate.weight") || hasKey(h, "blocks.35.mlp.all.proj_in_gate.weight") || hasKey(h, "blocks.35.mlp.vid.proj_out.weight"); + bool isSeedVR2Vae(JObject h) => h.ContainsKey("decoder.up_blocks.2.upsamplers.0.upscale_conv.weight"); bool isOmniGen(JObject h) => h.ContainsKey("time_caption_embed.timestep_embedder.linear_2.weight") && h.ContainsKey("context_refiner.0.attn.norm_k.weight") && !isBoogu(h); bool isBoogu(JObject h) => hasKey(h, "double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") && hasKey(h, "double_stream_layers.0.img_instruct_attn.processor.instruct_to_q.weight"); bool isQwenImage(JObject h) => (h.ContainsKey("time_text_embed.timestep_embedder.linear_1.bias") && h.ContainsKey("img_in.bias") && (h.ContainsKey("transformer_blocks.0.attn.add_k_proj.bias") || h.ContainsKey("transformer_blocks.0.attn.add_qkv_proj.bias"))) @@ -776,6 +779,15 @@ JToken GetEmbeddingKey(JObject h) { return isPixelDiT(h); }}); + // ====================== SeedVR2 ====================== + Register(new() { ID = "seedvr2", CompatClass = CompatSeedVR2, Name = "SeedVR2", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => + { + return isSeedVR2(h); + }}); + Register(new() { ID = "seedvr2/vae", CompatClass = CompatSeedVR2, Name = "SeedVR2 VAE", StandardWidth = 1024, StandardHeight = 1024, IsThisModelOfClass = (m, h) => + { + return isSeedVR2Vae(h); + }}); Register(new() { ID = "alt_diffusion_v1_512_placeholder", CompatClass = CompatAltDiffusion, Name = "Alt-Diffusion", StandardWidth = 512, StandardHeight = 512, IsThisModelOfClass = (m, h) => { return IsAlt(h); From 1556ded4a40ae2d518bb1f0d1b262bf3fa6dfb23 Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Fri, 31 Jul 2026 19:51:04 -0500 Subject: [PATCH 6/7] Respect refiner values --- .../ComfyUIBackend/WorkflowGenerator.cs | 18 +++++++++++------- .../ComfyUIBackend/WorkflowGeneratorSteps.cs | 6 +++--- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs index ec090bbd9..c67f41002 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGenerator.cs @@ -2674,8 +2674,8 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD return result; } - /// Creates a SeedVR2 restoration stage: converts to a SeedVR2-space latent and samples a one-step restore from it. - public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, WGNodeData decodeVae, long seed) + /// Creates a SeedVR2 restoration stage: converts to a SeedVR2-space latent and samples the restored image from it. + public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, WGNodeData decodeVae, long seed, bool isRefiner = false) { WGNodeData raw = media.AsRawImage(decodeVae); JArray resized = raw.Path; @@ -2686,13 +2686,12 @@ public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, W T2IModel priorFinalModel = FinalLoadedModel; List priorFinalModelList = FinalLoadedModelList; WGNodeData priorModel = CurrentModel, priorTextEnc = CurrentTextEnc, priorVae = CurrentVae; - bool priorNoVae = NoVAEOverride, priorRefinerStage = IsRefinerStage; + bool priorNoVae = NoVAEOverride; + int sectionId = isRefiner ? T2IParamInput.SectionID_Refiner : 0; FinalLoadedModel = seedVrModel; FinalLoadedModelList = [seedVrModel]; NoVAEOverride = true; - IsRefinerStage = false; - (FinalLoadedModel, CurrentModel, CurrentTextEnc, CurrentVae) = CreateModelLoader(seedVrModel, "SeedVR2"); - IsRefinerStage = priorRefinerStage; + (FinalLoadedModel, CurrentModel, CurrentTextEnc, CurrentVae) = CreateModelLoader(seedVrModel, "SeedVR2", sectionId: sectionId); NoVAEOverride = priorNoVae; WGNodeData encoded = raw.WithPath([preprocessed, 0]).EncodeToLatent(CurrentVae); JArray latent = encoded.Path; @@ -2714,7 +2713,12 @@ public WGNodeData CreateSeedVR2Restore(T2IModel seedVrModel, WGNodeData media, W ["model"] = CurrentModel.Path, ["vae_conditioning"] = latent }); - string sampled = CreateKSampler(CurrentModel.Path, [cond, 0], [cond, 1], latent, 1, 1, 0, 10000, seed, false, true, explicitSampler: "euler", explicitScheduler: "normal"); + int steps = UserInput.GetNullable(T2IParamTypes.Steps, sectionId, false) ?? (isRefiner ? UserInput.GetNullable(T2IParamTypes.RefinerSteps) : null) ?? 1; + double cfg = UserInput.GetNullable(T2IParamTypes.CFGScale, sectionId, false) ?? (isRefiner ? UserInput.GetNullable(T2IParamTypes.RefinerCFGScale) : null) ?? 1; + string explicitSampler = UserInput.Get(ComfyUIBackendExtension.SamplerParam, null, sectionId: sectionId, includeBase: false) ?? (isRefiner ? UserInput.Get(ComfyUIBackendExtension.RefinerSamplerParam, null) : null); + string explicitScheduler = UserInput.Get(ComfyUIBackendExtension.SchedulerParam, null, sectionId: sectionId, includeBase: false) ?? (isRefiner ? UserInput.Get(ComfyUIBackendExtension.RefinerSchedulerParam, null) : null); + int startStep = isRefiner ? (int)Math.Round(steps * (1 - UserInput.Get(T2IParamTypes.RefinerControl, 1))) : 0; + string sampled = CreateKSampler(CurrentModel.Path, [cond, 0], [cond, 1], latent, cfg, steps, startStep, 10000, seed, false, true, explicitSampler: explicitSampler ?? "euler", explicitScheduler: explicitScheduler ?? "normal", sectionId: sectionId); JArray sampledLatent = [sampled, 0]; if (chunkOverlap is not null) { diff --git a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs index 400e7dffa..c269723aa 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorSteps.cs @@ -1573,7 +1573,7 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) g.CurrentMedia = decoded; return; } - g.CurrentMedia = decoded.EncodeToLatent(g.CurrentVae, "25"); + g.CurrentMedia = isSeedVr ? decoded : decoded.EncodeToLatent(g.CurrentVae, "25"); } else if (modelMustReencode || doPixelUpscale || doSave || g.MaskShrunkInfo.BoundsNode is not null) { @@ -1629,7 +1629,7 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) } if (modelMustReencode || doPixelUpscale) { - g.CurrentMedia = decoded.EncodeToLatent(g.CurrentVae, "25"); + g.CurrentMedia = isSeedVr ? decoded : decoded.EncodeToLatent(g.CurrentVae, "25"); } } if (doUpscale && upscaleMethod.StartsWith("latent-")) @@ -1692,7 +1692,7 @@ JArray doMaskShrinkApply(WorkflowGenerator g, JArray imgIn) } if (isSeedVr) { - g.CurrentMedia = g.CreateSeedVR2Restore(refineModel, g.CurrentMedia, origVae, g.UserInput.Get(T2IParamTypes.Seed) + 1); + g.CurrentMedia = g.CreateSeedVR2Restore(refineModel, g.CurrentMedia, origVae, g.UserInput.Get(T2IParamTypes.Seed) + 1, isRefiner: true); g.IsRefinerStage = false; return; } From 5b03f3ca2b4ccf55cd4df23af7f9b63afc92c582 Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Fri, 31 Jul 2026 20:57:08 -0500 Subject: [PATCH 7/7] Fix model dropdown --- .../ComfyUIBackend/ComfyUIBackendExtension.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs index 05ef49575..fdefcd104 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIBackendExtension.cs @@ -805,8 +805,8 @@ public override void OnInit() // ================================================ SeedVR ================================================ GroupSeedVR = new T2IParamGroup("SeedVR", Toggles: true, Open: false, OrderPriority: -2.5, Description: "SeedVR2 is a one-step restoration model, run over the result of the normal generation."); SeedVRModel = T2IParamTypes.Register(new("SeedVR Model", "Which SeedVR2 model to restore with.", - "", Toggleable: true, FeatureFlag: "seedvr2", Group: GroupSeedVR, Subtype: "Stable-Diffusion", ChangeWeight: 9, DoNotPreview: true, OrderPriority: -10, - GetValues: (session) => T2IParamTypes.CleanModelList(Program.MainSDModels.ListModelsFor(session).Where(m => m.ModelClass?.CompatClass?.ID == "seedvr2").OrderBy(m => m.Name).Select(m => m.Name)) + "None", IgnoreIf: "None", FeatureFlag: "seedvr2", Group: GroupSeedVR, Subtype: "Stable-Diffusion", ChangeWeight: 9, DoNotPreview: true, OrderPriority: -10, + GetValues: (session) => ["None", .. T2IParamTypes.CleanModelList(Program.MainSDModels.ListModelsFor(session).Where(m => m.ModelClass?.CompatClass?.ID == "seedvr2").OrderBy(m => m.Name).Select(m => m.Name))] )); SeedVRUpscale = T2IParamTypes.Register(new("SeedVR Upscale", "Optional upscale of the image before SeedVR2 runs over it.\nSetting to '1' disables the upscale, and just restores at the current size.", "1", Min: 0.25, Max: 8, ViewMax: 4, Step: 0.25, OrderPriority: -9, ViewType: ParamViewType.SLIDER, FeatureFlag: "seedvr2", Group: GroupSeedVR, DoNotPreview: true, Examples: ["1", "1.5", "2"]