Adds native SeedVR2 support

docs/Features/Upscaling.md

@@ -2,8 +2,25 @@
 
 (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: <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

@@ -35,10 +35,10 @@ public record class ComfyCustomWorkflow(string Name, string Workflow, string Pro
     public static ConcurrentDictionary<string, ComfyCustomWorkflow> CustomWorkflows = new();
 
     /// <summary>Set of all feature-ids supported by ComfyUI backends.</summary>
-    public static HashSet<string> FeaturesSupported = ["comfyui", "refiners", "controlnet", "endstepsearly", "seamless", "video", "variation_seed", "freeu", "yolov8"];
+    public static HashSet<string> FeaturesSupported = ["comfyui", "refiners", "controlnet", "endstepsearly", "seamless", "video", "variation_seed", "freeu", "yolov8", "seedvr2"];
 
     /// <summary>Set of feature-ids that were added presumptively during loading and should be removed if the backend turns out to be missing them.</summary>
-    public static HashSet<string> FeaturesDiscardIfNotFound = ["variation_seed", "freeu", "yolov8"];
+    public static HashSet<string> FeaturesDiscardIfNotFound = ["variation_seed", "freeu", "yolov8", "seedvr2"];
 
     /// <summary>Extensible map of ComfyUI Node IDs to supported feature IDs.</summary>
     public static Dictionary<string, string> 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",
     };
 
     /// <inheritdoc/>
@@ -606,13 +607,13 @@ 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, FinalUpscaleMethod, SeedVR2ColorCorrectionBehavior, UseIPAdapterForRevision, IPAdapterWeightType, VideoPreviewType, VideoFrameInterpolationMethod, GligenModel, YoloModelInternal, PreferredDType, UseStyleModel, TeaCacheMode, EasyCacheMode, SetClipDevice;
 
-    public static T2IRegisteredParam<bool> AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG;
+    public static T2IRegisteredParam<bool> AITemplateParam, DebugRegionalPrompting, ShiftedLatentAverageInit, UseCfgZeroStar, UseTCFG, SeedVR2SplitLatent;
 
     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<int> RefinerHyperTile, VideoFrameInterpolationMultiplier;
+    public static T2IRegisteredParam<int> RefinerHyperTile, VideoFrameInterpolationMultiplier, SeedVR2TemporalVideoOverlap;
 
     public static T2IRegisteredParam<T2IModel> PixelDecoderModel;
 
@@ -656,6 +657,25 @@ public static T2IModel GetPidModel(string name, Session session)
         return model;
     }
 
+    /// <summary>Lists SeedVR2 upscaler models.</summary>
+    public static List<string> 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}")];
+
+    /// <summary>Resolves a SeedVR2 model from a model name.</summary>
+    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<string> IPAdapterModels = ["None"], IPAdapterWeightTypes = ["standard", "prompt is more important", "style transfer"];
 
     public static List<string> GligenModels = ["None"], YoloModels = [], StyleModels = ["None"], SetClipDevices = ["cpu"];
@@ -770,7 +790,11 @@ public override void OnInit()
             ));
         RefinerUpscaleMethod = T2IParamTypes.Register<string>(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<string>(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<T2IModel>(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,
@@ -799,6 +823,15 @@ 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
             ));
+        SeedVR2ColorCorrectionBehavior = T2IParamTypes.Register<string>(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<bool>(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<int>(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<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

@@ -169,6 +169,9 @@ public JArray FinalImageOut
     /// <summary>If true, the generator is currently working on the pixel-decoder stage.</summary>
     public bool IsPixelDecoderStage = false;
 
+    /// <summary>If true, the generator is currently working on the final stage upscale.</summary>
+    public bool IsFinalStage = false;
+
     /// <summary>If true, the generator is currently working on Image2Video.</summary>
     public bool IsImageToVideo = false;
 
@@ -2670,6 +2673,159 @@ public WGNodeData CreatePixelDecode(T2IModel pidModel, WGNodeData media, WGNodeD
         return result;
     }
 
+    /// <summary>Creates a SeedVR2 restoration stage.</summary>
+    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<T2IModel> 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;
+    }
+
+    /// <summary>Scales raw media to an exact pixel size, or does nothing if it's already that size.</summary>
+    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;
+    }
+
+    /// <summary>Runs the Final Stage upscale over the current media, if the user configured one and it hasn't run yet.</summary>
+    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;
+    }
+
+    /// <summary>Upscales raw media by any of the Final Upscale Methods.</summary>
+    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.");
+    }
+
     /// <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)
     {