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..fdefcd104 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.", + "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"] + )); + 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..c67f41002 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,154 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD return result; } + /// 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; + 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; + int sectionId = isRefiner ? T2IParamInput.SectionID_Refiner : 0; + FinalLoadedModel = seedVrModel; + FinalLoadedModelList = [seedVrModel]; + NoVAEOverride = true; + (FinalLoadedModel, CurrentModel, CurrentTextEnc, CurrentVae) = CreateModelLoader(seedVrModel, "SeedVR2", sectionId: sectionId); + 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 + }); + 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) + { + 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..c269723aa 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"); @@ -1569,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) { @@ -1625,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-")) @@ -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, isRefiner: true); + 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);