Adds native SeedVR2 support

docs/Features/Upscaling.md

@@ -7,3 +7,11 @@
 (TODO)
 
 Downloads here: <https://huggingface.co/Comfy-Org/PixelDiT/tree/main/diffusion_models>
+
+# SeedVR2
+
+- [SeedVR2](<https://huggingface.co/Comfy-Org/SeedVR2>) is a one-step restoration model for upscaling images or videos.
+    - Models can be downloaded here: [Comfy-Org/SeedVR2](<https://huggingface.co/Comfy-Org/SeedVR2/tree/main/diffusion_models>)
+    - 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)

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",
     };
 
     /// <inheritdoc/>
@@ -606,15 +607,17 @@ public static void AssignValuesFromRaw(JObject rawObjectInfo)
         }
     }
 
-    public static T2IRegisteredParam<string> CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice;
+    public static T2IRegisteredParam<string> CustomWorkflowParam, SamplerParam, SchedulerParam, RefinerSamplerParam, RefinerSchedulerParam, RefinerUpscaleMethod, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice, SeedVRUpscaleMethod, SeedVRColorCorrectionBehavior;
 
-    public static T2IRegisteredParam<bool> AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG;
+    public static T2IRegisteredParam<bool> AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG, SeedVRSplitLatent;
 
-    public static T2IRegisteredParam<double> IPAdapterWeight, IPAdapterStart, IPAdapterEnd, SelfAttentionGuidanceScale, SelfAttentionGuidanceSigmaBlur, PerturbedAttentionGuidanceScale, StyleModelMergeStrength, StyleModelApplyStart, StyleModelMultiplyStrength, RescaleCFGMultiplier, TeaCacheThreshold, TeaCacheStart, NunchakuCacheThreshold, EasyCacheThreshold, EasyCacheStart, EasyCacheEnd, RenormCFG, NormalizedAttentionGuidanceScale, NormalizedAttentionGuidanceAlpha, NormalizedAttentionGuidanceTau;
+    public static T2IRegisteredParam<double> IPAdapterWeight, IPAdapterStart, IPAdapterEnd, SelfAttentionGuidanceScale, SelfAttentionGuidanceSigmaBlur, PerturbedAttentionGuidanceScale, StyleModelMergeStrength, StyleModelApplyStart, StyleModelMultiplyStrength, RescaleCFGMultiplier, TeaCacheThreshold, TeaCacheStart, NunchakuCacheThreshold, EasyCacheThreshold, EasyCacheStart, EasyCacheEnd, RenormCFG, NormalizedAttentionGuidanceScale, NormalizedAttentionGuidanceAlpha, NormalizedAttentionGuidanceTau, SeedVRUpscale;
 
-    public static T2IRegisteredParam<int> RefinerHyperTile, VideoFrameInterpolationMultiplier;
+    public static T2IRegisteredParam<int> RefinerHyperTile, VideoFrameInterpolationMultiplier, SeedVRTemporalVideoOverlap;
 
-    public static T2IRegisteredParam<T2IModel> PixelDecoderModel;
+    public static T2IRegisteredParam<T2IModel> PixelDecoderModel, SeedVRModel;
+
+    public static T2IParamGroup GroupSeedVR;
 
     public static T2IRegisteredParam<string>[] ControlNetPreprocessorParams = new T2IRegisteredParam<string>[3], ControlNetUnionTypeParams = new T2IRegisteredParam<string>[3];
 
@@ -799,6 +802,28 @@ public override void OnInit()
         RefinerHyperTile = T2IParamTypes.Register<int>(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<T2IModel>(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<double>(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<string>(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<string>(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<bool>(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<int>(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<string> interpolators = ["RIFE", "FILM", "GIMM-VFI"];
         VideoPreviewType = T2IParamTypes.Register<string>(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"]

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<ImageToVideoGenInfo> altHandler in AltImageToVideoPreHandlers)
         {
             altHandler(genInfo);
@@ -2670,6 +2674,154 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD
         return result;
     }
 
+    /// <summary>Creates a SeedVR2 restoration stage: converts to a SeedVR2-space latent and samples the restored image from it.</summary>
+    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<T2IModel> 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;
+    }
+
+    /// <summary>Runs the SeedVR group's restoration pass over the current media.</summary>
+    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);
+    }
+
     /// <summary>Creates a "CLIPTextEncode" or equivalent node for the given input, applying prompt-given conditioning modifiers as relevant.</summary>
     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)
     {

src/BuiltinExtensions/ComfyUIBackend/WorkflowGeneratorModelSupport.cs

@@ -94,6 +94,9 @@ public bool IsKontext()
     /// <summary>Returns true if the current model is PiD.</summary>
     public bool IsPiD() => IsModelCompatClass(T2IModelClassSorter.CompatPiD);
 
+    /// <summary>Returns true if the current model is SeedVR2.</summary>
+    public bool IsSeedVR2() => IsModelCompatClass(T2IModelClassSorter.CompatSeedVR2);
+
     /// <summary>Returns true if the current model is HiDream-i1.</summary>
     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))
             {