convert nodes_lt_upsampler nodes to V3 schema

comfy_extras/nodes_lt_upsampler.py

@@ -1,32 +1,32 @@
 from comfy import model_management
+from comfy_api.latest import ComfyExtension, IO
+from typing_extensions import override
 import math
 
-class LTXVLatentUpsampler:
+
+class LTXVLatentUpsampler(IO.ComfyNode):
     """
     Upsamples a video latent by a factor of 2.
     """
 
     @classmethod
-    def INPUT_TYPES(s):
+    def define_schema(cls):
-        return {
+        return IO.Schema(
-            "required": {
+            node_id="LTXVLatentUpsampler",
-                "samples": ("LATENT",),
+            category="latent/video",
-                "upscale_model": ("LATENT_UPSCALE_MODEL",),
+            is_experimental=True,
-                "vae": ("VAE",),
+            inputs=[
-            }
+                IO.Latent.Input("samples"),
-        }
+                IO.LatentUpscaleModel.Input("upscale_model"),
+                IO.Vae.Input("vae"),
+            ],
+            outputs=[
+                IO.Latent.Output(),
+            ],
+        )
 
-    RETURN_TYPES = ("LATENT",)
+    @classmethod
-    FUNCTION = "upsample_latent"
+    def execute(cls, samples, upscale_model, vae) -> IO.NodeOutput:
-    CATEGORY = "latent/video"
-    EXPERIMENTAL = True
-
-    def upsample_latent(
-        self,
-        samples: dict,
-        upscale_model,
-        vae,
-    ) -> tuple:
         """
         Upsample the input latent using the provided model.
 
@@ -34,7 +34,6 @@ class LTXVLatentUpsampler:
             samples (dict): Input latent samples
             upscale_model (LatentUpsampler): Loaded upscale model
             vae: VAE model for normalization
-            auto_tiling (bool): Whether to automatically tile the input for processing
 
         Returns:
             tuple: Tuple containing the upsampled latent
@@ -67,9 +66,16 @@ class LTXVLatentUpsampler:
         return_dict = samples.copy()
         return_dict["samples"] = upsampled_latents
         return_dict.pop("noise_mask", None)
-        return (return_dict,)
+        return IO.NodeOutput(return_dict)
+
+    upsample_latent = execute  # TODO: remove
 
 
-NODE_CLASS_MAPPINGS = {
+class LTXVLatentUpsamplerExtension(ComfyExtension):
-    "LTXVLatentUpsampler": LTXVLatentUpsampler,
+    @override
-}
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [LTXVLatentUpsampler]
+
+
+async def comfy_entrypoint() -> LTXVLatentUpsamplerExtension:
+    return LTXVLatentUpsamplerExtension()