nodes_custom: Lazify the model

These little model reader bits need to lazily evaluate the model to
allow the model to be freed incase there is RAM pressure.

Author: rattus128

comfy_extras/nodes_custom_sampler.py

@@ -18,14 +18,17 @@ def define_schema(cls):
             node_id="BasicScheduler",
             category="sampling/custom_sampling/schedulers",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Combo.Input("scheduler", options=comfy.samplers.SCHEDULER_NAMES),
                 io.Int.Input("steps", default=20, min=1, max=10000),
                 io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01),
             ],
             outputs=[io.Sigmas.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, scheduler, steps, denoise) -> io.NodeOutput:
         total_steps = steps
@@ -137,13 +140,16 @@ def define_schema(cls):
             node_id="SDTurboScheduler",
             category="sampling/custom_sampling/schedulers",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Int.Input("steps", default=1, min=1, max=10),
                 io.Float.Input("denoise", default=1.0, min=0, max=1.0, step=0.01),
             ],
             outputs=[io.Sigmas.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, steps, denoise) -> io.NodeOutput:
         start_step = 10 - int(10 * denoise)
@@ -161,14 +167,17 @@ def define_schema(cls):
             node_id="BetaSamplingScheduler",
             category="sampling/custom_sampling/schedulers",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Int.Input("steps", default=20, min=1, max=10000),
                 io.Float.Input("alpha", default=0.6, min=0.0, max=50.0, step=0.01, round=False),
                 io.Float.Input("beta", default=0.6, min=0.0, max=50.0, step=0.01, round=False),
             ],
             outputs=[io.Sigmas.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, steps, alpha, beta) -> io.NodeOutput:
         sigmas = comfy.samplers.beta_scheduler(model.get_model_object("model_sampling"), steps, alpha=alpha, beta=beta)
@@ -351,13 +360,16 @@ def define_schema(cls):
             node_id="SamplingPercentToSigma",
             category="sampling/custom_sampling/sigmas",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Float.Input("sampling_percent", default=0.0, min=0.0, max=1.0, step=0.0001),
                 io.Boolean.Input("return_actual_sigma", default=False, tooltip="Return the actual sigma value instead of the value used for interval checks.\nThis only affects results at 0.0 and 1.0."),
             ],
             outputs=[io.Float.Output(display_name="sigma_value")]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, sampling_percent, return_actual_sigma) -> io.NodeOutput:
         model_sampling = model.get_model_object("model_sampling")
@@ -622,7 +634,7 @@ def define_schema(cls):
             node_id="SamplerSASolver",
             category="sampling/custom_sampling/samplers",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Float.Input("eta", default=1.0, min=0.0, max=10.0, step=0.01, round=False),
                 io.Float.Input("sde_start_percent", default=0.2, min=0.0, max=1.0, step=0.001),
                 io.Float.Input("sde_end_percent", default=0.8, min=0.0, max=1.0, step=0.001),
@@ -635,6 +647,9 @@ def define_schema(cls):
             outputs=[io.Sampler.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, eta, sde_start_percent, sde_end_percent, s_noise, predictor_order, corrector_order, use_pece, simple_order_2) -> io.NodeOutput:
         model_sampling = model.get_model_object("model_sampling")
@@ -684,7 +699,7 @@ def define_schema(cls):
             node_id="SamplerCustom",
             category="sampling/custom_sampling",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Boolean.Input("add_noise", default=True),
                 io.Int.Input("noise_seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
                 io.Float.Input("cfg", default=8.0, min=0.0, max=100.0, step=0.1, round=0.01),
@@ -700,6 +715,9 @@ def define_schema(cls):
             ]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image) -> io.NodeOutput:
         latent = latent_image
@@ -745,12 +763,15 @@ def define_schema(cls):
             node_id="BasicGuider",
             category="sampling/custom_sampling/guiders",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Conditioning.Input("conditioning"),
             ],
             outputs=[io.Guider.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, conditioning) -> io.NodeOutput:
         guider = Guider_Basic(model)
@@ -766,14 +787,17 @@ def define_schema(cls):
             node_id="CFGGuider",
             category="sampling/custom_sampling/guiders",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Conditioning.Input("positive"),
                 io.Conditioning.Input("negative"),
                 io.Float.Input("cfg", default=8.0, min=0.0, max=100.0, step=0.1, round=0.01),
             ],
             outputs=[io.Guider.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, positive, negative, cfg) -> io.NodeOutput:
         guider = comfy.samplers.CFGGuider(model)
@@ -819,7 +843,7 @@ def define_schema(cls):
             node_id="DualCFGGuider",
             category="sampling/custom_sampling/guiders",
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Conditioning.Input("cond1"),
                 io.Conditioning.Input("cond2"),
                 io.Conditioning.Input("negative"),
@@ -830,6 +854,9 @@ def define_schema(cls):
             outputs=[io.Guider.Output()]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, cond1, cond2, negative, cfg_conds, cfg_cond2_negative, style) -> io.NodeOutput:
         guider = Guider_DualCFG(model)
@@ -930,7 +957,7 @@ def define_schema(cls):
             category="_for_testing/custom_sampling/noise",
             is_experimental=True,
             inputs=[
-                io.Model.Input("model"),
+                io.Model.Input("model", lazy=True),
                 io.Noise.Input("noise"),
                 io.Sigmas.Input("sigmas"),
                 io.Latent.Input("latent_image"),
@@ -940,6 +967,9 @@ def define_schema(cls):
             ]
         )
 
+    def check_lazy_status(self, *args, **kwargs):
+        return ["model"]
+
     @classmethod
     def execute(cls, model, noise, sigmas, latent_image) -> io.NodeOutput:
         if len(sigmas) == 0: