ComfyUI 0.3.57

* convert RunAway API nodes to the V3 schema

* fixed small typo

* fix: add tooltip for "seed" input

comfy/cli_args.py

@@ -143,6 +143,7 @@ class PerformanceFeature(enum.Enum):
     Fp16Accumulation = "fp16_accumulation"
     Fp8MatrixMultiplication = "fp8_matrix_mult"
     CublasOps = "cublas_ops"
+    AutoTune = "autotune"
 
 parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: fp16_accumulation fp8_matrix_mult cublas_ops")
 

comfy/clip_model.py

@@ -61,8 +61,12 @@ class CLIPEncoder(torch.nn.Module):
     def forward(self, x, mask=None, intermediate_output=None):
         optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
 
+        all_intermediate = None
         if intermediate_output is not None:
-            if intermediate_output < 0:
+            if intermediate_output == "all":
+                all_intermediate = []
+                intermediate_output = None
+            elif intermediate_output < 0:
                 intermediate_output = len(self.layers) + intermediate_output
 
         intermediate = None
@@ -70,6 +74,12 @@ class CLIPEncoder(torch.nn.Module):
             x = l(x, mask, optimized_attention)
             if i == intermediate_output:
                 intermediate = x.clone()
+            if all_intermediate is not None:
+                all_intermediate.append(x.unsqueeze(1).clone())
+
+        if all_intermediate is not None:
+            intermediate = torch.cat(all_intermediate, dim=1)
+
         return x, intermediate
 
 class CLIPEmbeddings(torch.nn.Module):

comfy/clip_vision.py

@@ -50,7 +50,13 @@ class ClipVisionModel():
         self.image_size = config.get("image_size", 224)
         self.image_mean = config.get("image_mean", [0.48145466, 0.4578275, 0.40821073])
         self.image_std = config.get("image_std", [0.26862954, 0.26130258, 0.27577711])
-        model_class = IMAGE_ENCODERS.get(config.get("model_type", "clip_vision_model"))
+        model_type = config.get("model_type", "clip_vision_model")
+        model_class = IMAGE_ENCODERS.get(model_type)
+        if model_type == "siglip_vision_model":
+            self.return_all_hidden_states = True
+        else:
+            self.return_all_hidden_states = False
+
         self.load_device = comfy.model_management.text_encoder_device()
         offload_device = comfy.model_management.text_encoder_offload_device()
         self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
@@ -68,12 +74,18 @@ class ClipVisionModel():
     def encode_image(self, image, crop=True):
         comfy.model_management.load_model_gpu(self.patcher)
         pixel_values = clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=crop).float()
-        out = self.model(pixel_values=pixel_values, intermediate_output=-2)
+        out = self.model(pixel_values=pixel_values, intermediate_output='all' if self.return_all_hidden_states else -2)
 
         outputs = Output()
         outputs["last_hidden_state"] = out[0].to(comfy.model_management.intermediate_device())
         outputs["image_embeds"] = out[2].to(comfy.model_management.intermediate_device())
-        outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
+        if self.return_all_hidden_states:
+            all_hs = out[1].to(comfy.model_management.intermediate_device())
+            outputs["penultimate_hidden_states"] = all_hs[:, -2]
+            outputs["all_hidden_states"] = all_hs
+        else:
+            outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
+
         outputs["mm_projected"] = out[3]
         return outputs
 

comfy/k_diffusion/sampling.py

@@ -171,6 +171,16 @@ def offset_first_sigma_for_snr(sigmas, model_sampling, percent_offset=1e-4):
     return sigmas
 
 
+def ei_h_phi_1(h: torch.Tensor) -> torch.Tensor:
+    """Compute the result of h*phi_1(h) in exponential integrator methods."""
+    return torch.expm1(h)
+
+
+def ei_h_phi_2(h: torch.Tensor) -> torch.Tensor:
+    """Compute the result of h*phi_2(h) in exponential integrator methods."""
+    return (torch.expm1(h) - h) / h
+
+
 @torch.no_grad()
 def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
     """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
@@ -1550,13 +1560,12 @@ def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None
 @torch.no_grad()
 def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
     """SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
-    arXiv: https://arxiv.org/abs/2305.14267
+    arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
     """
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
     noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
-
     inject_noise = eta > 0 and s_noise > 0
 
     model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
@@ -1564,55 +1573,53 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
     lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
     sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
 
+    fac = 1 / (2 * r)
+
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised = model(x, sigmas[i] * s_in, **extra_args)
         if callback is not None:
             callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
         if sigmas[i + 1] == 0:
             x = denoised
-        else:
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            h_eta = h * (eta + 1)
-            lambda_s_1 = lambda_s + r * h
-            fac = 1 / (2 * r)
-            sigma_s_1 = sigma_fn(lambda_s_1)
+            continue
 
-            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
-            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
+        lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+        h = lambda_t - lambda_s
+        h_eta = h * (eta + 1)
+        lambda_s_1 = torch.lerp(lambda_s, lambda_t, r)
+        sigma_s_1 = sigma_fn(lambda_s_1)
 
-            coeff_1, coeff_2 = (-r * h_eta).expm1(), (-h_eta).expm1()
-            if inject_noise:
-                # 0 < r < 1
-                noise_coeff_1 = (-2 * r * h * eta).expm1().neg().sqrt()
-                noise_coeff_2 = (-r * h * eta).exp() * (-2 * (1 - r) * h * eta).expm1().neg().sqrt()
-                noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigmas[i + 1])
+        alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+        alpha_t = sigmas[i + 1] * lambda_t.exp()
 
-            # Step 1
-            x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
-            if inject_noise:
-                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
-            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+        # Step 1
+        x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * ei_h_phi_1(-r * h_eta) * denoised
+        if inject_noise:
+            sde_noise = (-2 * r * h * eta).expm1().neg().sqrt() * noise_sampler(sigmas[i], sigma_s_1)
+            x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
+        denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
-            # Step 2
-            denoised_d = (1 - fac) * denoised + fac * denoised_2
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_2 * denoised_d
-            if inject_noise:
-                x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
+        # Step 2
+        denoised_d = torch.lerp(denoised, denoised_2, fac)
+        x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
+        if inject_noise:
+            segment_factor = (r - 1) * h * eta
+            sde_noise = sde_noise * segment_factor.exp()
+            sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_1, sigmas[i + 1])
+            x = x + sde_noise * sigmas[i + 1] * s_noise
     return x
 
 
 @torch.no_grad()
 def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
     """SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 3.
-    arXiv: https://arxiv.org/abs/2305.14267
+    arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
     """
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
     noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
-
     inject_noise = eta > 0 and s_noise > 0
 
     model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
@@ -1624,45 +1631,49 @@ def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=Non
         denoised = model(x, sigmas[i] * s_in, **extra_args)
         if callback is not None:
             callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
         if sigmas[i + 1] == 0:
             x = denoised
-        else:
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            h_eta = h * (eta + 1)
-            lambda_s_1 = lambda_s + r_1 * h
-            lambda_s_2 = lambda_s + r_2 * h
-            sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
+            continue
 
-            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
-            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-            alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
+        lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+        h = lambda_t - lambda_s
+        h_eta = h * (eta + 1)
+        lambda_s_1 = torch.lerp(lambda_s, lambda_t, r_1)
+        lambda_s_2 = torch.lerp(lambda_s, lambda_t, r_2)
+        sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
 
-            coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
-            if inject_noise:
-                # 0 < r_1 < r_2 < 1
-                noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
-                noise_coeff_2 = (-r_1 * h * eta).exp() * (-2 * (r_2 - r_1) * h * eta).expm1().neg().sqrt()
-                noise_coeff_3 = (-r_2 * h * eta).exp() * (-2 * (1 - r_2) * h * eta).expm1().neg().sqrt()
-                noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
+        alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+        alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
+        alpha_t = sigmas[i + 1] * lambda_t.exp()
 
-            # Step 1
-            x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
-            if inject_noise:
-                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
-            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+        # Step 1
+        x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * ei_h_phi_1(-r_1 * h_eta) * denoised
+        if inject_noise:
+            sde_noise = (-2 * r_1 * h * eta).expm1().neg().sqrt() * noise_sampler(sigmas[i], sigma_s_1)
+            x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
+        denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
-            # Step 2
-            x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * coeff_2 * denoised + (r_2 / r_1) * alpha_s_2 * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
-            if inject_noise:
-                x_3 = x_3 + sigma_s_2 * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
-            denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
+        # Step 2
+        a3_2 = r_2 / r_1 * ei_h_phi_2(-r_2 * h_eta)
+        a3_1 = ei_h_phi_1(-r_2 * h_eta) - a3_2
+        x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * (a3_1 * denoised + a3_2 * denoised_2)
+        if inject_noise:
+            segment_factor = (r_1 - r_2) * h * eta
+            sde_noise = sde_noise * segment_factor.exp()
+            sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_1, sigma_s_2)
+            x_3 = x_3 + sde_noise * sigma_s_2 * s_noise
+        denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
 
-            # Step 3
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_3 * denoised + (1. / r_2) * alpha_t * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
-            if inject_noise:
-                x = x + sigmas[i + 1] * (noise_coeff_3 * noise_1 + noise_coeff_2 * noise_2 + noise_coeff_1 * noise_3) * s_noise
+        # Step 3
+        b3 = ei_h_phi_2(-h_eta) / r_2
+        b1 = ei_h_phi_1(-h_eta) - b3
+        x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * (b1 * denoised + b3 * denoised_3)
+        if inject_noise:
+            segment_factor = (r_2 - 1) * h * eta
+            sde_noise = sde_noise * segment_factor.exp()
+            sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_2, sigmas[i + 1])
+            x = x + sde_noise * sigmas[i + 1] * s_noise
     return x
 
 

comfy/ldm/audio/dit.py

@@ -632,7 +632,7 @@ class ContinuousTransformer(nn.Module):
         # Attention layers
 
         if self.rotary_pos_emb is not None:
-            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=x.dtype, device=x.device)
+            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=torch.float, device=x.device)
         else:
             rotary_pos_emb = None
 

comfy/ldm/flux/model.py

@@ -106,6 +106,7 @@ class Flux(nn.Module):
         if y is None:
             y = torch.zeros((img.shape[0], self.params.vec_in_dim), device=img.device, dtype=img.dtype)
 
+        patches = transformer_options.get("patches", {})
         patches_replace = transformer_options.get("patches_replace", {})
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -117,9 +118,17 @@ class Flux(nn.Module):
             if guidance is not None:
                 vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
 
-        vec = vec + self.vector_in(y[:,:self.params.vec_in_dim])
+        vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
         txt = self.txt_in(txt)
 
+        if "post_input" in patches:
+            for p in patches["post_input"]:
+                out = p({"img": img, "txt": txt, "img_ids": img_ids, "txt_ids": txt_ids})
+                img = out["img"]
+                txt = out["txt"]
+                img_ids = out["img_ids"]
+                txt_ids = out["txt_ids"]
+
         if img_ids is not None:
             ids = torch.cat((txt_ids, img_ids), dim=1)
             pe = self.pe_embedder(ids)
@@ -233,12 +242,18 @@ class Flux(nn.Module):
             h = 0
             w = 0
             index = 0
-            index_ref_method = kwargs.get("ref_latents_method", "offset") == "index"
+            ref_latents_method = kwargs.get("ref_latents_method", "offset")
             for ref in ref_latents:
-                if index_ref_method:
+                if ref_latents_method == "index":
                     index += 1
                     h_offset = 0
                     w_offset = 0
+                elif ref_latents_method == "uxo":
+                    index = 0
+                    h_offset = h_len * patch_size + h
+                    w_offset = w_len * patch_size + w
+                    h += ref.shape[-2]
+                    w += ref.shape[-1]
                 else:
                     index = 1
                     h_offset = 0

comfy/lora.py

@@ -260,6 +260,10 @@ def model_lora_keys_unet(model, key_map={}):
                 key_map["transformer.{}".format(k[:-len(".weight")])] = to #simpletrainer and probably regular diffusers flux lora format
                 key_map["lycoris_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #simpletrainer lycoris
                 key_map["lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #onetrainer
+        for k in sdk:
+            hidden_size = model.model_config.unet_config.get("hidden_size", 0)
+            if k.endswith(".weight") and ".linear1." in k:
+                key_map["{}".format(k.replace(".linear1.weight", ".linear1_qkv"))] = (k, (0, 0, hidden_size * 3))
 
     if isinstance(model, comfy.model_base.GenmoMochi):
         for k in sdk:

comfy/lora_convert.py

@@ -15,10 +15,29 @@ def convert_lora_bfl_control(sd): #BFL loras for Flux
 def convert_lora_wan_fun(sd): #Wan Fun loras
     return comfy.utils.state_dict_prefix_replace(sd, {"lora_unet__": "lora_unet_"})
 
+def convert_uso_lora(sd):
+    sd_out = {}
+    for k in sd:
+        tensor = sd[k]
+        k_to = "diffusion_model.{}".format(k.replace(".down.weight", ".lora_down.weight")
+                                           .replace(".up.weight", ".lora_up.weight")
+                                           .replace(".qkv_lora2.", ".txt_attn.qkv.")
+                                           .replace(".qkv_lora1.", ".img_attn.qkv.")
+                                           .replace(".proj_lora1.", ".img_attn.proj.")
+                                           .replace(".proj_lora2.", ".txt_attn.proj.")
+                                           .replace(".qkv_lora.", ".linear1_qkv.")
+                                           .replace(".proj_lora.", ".linear2.")
+                                           .replace(".processor.", ".")
+                                           )
+        sd_out[k_to] = tensor
+    return sd_out
+
 
 def convert_lora(sd):
     if "img_in.lora_A.weight" in sd and "single_blocks.0.norm.key_norm.scale" in sd:
         return convert_lora_bfl_control(sd)
     if "lora_unet__blocks_0_cross_attn_k.lora_down.weight" in sd:
         return convert_lora_wan_fun(sd)
+    if "single_blocks.37.processor.qkv_lora.up.weight" in sd and "double_blocks.18.processor.qkv_lora2.up.weight" in sd:
+        return convert_uso_lora(sd)
     return sd

comfy/model_patcher.py

@@ -433,6 +433,9 @@ class ModelPatcher:
     def set_model_double_block_patch(self, patch):
         self.set_model_patch(patch, "double_block")
 
+    def set_model_post_input_patch(self, patch):
+        self.set_model_patch(patch, "post_input")
+
     def add_object_patch(self, name, obj):
         self.object_patches[name] = obj
 

comfy/ops.py

@@ -52,6 +52,9 @@ except (ModuleNotFoundError, TypeError):
 
 cast_to = comfy.model_management.cast_to #TODO: remove once no more references
 
+if torch.cuda.is_available() and torch.backends.cudnn.is_available() and PerformanceFeature.AutoTune in args.fast:
+    torch.backends.cudnn.benchmark = True
+
 def cast_to_input(weight, input, non_blocking=False, copy=True):
     return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
 

comfy_api_nodes/apis/__init__.py

@@ -951,7 +951,11 @@ class MagicPrompt2(str, Enum):
 
 
 class StyleType1(str, Enum):
+    AUTO = 'AUTO'
     GENERAL = 'GENERAL'
+    REALISTIC = 'REALISTIC'
+    DESIGN = 'DESIGN'
+    FICTION = 'FICTION'
 
 
 class ImagenImageGenerationInstance(BaseModel):
@@ -2676,7 +2680,7 @@ class ReleaseNote(BaseModel):
 
 
 class RenderingSpeed(str, Enum):
-    BALANCED = 'BALANCED'
+    DEFAULT = 'DEFAULT'
     TURBO = 'TURBO'
     QUALITY = 'QUALITY'
 
@@ -4918,6 +4922,14 @@ class IdeogramV3EditRequest(BaseModel):
         None,
         description='A set of images to use as style references (maximum total size 10MB across all style references). The images should be in JPEG, PNG or WebP format.',
     )
+    character_reference_images: Optional[List[str]] = Field(
+        None,
+        description='Generations with character reference are subject to the character reference pricing. A set of images to use as character references (maximum total size 10MB across all character references), currently only supports 1 character reference image. The images should be in JPEG, PNG or WebP format.'
+    )
+    character_reference_images_mask: Optional[List[str]] = Field(
+        None,
+        description='Optional masks for character reference images. When provided, must match the number of character_reference_images. Each mask should be a grayscale image of the same dimensions as the corresponding character reference image. The images should be in JPEG, PNG or WebP format.'
+    )
 
 
 class IdeogramV3Request(BaseModel):
@@ -4951,6 +4963,14 @@ class IdeogramV3Request(BaseModel):
     style_type: Optional[StyleType1] = Field(
         None, description='The type of style to apply'
     )
+    character_reference_images: Optional[List[str]] = Field(
+        None,
+        description='Generations with character reference are subject to the character reference pricing. A set of images to use as character references (maximum total size 10MB across all character references), currently only supports 1 character reference image. The images should be in JPEG, PNG or WebP format.'
+    )
+    character_reference_images_mask: Optional[List[str]] = Field(
+        None,
+        description='Optional masks for character reference images. When provided, must match the number of character_reference_images. Each mask should be a grayscale image of the same dimensions as the corresponding character reference image. The images should be in JPEG, PNG or WebP format.'
+    )
 
 
 class ImagenGenerateImageResponse(BaseModel):

comfy_api_nodes/nodes_bytedance.py

@@ -0,0 +1,336 @@
+import logging
+from enum import Enum
+from typing import Optional
+from typing_extensions import override
+
+import torch
+from pydantic import BaseModel, Field
+
+from comfy_api.latest import ComfyExtension, io as comfy_io
+from comfy_api_nodes.util.validation_utils import (
+    validate_image_aspect_ratio_range,
+    get_number_of_images,
+)
+from comfy_api_nodes.apis.client import (
+    ApiEndpoint,
+    HttpMethod,
+    SynchronousOperation,
+)
+from comfy_api_nodes.apinode_utils import download_url_to_image_tensor, upload_images_to_comfyapi, validate_string
+
+
+BYTEPLUS_ENDPOINT = "/proxy/byteplus/api/v3/images/generations"
+
+
+class Text2ImageModelName(str, Enum):
+    seedream3 = "seedream-3-0-t2i-250415"
+
+
+class Image2ImageModelName(str, Enum):
+    seededit3 = "seededit-3-0-i2i-250628"
+
+
+class Text2ImageTaskCreationRequest(BaseModel):
+    model: Text2ImageModelName = Text2ImageModelName.seedream3
+    prompt: str = Field(...)
+    response_format: Optional[str] = Field("url")
+    size: Optional[str] = Field(None)
+    seed: Optional[int] = Field(0, ge=0, le=2147483647)
+    guidance_scale: Optional[float] = Field(..., ge=1.0, le=10.0)
+    watermark: Optional[bool] = Field(True)
+
+
+class Image2ImageTaskCreationRequest(BaseModel):
+    model: Image2ImageModelName = Image2ImageModelName.seededit3
+    prompt: str = Field(...)
+    response_format: Optional[str] = Field("url")
+    image: str = Field(..., description="Base64 encoded string or image URL")
+    size: Optional[str] = Field("adaptive")
+    seed: Optional[int] = Field(..., ge=0, le=2147483647)
+    guidance_scale: Optional[float] = Field(..., ge=1.0, le=10.0)
+    watermark: Optional[bool] = Field(True)
+
+
+class ImageTaskCreationResponse(BaseModel):
+    model: str = Field(...)
+    created: int = Field(..., description="Unix timestamp (in seconds) indicating time when the request was created.")
+    data: list = Field([], description="Contains information about the generated image(s).")
+    error: dict = Field({}, description="Contains `code` and `message` fields in case of error.")
+
+
+RECOMMENDED_PRESETS = [
+    ("1024x1024 (1:1)", 1024, 1024),
+    ("864x1152 (3:4)", 864, 1152),
+    ("1152x864 (4:3)", 1152, 864),
+    ("1280x720 (16:9)", 1280, 720),
+    ("720x1280 (9:16)", 720, 1280),
+    ("832x1248 (2:3)", 832, 1248),
+    ("1248x832 (3:2)", 1248, 832),
+    ("1512x648 (21:9)", 1512, 648),
+    ("2048x2048 (1:1)", 2048, 2048),
+    ("Custom", None, None),
+]
+
+
+def get_image_url_from_response(response: ImageTaskCreationResponse) -> str:
+    if response.error:
+        error_msg = f"ByteDance request failed. Code: {response.error['code']}, message: {response.error['message']}"
+        logging.info(error_msg)
+        raise RuntimeError(error_msg)
+    logging.info("ByteDance task succeeded, image URL: %s", response.data[0]["url"])
+    return response.data[0]["url"]
+
+
+class ByteDanceImageNode(comfy_io.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="ByteDanceImageNode",
+            display_name="ByteDance Image",
+            category="api node/image/ByteDance",
+            description="Generate images using ByteDance models via api based on prompt",
+            inputs=[
+                comfy_io.Combo.Input(
+                    "model",
+                    options=[model.value for model in Text2ImageModelName],
+                    default=Text2ImageModelName.seedream3.value,
+                    tooltip="Model name",
+                ),
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    tooltip="The text prompt used to generate the image",
+                ),
+                comfy_io.Combo.Input(
+                    "size_preset",
+                    options=[label for label, _, _ in RECOMMENDED_PRESETS],
+                    tooltip="Pick a recommended size. Select Custom to use the width and height below",
+                ),
+                comfy_io.Int.Input(
+                    "width",
+                    default=1024,
+                    min=512,
+                    max=2048,
+                    step=64,
+                    tooltip="Custom width for image. Value is working only if `size_preset` is set to `Custom`",
+                ),
+                comfy_io.Int.Input(
+                    "height",
+                    default=1024,
+                    min=512,
+                    max=2048,
+                    step=64,
+                    tooltip="Custom height for image. Value is working only if `size_preset` is set to `Custom`",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation",
+                    optional=True,
+                ),
+                comfy_io.Float.Input(
+                    "guidance_scale",
+                    default=2.5,
+                    min=1.0,
+                    max=10.0,
+                    step=0.01,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Higher value makes the image follow the prompt more closely",
+                    optional=True,
+                ),
+                comfy_io.Boolean.Input(
+                    "watermark",
+                    default=True,
+                    tooltip="Whether to add an \"AI generated\" watermark to the image",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: str,
+        prompt: str,
+        size_preset: str,
+        width: int,
+        height: int,
+        seed: int,
+        guidance_scale: float,
+        watermark: bool,
+    ) -> comfy_io.NodeOutput:
+        validate_string(prompt, strip_whitespace=True, min_length=1)
+        w = h = None
+        for label, tw, th in RECOMMENDED_PRESETS:
+            if label == size_preset:
+                w, h = tw, th
+                break
+
+        if w is None or h is None:
+            w, h = width, height
+            if not (512 <= w <= 2048) or not (512 <= h <= 2048):
+                raise ValueError(
+                    f"Custom size out of range: {w}x{h}. "
+                    "Both width and height must be between 512 and 2048 pixels."
+                )
+
+        payload = Text2ImageTaskCreationRequest(
+            model=model,
+            prompt=prompt,
+            size=f"{w}x{h}",
+            seed=seed,
+            guidance_scale=guidance_scale,
+            watermark=watermark,
+        )
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+        response = await SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=BYTEPLUS_ENDPOINT,
+                method=HttpMethod.POST,
+                request_model=Text2ImageTaskCreationRequest,
+                response_model=ImageTaskCreationResponse,
+            ),
+            request=payload,
+            auth_kwargs=auth_kwargs,
+        ).execute()
+        return comfy_io.NodeOutput(await download_url_to_image_tensor(get_image_url_from_response(response)))
+
+
+class ByteDanceImageEditNode(comfy_io.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="ByteDanceImageEditNode",
+            display_name="ByteDance Image Edit",
+            category="api node/video/ByteDance",
+            description="Edit images using ByteDance models via api based on prompt",
+            inputs=[
+                comfy_io.Combo.Input(
+                    "model",
+                    options=[model.value for model in Image2ImageModelName],
+                    default=Image2ImageModelName.seededit3.value,
+                    tooltip="Model name",
+                ),
+                comfy_io.Image.Input(
+                    "image",
+                    tooltip="The base image to edit",
+                ),
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Instruction to edit image",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation",
+                    optional=True,
+                ),
+                comfy_io.Float.Input(
+                    "guidance_scale",
+                    default=5.5,
+                    min=1.0,
+                    max=10.0,
+                    step=0.01,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Higher value makes the image follow the prompt more closely",
+                    optional=True,
+                ),
+                comfy_io.Boolean.Input(
+                    "watermark",
+                    default=True,
+                    tooltip="Whether to add an \"AI generated\" watermark to the image",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: str,
+        image: torch.Tensor,
+        prompt: str,
+        seed: int,
+        guidance_scale: float,
+        watermark: bool,
+    ) -> comfy_io.NodeOutput:
+        validate_string(prompt, strip_whitespace=True, min_length=1)
+        if get_number_of_images(image) != 1:
+            raise ValueError("Exactly one input image is required.")
+        validate_image_aspect_ratio_range(image, (1, 3), (3, 1))
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+        source_url = (await upload_images_to_comfyapi(
+            image,
+            max_images=1,
+            mime_type="image/png",
+            auth_kwargs=auth_kwargs,
+        ))[0]
+        payload = Image2ImageTaskCreationRequest(
+            model=model,
+            prompt=prompt,
+            image=source_url,
+            seed=seed,
+            guidance_scale=guidance_scale,
+            watermark=watermark,
+        )
+        response = await SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=BYTEPLUS_ENDPOINT,
+                method=HttpMethod.POST,
+                request_model=Image2ImageTaskCreationRequest,
+                response_model=ImageTaskCreationResponse,
+            ),
+            request=payload,
+            auth_kwargs=auth_kwargs,
+        ).execute()
+        return comfy_io.NodeOutput(await download_url_to_image_tensor(get_image_url_from_response(response)))
+
+
+class ByteDanceExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
+        return [
+            ByteDanceImageNode,
+            ByteDanceImageEditNode,
+        ]
+
+async def comfy_entrypoint() -> ByteDanceExtension:
+    return ByteDanceExtension()

comfy_api_nodes/nodes_ideogram.py

@@ -255,6 +255,7 @@ class IdeogramV1(comfy_io.ComfyNode):
             display_name="Ideogram V1",
             category="api node/image/Ideogram",
             description="Generates images using the Ideogram V1 model.",
+            is_api_node=True,
             inputs=[
                 comfy_io.String.Input(
                     "prompt",
@@ -383,6 +384,7 @@ class IdeogramV2(comfy_io.ComfyNode):
             display_name="Ideogram V2",
             category="api node/image/Ideogram",
             description="Generates images using the Ideogram V2 model.",
+            is_api_node=True,
             inputs=[
                 comfy_io.String.Input(
                     "prompt",
@@ -552,6 +554,7 @@ class IdeogramV3(comfy_io.ComfyNode):
             category="api node/image/Ideogram",
             description="Generates images using the Ideogram V3 model. "
                         "Supports both regular image generation from text prompts and image editing with mask.",
+            is_api_node=True,
             inputs=[
                 comfy_io.String.Input(
                     "prompt",
@@ -612,11 +615,21 @@ class IdeogramV3(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "rendering_speed",
-                    options=["BALANCED", "TURBO", "QUALITY"],
-                    default="BALANCED",
+                    options=["DEFAULT", "TURBO", "QUALITY"],
+                    default="DEFAULT",
                     tooltip="Controls the trade-off between generation speed and quality",
                     optional=True,
                 ),
+                comfy_io.Image.Input(
+                    "character_image",
+                    tooltip="Image to use as character reference.",
+                    optional=True,
+                ),
+                comfy_io.Mask.Input(
+                    "character_mask",
+                    tooltip="Optional mask for character reference image.",
+                    optional=True,
+                ),
             ],
             outputs=[
                 comfy_io.Image.Output(),
@@ -639,12 +652,46 @@ class IdeogramV3(comfy_io.ComfyNode):
         magic_prompt_option="AUTO",
         seed=0,
         num_images=1,
-        rendering_speed="BALANCED",
+        rendering_speed="DEFAULT",
+        character_image=None,
+        character_mask=None,
     ):
         auth = {
             "auth_token": cls.hidden.auth_token_comfy_org,
             "comfy_api_key": cls.hidden.api_key_comfy_org,
         }
+        if rendering_speed == "BALANCED":  # for backward compatibility
+            rendering_speed = "DEFAULT"
+
+        character_img_binary = None
+        character_mask_binary = None
+
+        if character_image is not None:
+            input_tensor = character_image.squeeze().cpu()
+            if character_mask is not None:
+                character_mask = resize_mask_to_image(character_mask, character_image, allow_gradient=False)
+                character_mask = 1.0 - character_mask
+                if character_mask.shape[1:] != character_image.shape[1:-1]:
+                    raise Exception("Character mask and image must be the same size")
+
+                mask_np = (character_mask.squeeze().cpu().numpy() * 255).astype(np.uint8)
+                mask_img = Image.fromarray(mask_np)
+                mask_byte_arr = BytesIO()
+                mask_img.save(mask_byte_arr, format="PNG")
+                mask_byte_arr.seek(0)
+                character_mask_binary = mask_byte_arr
+                character_mask_binary.name = "mask.png"
+
+            img_np = (input_tensor.numpy() * 255).astype(np.uint8)
+            img = Image.fromarray(img_np)
+            img_byte_arr = BytesIO()
+            img.save(img_byte_arr, format="PNG")
+            img_byte_arr.seek(0)
+            character_img_binary = img_byte_arr
+            character_img_binary.name = "image.png"
+        elif character_mask is not None:
+            raise Exception("Character mask requires character image to be present")
+
         # Check if both image and mask are provided for editing mode
         if image is not None and mask is not None:
             # Edit mode
@@ -693,6 +740,15 @@ class IdeogramV3(comfy_io.ComfyNode):
             if num_images > 1:
                 edit_request.num_images = num_images
 
+            files = {
+                "image": img_binary,
+                "mask": mask_binary,
+            }
+            if character_img_binary:
+                files["character_reference_images"] = character_img_binary
+            if character_mask_binary:
+                files["character_mask_binary"] = character_mask_binary
+
             # Execute the operation for edit mode
             operation = SynchronousOperation(
                 endpoint=ApiEndpoint(
@@ -702,10 +758,7 @@ class IdeogramV3(comfy_io.ComfyNode):
                     response_model=IdeogramGenerateResponse,
                 ),
                 request=edit_request,
-                files={
-                    "image": img_binary,
-                    "mask": mask_binary,
-                },
+                files=files,
                 content_type="multipart/form-data",
                 auth_kwargs=auth,
             )
@@ -739,6 +792,14 @@ class IdeogramV3(comfy_io.ComfyNode):
             if num_images > 1:
                 gen_request.num_images = num_images
 
+            files = {}
+            if character_img_binary:
+                files["character_reference_images"] = character_img_binary
+            if character_mask_binary:
+                files["character_mask_binary"] = character_mask_binary
+            if files:
+                gen_request.style_type = "AUTO"
+
             # Execute the operation for generation mode
             operation = SynchronousOperation(
                 endpoint=ApiEndpoint(
@@ -748,6 +809,8 @@ class IdeogramV3(comfy_io.ComfyNode):
                     response_model=IdeogramGenerateResponse,
                 ),
                 request=gen_request,
+                files=files if files else None,
+                content_type="multipart/form-data",
                 auth_kwargs=auth,
             )
 

comfy_api_nodes/nodes_runway.py

@@ -12,6 +12,7 @@ User Guides:
 """
 
 from typing import Union, Optional, Any
+from typing_extensions import override
 from enum import Enum
 
 import torch
@@ -46,9 +47,9 @@ from comfy_api_nodes.apinode_utils import (
     validate_string,
     download_url_to_image_tensor,
 )
-from comfy_api_nodes.mapper_utils import model_field_to_node_input
 from comfy_api.input_impl import VideoFromFile
-from comfy.comfy_types.node_typing import IO, ComfyNodeABC
+from comfy_api.latest import ComfyExtension, io as comfy_io
+from comfy_api_nodes.util.validation_utils import validate_image_dimensions, validate_image_aspect_ratio
 
 PATH_IMAGE_TO_VIDEO = "/proxy/runway/image_to_video"
 PATH_TEXT_TO_IMAGE = "/proxy/runway/text_to_image"
@@ -85,20 +86,11 @@ class RunwayGen3aAspectRatio(str, Enum):
 
 def get_video_url_from_task_status(response: TaskStatusResponse) -> Union[str, None]:
     """Returns the video URL from the task status response if it exists."""
-    if response.output and len(response.output) > 0:
+    if hasattr(response, "output") and len(response.output) > 0:
         return response.output[0]
     return None
 
 
-# TODO: replace with updated image validation utils (upstream)
-def validate_input_image(image: torch.Tensor) -> bool:
-    """
-    Validate the input image is within the size limits for the Runway API.
-    See: https://docs.dev.runwayml.com/assets/inputs/#common-error-reasons
-    """
-    return image.shape[2] < 8000 and image.shape[1] < 8000
-
-
 async def poll_until_finished(
     auth_kwargs: dict[str, str],
     api_endpoint: ApiEndpoint[Any, TaskStatusResponse],
@@ -134,458 +126,438 @@ def extract_progress_from_task_status(
 
 def get_image_url_from_task_status(response: TaskStatusResponse) -> Union[str, None]:
     """Returns the image URL from the task status response if it exists."""
-    if response.output and len(response.output) > 0:
+    if hasattr(response, "output") and len(response.output) > 0:
         return response.output[0]
     return None
 
 
-class RunwayVideoGenNode(ComfyNodeABC):
-    """Runway Video Node Base."""
+async def get_response(
+    task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None, estimated_duration: Optional[int] = None
+) -> TaskStatusResponse:
+    """Poll the task status until it is finished then get the response."""
+    return await poll_until_finished(
+        auth_kwargs,
+        ApiEndpoint(
+            path=f"{PATH_GET_TASK_STATUS}/{task_id}",
+            method=HttpMethod.GET,
+            request_model=EmptyRequest,
+            response_model=TaskStatusResponse,
+        ),
+        estimated_duration=estimated_duration,
+        node_id=node_id,
+    )
 
-    RETURN_TYPES = ("VIDEO",)
-    FUNCTION = "api_call"
-    CATEGORY = "api node/video/Runway"
-    API_NODE = True
 
-    def validate_task_created(self, response: RunwayImageToVideoResponse) -> bool:
-        """
-        Validate the task creation response from the Runway API matches
-        expected format.
-        """
-        if not bool(response.id):
-            raise RunwayApiError("Invalid initial response from Runway API.")
-        return True
+async def generate_video(
+    request: RunwayImageToVideoRequest,
+    auth_kwargs: dict[str, str],
+    node_id: Optional[str] = None,
+    estimated_duration: Optional[int] = None,
+) -> VideoFromFile:
+    initial_operation = SynchronousOperation(
+        endpoint=ApiEndpoint(
+            path=PATH_IMAGE_TO_VIDEO,
+            method=HttpMethod.POST,
+            request_model=RunwayImageToVideoRequest,
+            response_model=RunwayImageToVideoResponse,
+        ),
+        request=request,
+        auth_kwargs=auth_kwargs,
+    )
 
-    def validate_response(self, response: RunwayImageToVideoResponse) -> bool:
-        """
-        Validate the successful task status response from the Runway API
-        matches expected format.
-        """
-        if not response.output or len(response.output) == 0:
-            raise RunwayApiError(
-                "Runway task succeeded but no video data found in response."
-            )
-        return True
+    initial_response = await initial_operation.execute()
 
-    async def get_response(
-        self, task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None
-    ) -> RunwayImageToVideoResponse:
-        """Poll the task status until it is finished then get the response."""
-        return await poll_until_finished(
-            auth_kwargs,
-            ApiEndpoint(
-                path=f"{PATH_GET_TASK_STATUS}/{task_id}",
-                method=HttpMethod.GET,
-                request_model=EmptyRequest,
-                response_model=TaskStatusResponse,
-            ),
-            estimated_duration=AVERAGE_DURATION_FLF_SECONDS,
-            node_id=node_id,
+    final_response = await get_response(initial_response.id, auth_kwargs, node_id, estimated_duration)
+    if not final_response.output:
+        raise RunwayApiError("Runway task succeeded but no video data found in response.")
+
+    video_url = get_video_url_from_task_status(final_response)
+    return await download_url_to_video_output(video_url)
+
+
+class RunwayImageToVideoNodeGen3a(comfy_io.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="RunwayImageToVideoNodeGen3a",
+            display_name="Runway Image to Video (Gen3a Turbo)",
+            category="api node/video/Runway",
+            description="Generate a video from a single starting frame using Gen3a Turbo model. "
+                        "Before diving in, review these best practices to ensure that "
+                        "your input selections will set your generation up for success: "
+                        "https://help.runwayml.com/hc/en-us/articles/33927968552339-Creating-with-Act-One-on-Gen-3-Alpha-and-Turbo.",
+            inputs=[
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Text prompt for the generation",
+                ),
+                comfy_io.Image.Input(
+                    "start_frame",
+                    tooltip="Start frame to be used for the video",
+                ),
+                comfy_io.Combo.Input(
+                    "duration",
+                    options=[model.value for model in Duration],
+                ),
+                comfy_io.Combo.Input(
+                    "ratio",
+                    options=[model.value for model in RunwayGen3aAspectRatio],
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=4294967295,
+                    step=1,
+                    control_after_generate=True,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Random seed for generation",
+                ),
+            ],
+            outputs=[
+                comfy_io.Video.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
         )
 
-    async def generate_video(
-        self,
-        request: RunwayImageToVideoRequest,
-        auth_kwargs: dict[str, str],
-        node_id: Optional[str] = None,
-    ) -> tuple[VideoFromFile]:
-        initial_operation = SynchronousOperation(
-            endpoint=ApiEndpoint(
-                path=PATH_IMAGE_TO_VIDEO,
-                method=HttpMethod.POST,
-                request_model=RunwayImageToVideoRequest,
-                response_model=RunwayImageToVideoResponse,
-            ),
-            request=request,
+    @classmethod
+    async def execute(
+        cls,
+        prompt: str,
+        start_frame: torch.Tensor,
+        duration: str,
+        ratio: str,
+        seed: int,
+    ) -> comfy_io.NodeOutput:
+        validate_string(prompt, min_length=1)
+        validate_image_dimensions(start_frame, max_width=7999, max_height=7999)
+        validate_image_aspect_ratio(start_frame, min_aspect_ratio=0.5, max_aspect_ratio=2.0)
+
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
+        download_urls = await upload_images_to_comfyapi(
+            start_frame,
+            max_images=1,
+            mime_type="image/png",
             auth_kwargs=auth_kwargs,
         )
 
-        initial_response = await initial_operation.execute()
-        self.validate_task_created(initial_response)
-        task_id = initial_response.id
-
-        final_response = await self.get_response(task_id, auth_kwargs, node_id)
-        self.validate_response(final_response)
-
-        video_url = get_video_url_from_task_status(final_response)
-        return (await download_url_to_video_output(video_url),)
+        return comfy_io.NodeOutput(
+            await generate_video(
+                RunwayImageToVideoRequest(
+                    promptText=prompt,
+                    seed=seed,
+                    model=Model("gen3a_turbo"),
+                    duration=Duration(duration),
+                    ratio=AspectRatio(ratio),
+                    promptImage=RunwayPromptImageObject(
+                        root=[
+                            RunwayPromptImageDetailedObject(
+                                uri=str(download_urls[0]), position="first"
+                            )
+                        ]
+                    ),
+                ),
+                auth_kwargs=auth_kwargs,
+                node_id=cls.hidden.unique_id,
+            )
+        )
 
 
-class RunwayImageToVideoNodeGen3a(RunwayVideoGenNode):
-    """Runway Image to Video Node using Gen3a Turbo model."""
-
-    DESCRIPTION = "Generate a video from a single starting frame using Gen3a Turbo model. Before diving in, review these best practices to ensure that your input selections will set your generation up for success: https://help.runwayml.com/hc/en-us/articles/33927968552339-Creating-with-Act-One-on-Gen-3-Alpha-and-Turbo."
+class RunwayImageToVideoNodeGen4(comfy_io.ComfyNode):
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": model_field_to_node_input(
-                    IO.STRING, RunwayImageToVideoRequest, "promptText", multiline=True
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="RunwayImageToVideoNodeGen4",
+            display_name="Runway Image to Video (Gen4 Turbo)",
+            category="api node/video/Runway",
+            description="Generate a video from a single starting frame using Gen4 Turbo model. "
+                        "Before diving in, review these best practices to ensure that "
+                        "your input selections will set your generation up for success: "
+                        "https://help.runwayml.com/hc/en-us/articles/37327109429011-Creating-with-Gen-4-Video.",
+            inputs=[
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Text prompt for the generation",
                 ),
-                "start_frame": (
-                    IO.IMAGE,
-                    {"tooltip": "Start frame to be used for the video"},
+                comfy_io.Image.Input(
+                    "start_frame",
+                    tooltip="Start frame to be used for the video",
                 ),
-                "duration": model_field_to_node_input(
-                    IO.COMBO, RunwayImageToVideoRequest, "duration", enum_type=Duration
+                comfy_io.Combo.Input(
+                    "duration",
+                    options=[model.value for model in Duration],
                 ),
-                "ratio": model_field_to_node_input(
-                    IO.COMBO,
-                    RunwayImageToVideoRequest,
+                comfy_io.Combo.Input(
                     "ratio",
-                    enum_type=RunwayGen3aAspectRatio,
+                    options=[model.value for model in RunwayGen4TurboAspectRatio],
                 ),
-                "seed": model_field_to_node_input(
-                    IO.INT,
-                    RunwayImageToVideoRequest,
+                comfy_io.Int.Input(
                     "seed",
+                    default=0,
+                    min=0,
+                    max=4294967295,
+                    step=1,
                     control_after_generate=True,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Random seed for generation",
                 ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-                "unique_id": "UNIQUE_ID",
-            },
-        }
+            ],
+            outputs=[
+                comfy_io.Video.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
 
-    async def api_call(
-        self,
+    @classmethod
+    async def execute(
+        cls,
         prompt: str,
         start_frame: torch.Tensor,
         duration: str,
         ratio: str,
         seed: int,
-        unique_id: Optional[str] = None,
-        **kwargs,
-    ) -> tuple[VideoFromFile]:
-        # Validate inputs
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, min_length=1)
-        validate_input_image(start_frame)
+        validate_image_dimensions(start_frame, max_width=7999, max_height=7999)
+        validate_image_aspect_ratio(start_frame, min_aspect_ratio=0.5, max_aspect_ratio=2.0)
+
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
 
-        # Upload image
         download_urls = await upload_images_to_comfyapi(
             start_frame,
             max_images=1,
             mime_type="image/png",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth_kwargs,
         )
-        if len(download_urls) != 1:
-            raise RunwayApiError("Failed to upload one or more images to comfy api.")
 
-        return await self.generate_video(
-            RunwayImageToVideoRequest(
-                promptText=prompt,
-                seed=seed,
-                model=Model("gen3a_turbo"),
-                duration=Duration(duration),
-                ratio=AspectRatio(ratio),
-                promptImage=RunwayPromptImageObject(
-                    root=[
-                        RunwayPromptImageDetailedObject(
-                            uri=str(download_urls[0]), position="first"
-                        )
-                    ]
+        return comfy_io.NodeOutput(
+            await generate_video(
+                RunwayImageToVideoRequest(
+                    promptText=prompt,
+                    seed=seed,
+                    model=Model("gen4_turbo"),
+                    duration=Duration(duration),
+                    ratio=AspectRatio(ratio),
+                    promptImage=RunwayPromptImageObject(
+                        root=[
+                            RunwayPromptImageDetailedObject(
+                                uri=str(download_urls[0]), position="first"
+                            )
+                        ]
+                    ),
                 ),
-            ),
-            auth_kwargs=kwargs,
-            node_id=unique_id,
+                auth_kwargs=auth_kwargs,
+                node_id=cls.hidden.unique_id,
+                estimated_duration=AVERAGE_DURATION_FLF_SECONDS,
+            )
         )
 
 
-class RunwayImageToVideoNodeGen4(RunwayVideoGenNode):
-    """Runway Image to Video Node using Gen4 Turbo model."""
-
-    DESCRIPTION = "Generate a video from a single starting frame using Gen4 Turbo model. Before diving in, review these best practices to ensure that your input selections will set your generation up for success: https://help.runwayml.com/hc/en-us/articles/37327109429011-Creating-with-Gen-4-Video."
+class RunwayFirstLastFrameNode(comfy_io.ComfyNode):
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": model_field_to_node_input(
-                    IO.STRING, RunwayImageToVideoRequest, "promptText", multiline=True
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="RunwayFirstLastFrameNode",
+            display_name="Runway First-Last-Frame to Video",
+            category="api node/video/Runway",
+            description="Upload first and last keyframes, draft a prompt, and generate a video. "
+                        "More complex transitions, such as cases where the Last frame is completely different "
+                        "from the First frame, may benefit from the longer 10s duration. "
+                        "This would give the generation more time to smoothly transition between the two inputs. "
+                        "Before diving in, review these best practices to ensure that your input selections "
+                        "will set your generation up for success: "
+                        "https://help.runwayml.com/hc/en-us/articles/34170748696595-Creating-with-Keyframes-on-Gen-3.",
+            inputs=[
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Text prompt for the generation",
                 ),
-                "start_frame": (
-                    IO.IMAGE,
-                    {"tooltip": "Start frame to be used for the video"},
+                comfy_io.Image.Input(
+                    "start_frame",
+                    tooltip="Start frame to be used for the video",
                 ),
-                "duration": model_field_to_node_input(
-                    IO.COMBO, RunwayImageToVideoRequest, "duration", enum_type=Duration
+                comfy_io.Image.Input(
+                    "end_frame",
+                    tooltip="End frame to be used for the video. Supported for gen3a_turbo only.",
                 ),
-                "ratio": model_field_to_node_input(
-                    IO.COMBO,
-                    RunwayImageToVideoRequest,
+                comfy_io.Combo.Input(
+                    "duration",
+                    options=[model.value for model in Duration],
+                ),
+                comfy_io.Combo.Input(
                     "ratio",
-                    enum_type=RunwayGen4TurboAspectRatio,
+                    options=[model.value for model in RunwayGen3aAspectRatio],
                 ),
-                "seed": model_field_to_node_input(
-                    IO.INT,
-                    RunwayImageToVideoRequest,
+                comfy_io.Int.Input(
                     "seed",
+                    default=0,
+                    min=0,
+                    max=4294967295,
+                    step=1,
                     control_after_generate=True,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Random seed for generation",
                 ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-                "unique_id": "UNIQUE_ID",
-            },
-        }
-
-    async def api_call(
-        self,
-        prompt: str,
-        start_frame: torch.Tensor,
-        duration: str,
-        ratio: str,
-        seed: int,
-        unique_id: Optional[str] = None,
-        **kwargs,
-    ) -> tuple[VideoFromFile]:
-        # Validate inputs
-        validate_string(prompt, min_length=1)
-        validate_input_image(start_frame)
-
-        # Upload image
-        download_urls = await upload_images_to_comfyapi(
-            start_frame,
-            max_images=1,
-            mime_type="image/png",
-            auth_kwargs=kwargs,
-        )
-        if len(download_urls) != 1:
-            raise RunwayApiError("Failed to upload one or more images to comfy api.")
-
-        return await self.generate_video(
-            RunwayImageToVideoRequest(
-                promptText=prompt,
-                seed=seed,
-                model=Model("gen4_turbo"),
-                duration=Duration(duration),
-                ratio=AspectRatio(ratio),
-                promptImage=RunwayPromptImageObject(
-                    root=[
-                        RunwayPromptImageDetailedObject(
-                            uri=str(download_urls[0]), position="first"
-                        )
-                    ]
-                ),
-            ),
-            auth_kwargs=kwargs,
-            node_id=unique_id,
-        )
-
-
-class RunwayFirstLastFrameNode(RunwayVideoGenNode):
-    """Runway First-Last Frame Node."""
-
-    DESCRIPTION = "Upload first and last keyframes, draft a prompt, and generate a video. More complex transitions, such as cases where the Last frame is completely different from the First frame, may benefit from the longer 10s duration. This would give the generation more time to smoothly transition between the two inputs. Before diving in, review these best practices to ensure that your input selections will set your generation up for success: https://help.runwayml.com/hc/en-us/articles/34170748696595-Creating-with-Keyframes-on-Gen-3."
-
-    async def get_response(
-        self, task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None
-    ) -> RunwayImageToVideoResponse:
-        return await poll_until_finished(
-            auth_kwargs,
-            ApiEndpoint(
-                path=f"{PATH_GET_TASK_STATUS}/{task_id}",
-                method=HttpMethod.GET,
-                request_model=EmptyRequest,
-                response_model=TaskStatusResponse,
-            ),
-            estimated_duration=AVERAGE_DURATION_FLF_SECONDS,
-            node_id=node_id,
+            ],
+            outputs=[
+                comfy_io.Video.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
         )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": model_field_to_node_input(
-                    IO.STRING, RunwayImageToVideoRequest, "promptText", multiline=True
-                ),
-                "start_frame": (
-                    IO.IMAGE,
-                    {"tooltip": "Start frame to be used for the video"},
-                ),
-                "end_frame": (
-                    IO.IMAGE,
-                    {
-                        "tooltip": "End frame to be used for the video. Supported for gen3a_turbo only."
-                    },
-                ),
-                "duration": model_field_to_node_input(
-                    IO.COMBO, RunwayImageToVideoRequest, "duration", enum_type=Duration
-                ),
-                "ratio": model_field_to_node_input(
-                    IO.COMBO,
-                    RunwayImageToVideoRequest,
-                    "ratio",
-                    enum_type=RunwayGen3aAspectRatio,
-                ),
-                "seed": model_field_to_node_input(
-                    IO.INT,
-                    RunwayImageToVideoRequest,
-                    "seed",
-                    control_after_generate=True,
-                ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "unique_id": "UNIQUE_ID",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-            },
-        }
-
-    async def api_call(
-        self,
+    async def execute(
+        cls,
         prompt: str,
         start_frame: torch.Tensor,
         end_frame: torch.Tensor,
         duration: str,
         ratio: str,
         seed: int,
-        unique_id: Optional[str] = None,
-        **kwargs,
-    ) -> tuple[VideoFromFile]:
-        # Validate inputs
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, min_length=1)
-        validate_input_image(start_frame)
-        validate_input_image(end_frame)
+        validate_image_dimensions(start_frame, max_width=7999, max_height=7999)
+        validate_image_dimensions(end_frame, max_width=7999, max_height=7999)
+        validate_image_aspect_ratio(start_frame, min_aspect_ratio=0.5, max_aspect_ratio=2.0)
+        validate_image_aspect_ratio(end_frame, min_aspect_ratio=0.5, max_aspect_ratio=2.0)
+
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
 
-        # Upload images
         stacked_input_images = image_tensor_pair_to_batch(start_frame, end_frame)
         download_urls = await upload_images_to_comfyapi(
             stacked_input_images,
             max_images=2,
             mime_type="image/png",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth_kwargs,
         )
         if len(download_urls) != 2:
             raise RunwayApiError("Failed to upload one or more images to comfy api.")
 
-        return await self.generate_video(
-            RunwayImageToVideoRequest(
-                promptText=prompt,
-                seed=seed,
-                model=Model("gen3a_turbo"),
-                duration=Duration(duration),
-                ratio=AspectRatio(ratio),
-                promptImage=RunwayPromptImageObject(
-                    root=[
-                        RunwayPromptImageDetailedObject(
-                            uri=str(download_urls[0]), position="first"
-                        ),
-                        RunwayPromptImageDetailedObject(
-                            uri=str(download_urls[1]), position="last"
-                        ),
-                    ]
+        return comfy_io.NodeOutput(
+            await generate_video(
+                RunwayImageToVideoRequest(
+                    promptText=prompt,
+                    seed=seed,
+                    model=Model("gen3a_turbo"),
+                    duration=Duration(duration),
+                    ratio=AspectRatio(ratio),
+                    promptImage=RunwayPromptImageObject(
+                        root=[
+                            RunwayPromptImageDetailedObject(
+                                uri=str(download_urls[0]), position="first"
+                            ),
+                            RunwayPromptImageDetailedObject(
+                                uri=str(download_urls[1]), position="last"
+                            ),
+                        ]
+                    ),
                 ),
-            ),
-            auth_kwargs=kwargs,
-            node_id=unique_id,
+                auth_kwargs=auth_kwargs,
+                node_id=cls.hidden.unique_id,
+                estimated_duration=AVERAGE_DURATION_FLF_SECONDS,
+            )
         )
 
 
-class RunwayTextToImageNode(ComfyNodeABC):
-    """Runway Text to Image Node."""
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "api_call"
-    CATEGORY = "api node/image/Runway"
-    API_NODE = True
-    DESCRIPTION = "Generate an image from a text prompt using Runway's Gen 4 model. You can also include reference images to guide the generation."
+class RunwayTextToImageNode(comfy_io.ComfyNode):
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": model_field_to_node_input(
-                    IO.STRING, RunwayTextToImageRequest, "promptText", multiline=True
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="RunwayTextToImageNode",
+            display_name="Runway Text to Image",
+            category="api node/image/Runway",
+            description="Generate an image from a text prompt using Runway's Gen 4 model. "
+                        "You can also include reference image to guide the generation.",
+            inputs=[
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Text prompt for the generation",
                 ),
-                "ratio": model_field_to_node_input(
-                    IO.COMBO,
-                    RunwayTextToImageRequest,
+                comfy_io.Combo.Input(
                     "ratio",
-                    enum_type=RunwayTextToImageAspectRatioEnum,
+                    options=[model.value for model in RunwayTextToImageAspectRatioEnum],
                 ),
-            },
-            "optional": {
-                "reference_image": (
-                    IO.IMAGE,
-                    {"tooltip": "Optional reference image to guide the generation"},
-                )
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-                "unique_id": "UNIQUE_ID",
-            },
-        }
-
-    def validate_task_created(self, response: RunwayTextToImageResponse) -> bool:
-        """
-        Validate the task creation response from the Runway API matches
-        expected format.
-        """
-        if not bool(response.id):
-            raise RunwayApiError("Invalid initial response from Runway API.")
-        return True
-
-    def validate_response(self, response: TaskStatusResponse) -> bool:
-        """
-        Validate the successful task status response from the Runway API
-        matches expected format.
-        """
-        if not response.output or len(response.output) == 0:
-            raise RunwayApiError(
-                "Runway task succeeded but no image data found in response."
-            )
-        return True
-
-    async def get_response(
-        self, task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None
-    ) -> TaskStatusResponse:
-        """Poll the task status until it is finished then get the response."""
-        return await poll_until_finished(
-            auth_kwargs,
-            ApiEndpoint(
-                path=f"{PATH_GET_TASK_STATUS}/{task_id}",
-                method=HttpMethod.GET,
-                request_model=EmptyRequest,
-                response_model=TaskStatusResponse,
-            ),
-            estimated_duration=AVERAGE_DURATION_T2I_SECONDS,
-            node_id=node_id,
+                comfy_io.Image.Input(
+                    "reference_image",
+                    tooltip="Optional reference image to guide the generation",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
         )
 
-    async def api_call(
-        self,
+    @classmethod
+    async def execute(
+        cls,
         prompt: str,
         ratio: str,
         reference_image: Optional[torch.Tensor] = None,
-        unique_id: Optional[str] = None,
-        **kwargs,
-    ) -> tuple[torch.Tensor]:
-        # Validate inputs
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, min_length=1)
 
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
         # Prepare reference images if provided
         reference_images = None
         if reference_image is not None:
-            validate_input_image(reference_image)
+            validate_image_dimensions(reference_image, max_width=7999, max_height=7999)
+            validate_image_aspect_ratio(reference_image, min_aspect_ratio=0.5, max_aspect_ratio=2.0)
             download_urls = await upload_images_to_comfyapi(
                 reference_image,
                 max_images=1,
                 mime_type="image/png",
-                auth_kwargs=kwargs,
+                auth_kwargs=auth_kwargs,
             )
-            if len(download_urls) != 1:
-                raise RunwayApiError("Failed to upload reference image to comfy api.")
-
             reference_images = [ReferenceImage(uri=str(download_urls[0]))]
 
-        # Create request
         request = RunwayTextToImageRequest(
             promptText=prompt,
             model=Model4.gen4_image,
@@ -593,7 +565,6 @@ class RunwayTextToImageNode(ComfyNodeABC):
             referenceImages=reference_images,
         )
 
-        # Execute initial request
         initial_operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path=PATH_TEXT_TO_IMAGE,
@@ -602,34 +573,33 @@ class RunwayTextToImageNode(ComfyNodeABC):
                 response_model=RunwayTextToImageResponse,
             ),
             request=request,
-            auth_kwargs=kwargs,
+            auth_kwargs=auth_kwargs,
         )
 
         initial_response = await initial_operation.execute()
-        self.validate_task_created(initial_response)
-        task_id = initial_response.id
 
         # Poll for completion
-        final_response = await self.get_response(
-            task_id, auth_kwargs=kwargs, node_id=unique_id
+        final_response = await get_response(
+            initial_response.id,
+            auth_kwargs=auth_kwargs,
+            node_id=cls.hidden.unique_id,
+            estimated_duration=AVERAGE_DURATION_T2I_SECONDS,
         )
-        self.validate_response(final_response)
+        if not final_response.output:
+            raise RunwayApiError("Runway task succeeded but no image data found in response.")
 
-        # Download and return image
-        image_url = get_image_url_from_task_status(final_response)
-        return (await download_url_to_image_tensor(image_url),)
+        return comfy_io.NodeOutput(await download_url_to_image_tensor(get_image_url_from_task_status(final_response)))
 
 
-NODE_CLASS_MAPPINGS = {
-    "RunwayFirstLastFrameNode": RunwayFirstLastFrameNode,
-    "RunwayImageToVideoNodeGen3a": RunwayImageToVideoNodeGen3a,
-    "RunwayImageToVideoNodeGen4": RunwayImageToVideoNodeGen4,
-    "RunwayTextToImageNode": RunwayTextToImageNode,
-}
+class RunwayExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
+        return [
+            RunwayFirstLastFrameNode,
+            RunwayImageToVideoNodeGen3a,
+            RunwayImageToVideoNodeGen4,
+            RunwayTextToImageNode,
+        ]
 
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "RunwayFirstLastFrameNode": "Runway First-Last-Frame to Video",
-    "RunwayImageToVideoNodeGen3a": "Runway Image to Video (Gen3a Turbo)",
-    "RunwayImageToVideoNodeGen4": "Runway Image to Video (Gen4 Turbo)",
-    "RunwayTextToImageNode": "Runway Text to Image",
-}
+async def comfy_entrypoint() -> RunwayExtension:
+    return RunwayExtension()

comfy_api_nodes/nodes_stability.py

@@ -1,5 +1,8 @@
 from inspect import cleandoc
-from comfy.comfy_types.node_typing import IO
+from typing import Optional
+from typing_extensions import override
+
+from comfy_api.latest import ComfyExtension, io as comfy_io
 from comfy_api_nodes.apis.stability_api import (
     StabilityUpscaleConservativeRequest,
     StabilityUpscaleCreativeRequest,
@@ -46,87 +49,94 @@ def get_async_dummy_status(x: StabilityResultsGetResponse):
     return StabilityPollStatus.in_progress
 
 
-class StabilityStableImageUltraNode:
+class StabilityStableImageUltraNode(comfy_io.ComfyNode):
     """
     Generates images synchronously based on prompt and resolution.
     """
 
-    RETURN_TYPES = (IO.IMAGE,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "api_call"
-    API_NODE = True
-    CATEGORY = "api node/image/Stability AI"
-
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": (
-                    IO.STRING,
-                    {
-                        "multiline": True,
-                        "default": "",
-                        "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines" +
-                                    "What you wish to see in the output image. A strong, descriptive prompt that clearly defines" +
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="StabilityStableImageUltraNode",
+            display_name="Stability AI Stable Image Ultra",
+            category="api node/image/Stability AI",
+            description=cleandoc(cls.__doc__ or ""),
+            inputs=[
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines" +
                                     "elements, colors, and subjects will lead to better results. " +
                                     "To control the weight of a given word use the format `(word:weight)`," +
                                     "where `word` is the word you'd like to control the weight of and `weight`" +
                                     "is a value between 0 and 1. For example: `The sky was a crisp (blue:0.3) and (green:0.8)`" +
-                                    "would convey a sky that was blue and green, but more green than blue."
-                    },
+                                    "would convey a sky that was blue and green, but more green than blue.",
                 ),
-                "aspect_ratio": ([x.value for x in StabilityAspectRatio],
-                    {
-                        "default": StabilityAspectRatio.ratio_1_1,
-                        "tooltip": "Aspect ratio of generated image.",
-                    },
+                comfy_io.Combo.Input(
+                    "aspect_ratio",
+                    options=[x.value for x in StabilityAspectRatio],
+                    default=StabilityAspectRatio.ratio_1_1.value,
+                    tooltip="Aspect ratio of generated image.",
                 ),
-                "style_preset": (get_stability_style_presets(),
-                    {
-                        "tooltip": "Optional desired style of generated image.",
-                    },
+                comfy_io.Combo.Input(
+                    "style_preset",
+                    options=get_stability_style_presets(),
+                    tooltip="Optional desired style of generated image.",
                 ),
-                "seed": (
-                    IO.INT,
-                    {
-                        "default": 0,
-                        "min": 0,
-                        "max": 4294967294,
-                        "control_after_generate": True,
-                        "tooltip": "The random seed used for creating the noise.",
-                    },
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=4294967294,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="The random seed used for creating the noise.",
                 ),
-            },
-            "optional": {
-                "image": (IO.IMAGE,),
-                "negative_prompt": (
-                    IO.STRING,
-                    {
-                        "default": "",
-                        "forceInput": True,
-                        "tooltip": "A blurb of text describing what you do not wish to see in the output image. This is an advanced feature."
-                    },
+                comfy_io.Image.Input(
+                    "image",
+                    optional=True,
                 ),
-                "image_denoise": (
-                    IO.FLOAT,
-                    {
-                        "default": 0.5,
-                        "min": 0.0,
-                        "max": 1.0,
-                        "step": 0.01,
-                        "tooltip": "Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.",
-                    },
+                comfy_io.String.Input(
+                    "negative_prompt",
+                    default="",
+                    tooltip="A blurb of text describing what you do not wish to see in the output image. This is an advanced feature.",
+                    force_input=True,
+                    optional=True,
                 ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-            },
-        }
+                comfy_io.Float.Input(
+                    "image_denoise",
+                    default=0.5,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip="Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
 
-    async def api_call(self, prompt: str, aspect_ratio: str, style_preset: str, seed: int,
-                 negative_prompt: str=None, image: torch.Tensor = None, image_denoise: float=None,
-                 **kwargs):
+    @classmethod
+    async def execute(
+        cls,
+        prompt: str,
+        aspect_ratio: str,
+        style_preset: str,
+        seed: int,
+        image: Optional[torch.Tensor] = None,
+        negative_prompt: str = "",
+        image_denoise: Optional[float] = 0.5,
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, strip_whitespace=False)
         # prepare image binary if image present
         image_binary = None
@@ -144,6 +154,11 @@ class StabilityStableImageUltraNode:
             "image": image_binary
         }
 
+        auth = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/stability/v2beta/stable-image/generate/ultra",
@@ -161,7 +176,7 @@ class StabilityStableImageUltraNode:
             ),
             files=files,
             content_type="multipart/form-data",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth,
         )
         response_api = await operation.execute()
 
@@ -171,95 +186,106 @@ class StabilityStableImageUltraNode:
         image_data = base64.b64decode(response_api.image)
         returned_image = bytesio_to_image_tensor(BytesIO(image_data))
 
-        return (returned_image,)
+        return comfy_io.NodeOutput(returned_image)
 
 
-class StabilityStableImageSD_3_5Node:
+class StabilityStableImageSD_3_5Node(comfy_io.ComfyNode):
     """
     Generates images synchronously based on prompt and resolution.
     """
 
-    RETURN_TYPES = (IO.IMAGE,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "api_call"
-    API_NODE = True
-    CATEGORY = "api node/image/Stability AI"
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="StabilityStableImageSD_3_5Node",
+            display_name="Stability AI Stable Diffusion 3.5 Image",
+            category="api node/image/Stability AI",
+            description=cleandoc(cls.__doc__ or ""),
+            inputs=[
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results.",
+                ),
+                comfy_io.Combo.Input(
+                    "model",
+                    options=[x.value for x in Stability_SD3_5_Model],
+                ),
+                comfy_io.Combo.Input(
+                    "aspect_ratio",
+                    options=[x.value for x in StabilityAspectRatio],
+                    default=StabilityAspectRatio.ratio_1_1.value,
+                    tooltip="Aspect ratio of generated image.",
+                ),
+                comfy_io.Combo.Input(
+                    "style_preset",
+                    options=get_stability_style_presets(),
+                    tooltip="Optional desired style of generated image.",
+                ),
+                comfy_io.Float.Input(
+                    "cfg_scale",
+                    default=4.0,
+                    min=1.0,
+                    max=10.0,
+                    step=0.1,
+                    tooltip="How strictly the diffusion process adheres to the prompt text (higher values keep your image closer to your prompt)",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=4294967294,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="The random seed used for creating the noise.",
+                ),
+                comfy_io.Image.Input(
+                    "image",
+                    optional=True,
+                ),
+                comfy_io.String.Input(
+                    "negative_prompt",
+                    default="",
+                    tooltip="Keywords of what you do not wish to see in the output image. This is an advanced feature.",
+                    force_input=True,
+                    optional=True,
+                ),
+                comfy_io.Float.Input(
+                    "image_denoise",
+                    default=0.5,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip="Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": (
-                    IO.STRING,
-                    {
-                        "multiline": True,
-                        "default": "",
-                        "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results."
-                    },
-                ),
-                "model": ([x.value for x in Stability_SD3_5_Model],),
-                "aspect_ratio": ([x.value for x in StabilityAspectRatio],
-                    {
-                        "default": StabilityAspectRatio.ratio_1_1,
-                        "tooltip": "Aspect ratio of generated image.",
-                    },
-                ),
-                "style_preset": (get_stability_style_presets(),
-                    {
-                        "tooltip": "Optional desired style of generated image.",
-                    },
-                ),
-                "cfg_scale": (
-                    IO.FLOAT,
-                    {
-                        "default": 4.0,
-                        "min": 1.0,
-                        "max": 10.0,
-                        "step": 0.1,
-                        "tooltip": "How strictly the diffusion process adheres to the prompt text (higher values keep your image closer to your prompt)",
-                    },
-                ),
-                "seed": (
-                    IO.INT,
-                    {
-                        "default": 0,
-                        "min": 0,
-                        "max": 4294967294,
-                        "control_after_generate": True,
-                        "tooltip": "The random seed used for creating the noise.",
-                    },
-                ),
-            },
-            "optional": {
-                "image": (IO.IMAGE,),
-                "negative_prompt": (
-                    IO.STRING,
-                    {
-                        "default": "",
-                        "forceInput": True,
-                        "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature."
-                    },
-                ),
-                "image_denoise": (
-                    IO.FLOAT,
-                    {
-                        "default": 0.5,
-                        "min": 0.0,
-                        "max": 1.0,
-                        "step": 0.01,
-                        "tooltip": "Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.",
-                    },
-                ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-            },
-        }
-
-    async def api_call(self, model: str, prompt: str, aspect_ratio: str, style_preset: str, seed: int, cfg_scale: float,
-                 negative_prompt: str=None, image: torch.Tensor = None, image_denoise: float=None,
-                 **kwargs):
+    async def execute(
+        cls,
+        model: str,
+        prompt: str,
+        aspect_ratio: str,
+        style_preset: str,
+        seed: int,
+        cfg_scale: float,
+        image: Optional[torch.Tensor] = None,
+        negative_prompt: str = "",
+        image_denoise: Optional[float] = 0.5,
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, strip_whitespace=False)
         # prepare image binary if image present
         image_binary = None
@@ -280,6 +306,11 @@ class StabilityStableImageSD_3_5Node:
             "image": image_binary
         }
 
+        auth = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/stability/v2beta/stable-image/generate/sd3",
@@ -300,7 +331,7 @@ class StabilityStableImageSD_3_5Node:
             ),
             files=files,
             content_type="multipart/form-data",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth,
         )
         response_api = await operation.execute()
 
@@ -310,72 +341,75 @@ class StabilityStableImageSD_3_5Node:
         image_data = base64.b64decode(response_api.image)
         returned_image = bytesio_to_image_tensor(BytesIO(image_data))
 
-        return (returned_image,)
+        return comfy_io.NodeOutput(returned_image)
 
 
-class StabilityUpscaleConservativeNode:
+class StabilityUpscaleConservativeNode(comfy_io.ComfyNode):
     """
     Upscale image with minimal alterations to 4K resolution.
     """
 
-    RETURN_TYPES = (IO.IMAGE,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "api_call"
-    API_NODE = True
-    CATEGORY = "api node/image/Stability AI"
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="StabilityUpscaleConservativeNode",
+            display_name="Stability AI Upscale Conservative",
+            category="api node/image/Stability AI",
+            description=cleandoc(cls.__doc__ or ""),
+            inputs=[
+                comfy_io.Image.Input("image"),
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results.",
+                ),
+                comfy_io.Float.Input(
+                    "creativity",
+                    default=0.35,
+                    min=0.2,
+                    max=0.5,
+                    step=0.01,
+                    tooltip="Controls the likelihood of creating additional details not heavily conditioned by the init image.",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=4294967294,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="The random seed used for creating the noise.",
+                ),
+                comfy_io.String.Input(
+                    "negative_prompt",
+                    default="",
+                    tooltip="Keywords of what you do not wish to see in the output image. This is an advanced feature.",
+                    force_input=True,
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": (IO.IMAGE,),
-                "prompt": (
-                    IO.STRING,
-                    {
-                        "multiline": True,
-                        "default": "",
-                        "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results."
-                    },
-                ),
-                "creativity": (
-                    IO.FLOAT,
-                    {
-                        "default": 0.35,
-                        "min": 0.2,
-                        "max": 0.5,
-                        "step": 0.01,
-                        "tooltip": "Controls the likelihood of creating additional details not heavily conditioned by the init image.",
-                    },
-                ),
-                "seed": (
-                    IO.INT,
-                    {
-                        "default": 0,
-                        "min": 0,
-                        "max": 4294967294,
-                        "control_after_generate": True,
-                        "tooltip": "The random seed used for creating the noise.",
-                    },
-                ),
-            },
-            "optional": {
-                "negative_prompt": (
-                    IO.STRING,
-                    {
-                        "default": "",
-                        "forceInput": True,
-                        "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature."
-                    },
-                ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-            },
-        }
-
-    async def api_call(self, image: torch.Tensor, prompt: str, creativity: float, seed: int, negative_prompt: str=None,
-                 **kwargs):
+    async def execute(
+        cls,
+        image: torch.Tensor,
+        prompt: str,
+        creativity: float,
+        seed: int,
+        negative_prompt: str = "",
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, strip_whitespace=False)
         image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read()
 
@@ -386,6 +420,11 @@ class StabilityUpscaleConservativeNode:
             "image": image_binary
         }
 
+        auth = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/stability/v2beta/stable-image/upscale/conservative",
@@ -401,7 +440,7 @@ class StabilityUpscaleConservativeNode:
             ),
             files=files,
             content_type="multipart/form-data",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth,
         )
         response_api = await operation.execute()
 
@@ -411,77 +450,81 @@ class StabilityUpscaleConservativeNode:
         image_data = base64.b64decode(response_api.image)
         returned_image = bytesio_to_image_tensor(BytesIO(image_data))
 
-        return (returned_image,)
+        return comfy_io.NodeOutput(returned_image)
 
 
-class StabilityUpscaleCreativeNode:
+class StabilityUpscaleCreativeNode(comfy_io.ComfyNode):
     """
     Upscale image with minimal alterations to 4K resolution.
     """
 
-    RETURN_TYPES = (IO.IMAGE,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "api_call"
-    API_NODE = True
-    CATEGORY = "api node/image/Stability AI"
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="StabilityUpscaleCreativeNode",
+            display_name="Stability AI Upscale Creative",
+            category="api node/image/Stability AI",
+            description=cleandoc(cls.__doc__ or ""),
+            inputs=[
+                comfy_io.Image.Input("image"),
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results.",
+                ),
+                comfy_io.Float.Input(
+                    "creativity",
+                    default=0.3,
+                    min=0.1,
+                    max=0.5,
+                    step=0.01,
+                    tooltip="Controls the likelihood of creating additional details not heavily conditioned by the init image.",
+                ),
+                comfy_io.Combo.Input(
+                    "style_preset",
+                    options=get_stability_style_presets(),
+                    tooltip="Optional desired style of generated image.",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=4294967294,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="The random seed used for creating the noise.",
+                ),
+                comfy_io.String.Input(
+                    "negative_prompt",
+                    default="",
+                    tooltip="Keywords of what you do not wish to see in the output image. This is an advanced feature.",
+                    force_input=True,
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": (IO.IMAGE,),
-                "prompt": (
-                    IO.STRING,
-                    {
-                        "multiline": True,
-                        "default": "",
-                        "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results."
-                    },
-                ),
-                "creativity": (
-                    IO.FLOAT,
-                    {
-                        "default": 0.3,
-                        "min": 0.1,
-                        "max": 0.5,
-                        "step": 0.01,
-                        "tooltip": "Controls the likelihood of creating additional details not heavily conditioned by the init image.",
-                    },
-                ),
-                "style_preset": (get_stability_style_presets(),
-                    {
-                        "tooltip": "Optional desired style of generated image.",
-                    },
-                ),
-                "seed": (
-                    IO.INT,
-                    {
-                        "default": 0,
-                        "min": 0,
-                        "max": 4294967294,
-                        "control_after_generate": True,
-                        "tooltip": "The random seed used for creating the noise.",
-                    },
-                ),
-            },
-            "optional": {
-                "negative_prompt": (
-                    IO.STRING,
-                    {
-                        "default": "",
-                        "forceInput": True,
-                        "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature."
-                    },
-                ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-            },
-        }
-
-    async def api_call(self, image: torch.Tensor, prompt: str, creativity: float, style_preset: str, seed: int, negative_prompt: str=None,
-                 **kwargs):
+    async def execute(
+        cls,
+        image: torch.Tensor,
+        prompt: str,
+        creativity: float,
+        style_preset: str,
+        seed: int,
+        negative_prompt: str = "",
+    ) -> comfy_io.NodeOutput:
         validate_string(prompt, strip_whitespace=False)
         image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read()
 
@@ -494,6 +537,11 @@ class StabilityUpscaleCreativeNode:
             "image": image_binary
         }
 
+        auth = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/stability/v2beta/stable-image/upscale/creative",
@@ -510,7 +558,7 @@ class StabilityUpscaleCreativeNode:
             ),
             files=files,
             content_type="multipart/form-data",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth,
         )
         response_api = await operation.execute()
 
@@ -525,7 +573,8 @@ class StabilityUpscaleCreativeNode:
             completed_statuses=[StabilityPollStatus.finished],
             failed_statuses=[StabilityPollStatus.failed],
             status_extractor=lambda x: get_async_dummy_status(x),
-            auth_kwargs=kwargs,
+            auth_kwargs=auth,
+            node_id=cls.hidden.unique_id,
         )
         response_poll: StabilityResultsGetResponse = await operation.execute()
 
@@ -535,41 +584,48 @@ class StabilityUpscaleCreativeNode:
         image_data = base64.b64decode(response_poll.result)
         returned_image = bytesio_to_image_tensor(BytesIO(image_data))
 
-        return (returned_image,)
+        return comfy_io.NodeOutput(returned_image)
 
 
-class StabilityUpscaleFastNode:
+class StabilityUpscaleFastNode(comfy_io.ComfyNode):
     """
     Quickly upscales an image via Stability API call to 4x its original size; intended for upscaling low-quality/compressed images.
     """
 
-    RETURN_TYPES = (IO.IMAGE,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "api_call"
-    API_NODE = True
-    CATEGORY = "api node/image/Stability AI"
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="StabilityUpscaleFastNode",
+            display_name="Stability AI Upscale Fast",
+            category="api node/image/Stability AI",
+            description=cleandoc(cls.__doc__ or ""),
+            inputs=[
+                comfy_io.Image.Input("image"),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": (IO.IMAGE,),
-            },
-            "optional": {
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-                "comfy_api_key": "API_KEY_COMFY_ORG",
-            },
-        }
-
-    async def api_call(self, image: torch.Tensor, **kwargs):
+    async def execute(cls, image: torch.Tensor) -> comfy_io.NodeOutput:
         image_binary = tensor_to_bytesio(image, total_pixels=4096*4096).read()
 
         files = {
             "image": image_binary
         }
 
+        auth = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+
         operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path="/proxy/stability/v2beta/stable-image/upscale/fast",
@@ -580,7 +636,7 @@ class StabilityUpscaleFastNode:
             request=EmptyRequest(),
             files=files,
             content_type="multipart/form-data",
-            auth_kwargs=kwargs,
+            auth_kwargs=auth,
         )
         response_api = await operation.execute()
 
@@ -590,24 +646,20 @@ class StabilityUpscaleFastNode:
         image_data = base64.b64decode(response_api.image)
         returned_image = bytesio_to_image_tensor(BytesIO(image_data))
 
-        return (returned_image,)
+        return comfy_io.NodeOutput(returned_image)
 
 
-# A dictionary that contains all nodes you want to export with their names
-# NOTE: names should be globally unique
-NODE_CLASS_MAPPINGS = {
-    "StabilityStableImageUltraNode": StabilityStableImageUltraNode,
-    "StabilityStableImageSD_3_5Node": StabilityStableImageSD_3_5Node,
-    "StabilityUpscaleConservativeNode": StabilityUpscaleConservativeNode,
-    "StabilityUpscaleCreativeNode": StabilityUpscaleCreativeNode,
-    "StabilityUpscaleFastNode": StabilityUpscaleFastNode,
-}
+class StabilityExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
+        return [
+            StabilityStableImageUltraNode,
+            StabilityStableImageSD_3_5Node,
+            StabilityUpscaleConservativeNode,
+            StabilityUpscaleCreativeNode,
+            StabilityUpscaleFastNode,
+        ]
 
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "StabilityStableImageUltraNode": "Stability AI Stable Image Ultra",
-    "StabilityStableImageSD_3_5Node": "Stability AI Stable Diffusion 3.5 Image",
-    "StabilityUpscaleConservativeNode": "Stability AI Upscale Conservative",
-    "StabilityUpscaleCreativeNode": "Stability AI Upscale Creative",
-    "StabilityUpscaleFastNode": "Stability AI Upscale Fast",
-}
+
+async def comfy_entrypoint() -> StabilityExtension:
+    return StabilityExtension()

comfy_extras/nodes_align_your_steps.py

@@ -1,6 +1,10 @@
 #from: https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
 import numpy as np
 import torch
+from typing_extensions import override
+
+from comfy_api.latest import ComfyExtension, io
+
 
 def loglinear_interp(t_steps, num_steps):
     """
@@ -19,25 +23,30 @@ NOISE_LEVELS = {"SD1": [14.6146412293, 6.4745760956,  3.8636745985,  2.694615152
                 "SDXL":[14.6146412293, 6.3184485287,  3.7681790315,  2.1811480769, 1.3405244945,  0.8620721141,  0.5550693289,  0.3798540708, 0.2332364134,  0.1114188177,  0.0291671582],
                 "SVD": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002]}
 
-class AlignYourStepsScheduler:
+class AlignYourStepsScheduler(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model_type": (["SD1", "SDXL", "SVD"], ),
-                     "steps": ("INT", {"default": 10, "min": 1, "max": 10000}),
-                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                      }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="AlignYourStepsScheduler",
+            category="sampling/custom_sampling/schedulers",
+            inputs=[
+                io.Combo.Input("model_type", options=["SD1", "SDXL", "SVD"]),
+                io.Int.Input("steps", default=10, 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 get_sigmas(self, model_type, steps, denoise):
+        # Deprecated: use the V3 schema's `execute` method instead of this.
+        return AlignYourStepsScheduler().execute(model_type, steps, denoise).result
+
+    @classmethod
+    def execute(cls, model_type, steps, denoise) -> io.NodeOutput:
         total_steps = steps
         if denoise < 1.0:
             if denoise <= 0.0:
-                return (torch.FloatTensor([]),)
+                return io.NodeOutput(torch.FloatTensor([]))
             total_steps = round(steps * denoise)
 
         sigmas = NOISE_LEVELS[model_type][:]
@@ -46,8 +55,15 @@ class AlignYourStepsScheduler:
 
         sigmas = sigmas[-(total_steps + 1):]
         sigmas[-1] = 0
-        return (torch.FloatTensor(sigmas), )
+        return io.NodeOutput(torch.FloatTensor(sigmas))
 
-NODE_CLASS_MAPPINGS = {
-    "AlignYourStepsScheduler": AlignYourStepsScheduler,
-}
+
+class AlignYourStepsExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            AlignYourStepsScheduler,
+        ]
+
+async def comfy_entrypoint() -> AlignYourStepsExtension:
+    return AlignYourStepsExtension()

comfy_extras/nodes_flux.py

@@ -105,7 +105,7 @@ class FluxKontextMultiReferenceLatentMethod:
     def INPUT_TYPES(s):
         return {"required": {
             "conditioning": ("CONDITIONING", ),
-            "reference_latents_method": (("offset", "index"), ),
+            "reference_latents_method": (("offset", "index", "uxo/uno"), ),
             }}
 
     RETURN_TYPES = ("CONDITIONING",)
@@ -115,6 +115,8 @@ class FluxKontextMultiReferenceLatentMethod:
     CATEGORY = "advanced/conditioning/flux"
 
     def append(self, conditioning, reference_latents_method):
+        if "uxo" in reference_latents_method or "uso" in reference_latents_method:
+            reference_latents_method = "uxo"
         c = node_helpers.conditioning_set_values(conditioning, {"reference_latents_method": reference_latents_method})
         return (c, )
 

comfy_extras/nodes_images.py

@@ -625,6 +625,37 @@ class ImageFlip:
 
         return (image,)
 
+class ImageScaleToMaxDimension:
+    upscale_methods = ["area", "lanczos", "bilinear", "nearest-exact", "bilinear", "bicubic"]
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"image": ("IMAGE",),
+                             "upscale_method": (s.upscale_methods,),
+                             "largest_size": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1})}}
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "upscale"
+
+    CATEGORY = "image/upscaling"
+
+    def upscale(self, image, upscale_method, largest_size):
+        height = image.shape[1]
+        width = image.shape[2]
+
+        if height > width:
+            width = round((width / height) * largest_size)
+            height = largest_size
+        elif width > height:
+            height = round((height / width) * largest_size)
+            width = largest_size
+        else:
+            height = largest_size
+            width = largest_size
+
+        samples = image.movedim(-1, 1)
+        s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
+        s = s.movedim(1, -1)
+        return (s,)
 
 NODE_CLASS_MAPPINGS = {
     "ImageCrop": ImageCrop,
@@ -639,4 +670,5 @@ NODE_CLASS_MAPPINGS = {
     "GetImageSize": GetImageSize,
     "ImageRotate": ImageRotate,
     "ImageFlip": ImageFlip,
+    "ImageScaleToMaxDimension": ImageScaleToMaxDimension,
 }

comfy_extras/nodes_model_patch.py

@@ -1,4 +1,5 @@
 import torch
+from torch import nn
 import folder_paths
 import comfy.utils
 import comfy.ops
@@ -58,6 +59,136 @@ class QwenImageBlockWiseControlNet(torch.nn.Module):
         return self.controlnet_blocks[block_id](img, controlnet_conditioning)
 
 
+class SigLIPMultiFeatProjModel(torch.nn.Module):
+    """
+    SigLIP Multi-Feature Projection Model for processing style features from different layers
+    and projecting them into a unified hidden space.
+
+    Args:
+        siglip_token_nums (int): Number of SigLIP tokens, default 257
+        style_token_nums (int): Number of style tokens, default 256
+        siglip_token_dims (int): Dimension of SigLIP tokens, default 1536
+        hidden_size (int): Hidden layer size, default 3072
+        context_layer_norm (bool): Whether to use context layer normalization, default False
+    """
+
+    def __init__(
+        self,
+        siglip_token_nums: int = 729,
+        style_token_nums: int = 64,
+        siglip_token_dims: int = 1152,
+        hidden_size: int = 3072,
+        context_layer_norm: bool = True,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+
+        # High-level feature processing (layer -2)
+        self.high_embedding_linear = nn.Sequential(
+            operations.Linear(siglip_token_nums, style_token_nums),
+            nn.SiLU()
+        )
+        self.high_layer_norm = (
+            operations.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()
+        )
+        self.high_projection = operations.Linear(siglip_token_dims, hidden_size, bias=True)
+
+        # Mid-level feature processing (layer -11)
+        self.mid_embedding_linear = nn.Sequential(
+            operations.Linear(siglip_token_nums, style_token_nums),
+            nn.SiLU()
+        )
+        self.mid_layer_norm = (
+            operations.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()
+        )
+        self.mid_projection = operations.Linear(siglip_token_dims, hidden_size, bias=True)
+
+        # Low-level feature processing (layer -20)
+        self.low_embedding_linear = nn.Sequential(
+            operations.Linear(siglip_token_nums, style_token_nums),
+            nn.SiLU()
+        )
+        self.low_layer_norm = (
+            operations.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()
+        )
+        self.low_projection = operations.Linear(siglip_token_dims, hidden_size, bias=True)
+
+    def forward(self, siglip_outputs):
+        """
+        Forward pass function
+
+        Args:
+            siglip_outputs: Output from SigLIP model, containing hidden_states
+
+        Returns:
+            torch.Tensor: Concatenated multi-layer features with shape [bs, 3*style_token_nums, hidden_size]
+        """
+        dtype = next(self.high_embedding_linear.parameters()).dtype
+
+        # Process high-level features (layer -2)
+        high_embedding = self._process_layer_features(
+            siglip_outputs[2],
+            self.high_embedding_linear,
+            self.high_layer_norm,
+            self.high_projection,
+            dtype
+        )
+
+        # Process mid-level features (layer -11)
+        mid_embedding = self._process_layer_features(
+            siglip_outputs[1],
+            self.mid_embedding_linear,
+            self.mid_layer_norm,
+            self.mid_projection,
+            dtype
+        )
+
+        # Process low-level features (layer -20)
+        low_embedding = self._process_layer_features(
+            siglip_outputs[0],
+            self.low_embedding_linear,
+            self.low_layer_norm,
+            self.low_projection,
+            dtype
+        )
+
+        # Concatenate features from all layersmodel_patch
+        return torch.cat((high_embedding, mid_embedding, low_embedding), dim=1)
+
+    def _process_layer_features(
+        self,
+        hidden_states: torch.Tensor,
+        embedding_linear: nn.Module,
+        layer_norm: nn.Module,
+        projection: nn.Module,
+        dtype: torch.dtype
+    ) -> torch.Tensor:
+        """
+        Helper function to process features from a single layer
+
+        Args:
+            hidden_states: Input hidden states [bs, seq_len, dim]
+            embedding_linear: Embedding linear layer
+            layer_norm: Layer normalization
+            projection: Projection layer
+            dtype: Target data type
+
+        Returns:
+            torch.Tensor: Processed features [bs, style_token_nums, hidden_size]
+        """
+        # Transform dimensions: [bs, seq_len, dim] -> [bs, dim, seq_len] -> [bs, dim, style_token_nums] -> [bs, style_token_nums, dim]
+        embedding = embedding_linear(
+            hidden_states.to(dtype).transpose(1, 2)
+        ).transpose(1, 2)
+
+        # Apply layer normalization
+        embedding = layer_norm(embedding)
+
+        # Project to target hidden space
+        embedding = projection(embedding)
+
+        return embedding
+
 class ModelPatchLoader:
     @classmethod
     def INPUT_TYPES(s):
@@ -73,9 +204,14 @@ class ModelPatchLoader:
         model_patch_path = folder_paths.get_full_path_or_raise("model_patches", name)
         sd = comfy.utils.load_torch_file(model_patch_path, safe_load=True)
         dtype = comfy.utils.weight_dtype(sd)
-        # TODO: this node will work with more types of model patches
-        additional_in_dim = sd["img_in.weight"].shape[1] - 64
-        model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
+
+        if 'controlnet_blocks.0.y_rms.weight' in sd:
+            additional_in_dim = sd["img_in.weight"].shape[1] - 64
+            model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
+        elif 'feature_embedder.mid_layer_norm.bias' in sd:
+            sd = comfy.utils.state_dict_prefix_replace(sd, {"feature_embedder.": ""}, filter_keys=True)
+            model = SigLIPMultiFeatProjModel(device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
+
         model.load_state_dict(sd)
         model = comfy.model_patcher.ModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
         return (model,)
@@ -157,7 +293,51 @@ class QwenImageDiffsynthControlnet:
         return (model_patched,)
 
 
+class UsoStyleProjectorPatch:
+    def __init__(self, model_patch, encoded_image):
+        self.model_patch = model_patch
+        self.encoded_image = encoded_image
+
+    def __call__(self, kwargs):
+        txt_ids = kwargs.get("txt_ids")
+        txt = kwargs.get("txt")
+        siglip_embedding = self.model_patch.model(self.encoded_image.to(txt.dtype)).to(txt.dtype)
+        txt = torch.cat([siglip_embedding, txt], dim=1)
+        kwargs['txt'] = txt
+        kwargs['txt_ids'] = torch.cat([torch.zeros(siglip_embedding.shape[0], siglip_embedding.shape[1], 3, dtype=txt_ids.dtype, device=txt_ids.device), txt_ids], dim=1)
+        return kwargs
+
+    def to(self, device_or_dtype):
+        if isinstance(device_or_dtype, torch.device):
+            self.encoded_image = self.encoded_image.to(device_or_dtype)
+        return self
+
+    def models(self):
+        return [self.model_patch]
+
+
+class USOStyleReference:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL",),
+                             "model_patch": ("MODEL_PATCH",),
+                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply_patch"
+    EXPERIMENTAL = True
+
+    CATEGORY = "advanced/model_patches/flux"
+
+    def apply_patch(self, model, model_patch, clip_vision_output):
+        encoded_image = torch.stack((clip_vision_output.all_hidden_states[:, -20], clip_vision_output.all_hidden_states[:, -11], clip_vision_output.penultimate_hidden_states))
+        model_patched = model.clone()
+        model_patched.set_model_post_input_patch(UsoStyleProjectorPatch(model_patch, encoded_image))
+        return (model_patched,)
+
+
 NODE_CLASS_MAPPINGS = {
     "ModelPatchLoader": ModelPatchLoader,
     "QwenImageDiffsynthControlnet": QwenImageDiffsynthControlnet,
+    "USOStyleReference": USOStyleReference,
 }

comfy_extras/nodes_primitive.py

@@ -1,98 +1,109 @@
-# Primitive nodes that are evaluated at backend.
-from __future__ import annotations
-
 import sys
+from typing_extensions import override
 
-from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, IO
+from comfy_api.latest import ComfyExtension, io
 
 
-class String(ComfyNodeABC):
+class String(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls) -> InputTypeDict:
-        return {
-            "required": {"value": (IO.STRING, {})},
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PrimitiveString",
+            display_name="String",
+            category="utils/primitive",
+            inputs=[
+                io.String.Input("value"),
+            ],
+            outputs=[io.String.Output()],
+        )
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/primitive"
-
-    def execute(self, value: str) -> tuple[str]:
-        return (value,)
-
-
-class StringMultiline(ComfyNodeABC):
     @classmethod
-    def INPUT_TYPES(cls) -> InputTypeDict:
-        return {
-            "required": {"value": (IO.STRING, {"multiline": True,},)},
-        }
-
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/primitive"
-
-    def execute(self, value: str) -> tuple[str]:
-        return (value,)
+    def execute(cls, value: str) -> io.NodeOutput:
+        return io.NodeOutput(value)
 
 
-class Int(ComfyNodeABC):
+class StringMultiline(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls) -> InputTypeDict:
-        return {
-            "required": {"value": (IO.INT, {"min": -sys.maxsize, "max": sys.maxsize, "control_after_generate": True})},
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PrimitiveStringMultiline",
+            display_name="String (Multiline)",
+            category="utils/primitive",
+            inputs=[
+                io.String.Input("value", multiline=True),
+            ],
+            outputs=[io.String.Output()],
+        )
 
-    RETURN_TYPES = (IO.INT,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/primitive"
-
-    def execute(self, value: int) -> tuple[int]:
-        return (value,)
-
-
-class Float(ComfyNodeABC):
     @classmethod
-    def INPUT_TYPES(cls) -> InputTypeDict:
-        return {
-            "required": {"value": (IO.FLOAT, {"min": -sys.maxsize, "max": sys.maxsize})},
-        }
-
-    RETURN_TYPES = (IO.FLOAT,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/primitive"
-
-    def execute(self, value: float) -> tuple[float]:
-        return (value,)
+    def execute(cls, value: str) -> io.NodeOutput:
+        return io.NodeOutput(value)
 
 
-class Boolean(ComfyNodeABC):
+class Int(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls) -> InputTypeDict:
-        return {
-            "required": {"value": (IO.BOOLEAN, {})},
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PrimitiveInt",
+            display_name="Int",
+            category="utils/primitive",
+            inputs=[
+                io.Int.Input("value", min=-sys.maxsize, max=sys.maxsize, control_after_generate=True),
+            ],
+            outputs=[io.Int.Output()],
+        )
 
-    RETURN_TYPES = (IO.BOOLEAN,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/primitive"
-
-    def execute(self, value: bool) -> tuple[bool]:
-        return (value,)
+    @classmethod
+    def execute(cls, value: int) -> io.NodeOutput:
+        return io.NodeOutput(value)
 
 
-NODE_CLASS_MAPPINGS = {
-    "PrimitiveString": String,
-    "PrimitiveStringMultiline": StringMultiline,
-    "PrimitiveInt": Int,
-    "PrimitiveFloat": Float,
-    "PrimitiveBoolean": Boolean,
-}
+class Float(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PrimitiveFloat",
+            display_name="Float",
+            category="utils/primitive",
+            inputs=[
+                io.Float.Input("value", min=-sys.maxsize, max=sys.maxsize),
+            ],
+            outputs=[io.Float.Output()],
+        )
 
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "PrimitiveString": "String",
-    "PrimitiveStringMultiline": "String (Multiline)",
-    "PrimitiveInt": "Int",
-    "PrimitiveFloat": "Float",
-    "PrimitiveBoolean": "Boolean",
-}
+    @classmethod
+    def execute(cls, value: float) -> io.NodeOutput:
+        return io.NodeOutput(value)
+
+
+class Boolean(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="PrimitiveBoolean",
+            display_name="Boolean",
+            category="utils/primitive",
+            inputs=[
+                io.Boolean.Input("value"),
+            ],
+            outputs=[io.Boolean.Output()],
+        )
+
+    @classmethod
+    def execute(cls, value: bool) -> io.NodeOutput:
+        return io.NodeOutput(value)
+
+
+class PrimitivesExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            String,
+            StringMultiline,
+            Int,
+            Float,
+            Boolean,
+        ]
+
+async def comfy_entrypoint() -> PrimitivesExtension:
+    return PrimitivesExtension()

comfy_extras/nodes_stable_cascade.py

@@ -17,55 +17,61 @@
 """
 
 import torch
-import nodes
+from typing_extensions import override
+
 import comfy.utils
+import nodes
+from comfy_api.latest import ComfyExtension, io
 
 
-class StableCascade_EmptyLatentImage:
-    def __init__(self, device="cpu"):
-        self.device = device
+class StableCascade_EmptyLatentImage(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StableCascade_EmptyLatentImage",
+            category="latent/stable_cascade",
+            inputs=[
+                io.Int.Input("width", default=1024, min=256, max=nodes.MAX_RESOLUTION, step=8),
+                io.Int.Input("height", default=1024, min=256, max=nodes.MAX_RESOLUTION, step=8),
+                io.Int.Input("compression", default=42, min=4, max=128, step=1),
+                io.Int.Input("batch_size", default=1, min=1, max=4096),
+            ],
+            outputs=[
+                io.Latent.Output(display_name="stage_c"),
+                io.Latent.Output(display_name="stage_b"),
+            ],
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "width": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}),
-            "height": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}),
-            "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}),
-            "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})
-        }}
-    RETURN_TYPES = ("LATENT", "LATENT")
-    RETURN_NAMES = ("stage_c", "stage_b")
-    FUNCTION = "generate"
-
-    CATEGORY = "latent/stable_cascade"
-
-    def generate(self, width, height, compression, batch_size=1):
+    def execute(cls, width, height, compression, batch_size=1):
         c_latent = torch.zeros([batch_size, 16, height // compression, width // compression])
         b_latent = torch.zeros([batch_size, 4, height // 4, width // 4])
-        return ({
+        return io.NodeOutput({
             "samples": c_latent,
         }, {
             "samples": b_latent,
         })
 
-class StableCascade_StageC_VAEEncode:
-    def __init__(self, device="cpu"):
-        self.device = device
+
+class StableCascade_StageC_VAEEncode(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StableCascade_StageC_VAEEncode",
+            category="latent/stable_cascade",
+            inputs=[
+                io.Image.Input("image"),
+                io.Vae.Input("vae"),
+                io.Int.Input("compression", default=42, min=4, max=128, step=1),
+            ],
+            outputs=[
+                io.Latent.Output(display_name="stage_c"),
+                io.Latent.Output(display_name="stage_b"),
+            ],
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "image": ("IMAGE",),
-            "vae": ("VAE", ),
-            "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}),
-        }}
-    RETURN_TYPES = ("LATENT", "LATENT")
-    RETURN_NAMES = ("stage_c", "stage_b")
-    FUNCTION = "generate"
-
-    CATEGORY = "latent/stable_cascade"
-
-    def generate(self, image, vae, compression):
+    def execute(cls, image, vae, compression):
         width = image.shape[-2]
         height = image.shape[-3]
         out_width = (width // compression) * vae.downscale_ratio
@@ -75,51 +81,59 @@ class StableCascade_StageC_VAEEncode:
 
         c_latent = vae.encode(s[:,:,:,:3])
         b_latent = torch.zeros([c_latent.shape[0], 4, (height // 8) * 2, (width // 8) * 2])
-        return ({
+        return io.NodeOutput({
             "samples": c_latent,
         }, {
             "samples": b_latent,
         })
 
-class StableCascade_StageB_Conditioning:
+
+class StableCascade_StageB_Conditioning(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "conditioning": ("CONDITIONING",),
-                              "stage_c": ("LATENT",),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StableCascade_StageB_Conditioning",
+            category="conditioning/stable_cascade",
+            inputs=[
+                io.Conditioning.Input("conditioning"),
+                io.Latent.Input("stage_c"),
+            ],
+            outputs=[
+                io.Conditioning.Output(),
+            ],
+        )
 
-    FUNCTION = "set_prior"
-
-    CATEGORY = "conditioning/stable_cascade"
-
-    def set_prior(self, conditioning, stage_c):
+    @classmethod
+    def execute(cls, conditioning, stage_c):
         c = []
         for t in conditioning:
             d = t[1].copy()
-            d['stable_cascade_prior'] = stage_c['samples']
+            d["stable_cascade_prior"] = stage_c["samples"]
             n = [t[0], d]
             c.append(n)
-        return (c, )
+        return io.NodeOutput(c)
 
-class StableCascade_SuperResolutionControlnet:
-    def __init__(self, device="cpu"):
-        self.device = device
+
+class StableCascade_SuperResolutionControlnet(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StableCascade_SuperResolutionControlnet",
+            category="_for_testing/stable_cascade",
+            is_experimental=True,
+            inputs=[
+                io.Image.Input("image"),
+                io.Vae.Input("vae"),
+            ],
+            outputs=[
+                io.Image.Output(display_name="controlnet_input"),
+                io.Latent.Output(display_name="stage_c"),
+                io.Latent.Output(display_name="stage_b"),
+            ],
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "image": ("IMAGE",),
-            "vae": ("VAE", ),
-        }}
-    RETURN_TYPES = ("IMAGE", "LATENT", "LATENT")
-    RETURN_NAMES = ("controlnet_input", "stage_c", "stage_b")
-    FUNCTION = "generate"
-
-    EXPERIMENTAL = True
-    CATEGORY = "_for_testing/stable_cascade"
-
-    def generate(self, image, vae):
+    def execute(cls, image, vae):
         width = image.shape[-2]
         height = image.shape[-3]
         batch_size = image.shape[0]
@@ -127,15 +141,22 @@ class StableCascade_SuperResolutionControlnet:
 
         c_latent = torch.zeros([batch_size, 16, height // 16, width // 16])
         b_latent = torch.zeros([batch_size, 4, height // 2, width // 2])
-        return (controlnet_input, {
+        return io.NodeOutput(controlnet_input, {
             "samples": c_latent,
         }, {
             "samples": b_latent,
         })
 
-NODE_CLASS_MAPPINGS = {
-    "StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage,
-    "StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning,
-    "StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode,
-    "StableCascade_SuperResolutionControlnet": StableCascade_SuperResolutionControlnet,
-}
+
+class StableCascadeExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            StableCascade_EmptyLatentImage,
+            StableCascade_StageB_Conditioning,
+            StableCascade_StageC_VAEEncode,
+            StableCascade_SuperResolutionControlnet,
+        ]
+
+async def comfy_entrypoint() -> StableCascadeExtension:
+    return StableCascadeExtension()

comfy_extras/nodes_video.py

@@ -5,52 +5,49 @@ import av
 import torch
 import folder_paths
 import json
-from typing import Optional, Literal
+from typing import Optional
+from typing_extensions import override
 from fractions import Fraction
-from comfy.comfy_types import IO, FileLocator, ComfyNodeABC
-from comfy_api.latest import Input, InputImpl, Types
+from comfy_api.input import AudioInput, ImageInput, VideoInput
+from comfy_api.input_impl import VideoFromComponents, VideoFromFile
+from comfy_api.util import VideoCodec, VideoComponents, VideoContainer
+from comfy_api.latest import ComfyExtension, io, ui
 from comfy.cli_args import args
 
-class SaveWEBM:
-    def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
+class SaveWEBM(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SaveWEBM",
+            category="image/video",
+            is_experimental=True,
+            inputs=[
+                io.Image.Input("images"),
+                io.String.Input("filename_prefix", default="ComfyUI"),
+                io.Combo.Input("codec", options=["vp9", "av1"]),
+                io.Float.Input("fps", default=24.0, min=0.01, max=1000.0, step=0.01),
+                io.Float.Input("crf", default=32.0, min=0, max=63.0, step=1, tooltip="Higher crf means lower quality with a smaller file size, lower crf means higher quality higher filesize."),
+            ],
+            outputs=[],
+            hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
+            is_output_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"}),
-                     "codec": (["vp9", "av1"],),
-                     "fps": ("FLOAT", {"default": 24.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
-                     "crf": ("FLOAT", {"default": 32.0, "min": 0, "max": 63.0, "step": 1, "tooltip": "Higher crf means lower quality with a smaller file size, lower crf means higher quality higher filesize."}),
-                     },
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_images"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image/video"
-
-    EXPERIMENTAL = True
-
-    def save_images(self, images, codec, fps, filename_prefix, crf, prompt=None, extra_pnginfo=None):
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
+    def execute(cls, images, codec, fps, filename_prefix, crf) -> io.NodeOutput:
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
+            filename_prefix, folder_paths.get_output_directory(), images[0].shape[1], images[0].shape[0]
+        )
 
         file = f"{filename}_{counter:05}_.webm"
         container = av.open(os.path.join(full_output_folder, file), mode="w")
 
-        if prompt is not None:
-            container.metadata["prompt"] = json.dumps(prompt)
+        if cls.hidden.prompt is not None:
+            container.metadata["prompt"] = json.dumps(cls.hidden.prompt)
 
-        if extra_pnginfo is not None:
-            for x in extra_pnginfo:
-                container.metadata[x] = json.dumps(extra_pnginfo[x])
+        if cls.hidden.extra_pnginfo is not None:
+            for x in cls.hidden.extra_pnginfo:
+                container.metadata[x] = json.dumps(cls.hidden.extra_pnginfo[x])
 
         codec_map = {"vp9": "libvpx-vp9", "av1": "libsvtav1"}
         stream = container.add_stream(codec_map[codec], rate=Fraction(round(fps * 1000), 1000))
@@ -69,63 +66,46 @@ class SaveWEBM:
         container.mux(stream.encode())
         container.close()
 
-        results: list[FileLocator] = [{
-            "filename": file,
-            "subfolder": subfolder,
-            "type": self.type
-        }]
+        return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
 
-        return {"ui": {"images": results, "animated": (True,)}}  # TODO: frontend side
-
-class SaveVideo(ComfyNodeABC):
-    def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
-        self.type: Literal["output"] = "output"
-        self.prefix_append = ""
+class SaveVideo(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SaveVideo",
+            display_name="Save Video",
+            category="image/video",
+            description="Saves the input images to your ComfyUI output directory.",
+            inputs=[
+                io.Video.Input("video", tooltip="The video to save."),
+                io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),
+                io.Combo.Input("format", options=VideoContainer.as_input(), default="auto", tooltip="The format to save the video as."),
+                io.Combo.Input("codec", options=VideoCodec.as_input(), default="auto", tooltip="The codec to use for the video."),
+            ],
+            outputs=[],
+            hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
+            is_output_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "video": (IO.VIDEO, {"tooltip": "The video to save."}),
-                "filename_prefix": ("STRING", {"default": "video/ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}),
-                "format": (Types.VideoContainer.as_input(), {"default": "auto", "tooltip": "The format to save the video as."}),
-                "codec": (Types.VideoCodec.as_input(), {"default": "auto", "tooltip": "The codec to use for the video."}),
-            },
-            "hidden": {
-                "prompt": "PROMPT",
-                "extra_pnginfo": "EXTRA_PNGINFO"
-            },
-        }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_video"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image/video"
-    DESCRIPTION = "Saves the input images to your ComfyUI output directory."
-
-    def save_video(self, video: Input.Video, filename_prefix, format, codec, prompt=None, extra_pnginfo=None):
-        filename_prefix += self.prefix_append
+    def execute(cls, video: VideoInput, filename_prefix, format, codec) -> io.NodeOutput:
         width, height = video.get_dimensions()
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
             filename_prefix,
-            self.output_dir,
+            folder_paths.get_output_directory(),
             width,
             height
         )
-        results: list[FileLocator] = list()
         saved_metadata = None
         if not args.disable_metadata:
             metadata = {}
-            if extra_pnginfo is not None:
-                metadata.update(extra_pnginfo)
-            if prompt is not None:
-                metadata["prompt"] = prompt
+            if cls.hidden.extra_pnginfo is not None:
+                metadata.update(cls.hidden.extra_pnginfo)
+            if cls.hidden.prompt is not None:
+                metadata["prompt"] = cls.hidden.prompt
             if len(metadata) > 0:
                 saved_metadata = metadata
-        file = f"{filename}_{counter:05}_.{Types.VideoContainer.get_extension(format)}"
+        file = f"{filename}_{counter:05}_.{VideoContainer.get_extension(format)}"
         video.save_to(
             os.path.join(full_output_folder, file),
             format=format,
@@ -133,83 +113,82 @@ class SaveVideo(ComfyNodeABC):
             metadata=saved_metadata
         )
 
-        results.append({
-            "filename": file,
-            "subfolder": subfolder,
-            "type": self.type
-        })
-        counter += 1
+        return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
 
-        return { "ui": { "images": results, "animated": (True,) } }
 
-class CreateVideo(ComfyNodeABC):
+class CreateVideo(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "images": (IO.IMAGE, {"tooltip": "The images to create a video from."}),
-                "fps": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 120.0, "step": 1.0}),
-            },
-            "optional": {
-                "audio": (IO.AUDIO, {"tooltip": "The audio to add to the video."}),
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="CreateVideo",
+            display_name="Create Video",
+            category="image/video",
+            description="Create a video from images.",
+            inputs=[
+                io.Image.Input("images", tooltip="The images to create a video from."),
+                io.Float.Input("fps", default=30.0, min=1.0, max=120.0, step=1.0),
+                io.Audio.Input("audio", optional=True, tooltip="The audio to add to the video."),
+            ],
+            outputs=[
+                io.Video.Output(),
+            ],
+        )
 
-    RETURN_TYPES = (IO.VIDEO,)
-    FUNCTION = "create_video"
-
-    CATEGORY = "image/video"
-    DESCRIPTION = "Create a video from images."
-
-    def create_video(self, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None):
-        return (InputImpl.VideoFromComponents(
-            Types.VideoComponents(
-            images=images,
-            audio=audio,
-            frame_rate=Fraction(fps),
-            )
-        ),)
-
-class GetVideoComponents(ComfyNodeABC):
     @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "video": (IO.VIDEO, {"tooltip": "The video to extract components from."}),
-            }
-        }
-    RETURN_TYPES = (IO.IMAGE, IO.AUDIO, IO.FLOAT)
-    RETURN_NAMES = ("images", "audio", "fps")
-    FUNCTION = "get_components"
+    def execute(cls, images: ImageInput, fps: float, audio: Optional[AudioInput] = None) -> io.NodeOutput:
+        return io.NodeOutput(
+            VideoFromComponents(VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps)))
+        )
 
-    CATEGORY = "image/video"
-    DESCRIPTION = "Extracts all components from a video: frames, audio, and framerate."
+class GetVideoComponents(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="GetVideoComponents",
+            display_name="Get Video Components",
+            category="image/video",
+            description="Extracts all components from a video: frames, audio, and framerate.",
+            inputs=[
+                io.Video.Input("video", tooltip="The video to extract components from."),
+            ],
+            outputs=[
+                io.Image.Output(display_name="images"),
+                io.Audio.Output(display_name="audio"),
+                io.Float.Output(display_name="fps"),
+            ],
+        )
 
-    def get_components(self, video: Input.Video):
+    @classmethod
+    def execute(cls, video: VideoInput) -> io.NodeOutput:
         components = video.get_components()
 
-        return (components.images, components.audio, float(components.frame_rate))
+        return io.NodeOutput(components.images, components.audio, float(components.frame_rate))
 
-class LoadVideo(ComfyNodeABC):
+class LoadVideo(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls):
+    def define_schema(cls):
         input_dir = folder_paths.get_input_directory()
         files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
         files = folder_paths.filter_files_content_types(files, ["video"])
-        return {"required":
-                    {"file": (sorted(files), {"video_upload": True})},
-                }
-
-    CATEGORY = "image/video"
-
-    RETURN_TYPES = (IO.VIDEO,)
-    FUNCTION = "load_video"
-    def load_video(self, file):
-        video_path = folder_paths.get_annotated_filepath(file)
-        return (InputImpl.VideoFromFile(video_path),)
+        return io.Schema(
+            node_id="LoadVideo",
+            display_name="Load Video",
+            category="image/video",
+            inputs=[
+                io.Combo.Input("file", options=sorted(files), upload=io.UploadType.video),
+            ],
+            outputs=[
+                io.Video.Output(),
+            ],
+        )
 
     @classmethod
-    def IS_CHANGED(cls, file):
+    def execute(cls, file) -> io.NodeOutput:
+        video_path = folder_paths.get_annotated_filepath(file)
+        return io.NodeOutput(VideoFromFile(video_path))
+
+    @classmethod
+    def fingerprint_inputs(s, file):
         video_path = folder_paths.get_annotated_filepath(file)
         mod_time = os.path.getmtime(video_path)
         # Instead of hashing the file, we can just use the modification time to avoid
@@ -217,24 +196,23 @@ class LoadVideo(ComfyNodeABC):
         return mod_time
 
     @classmethod
-    def VALIDATE_INPUTS(cls, file):
+    def validate_inputs(s, file):
         if not folder_paths.exists_annotated_filepath(file):
             return "Invalid video file: {}".format(file)
 
         return True
 
-NODE_CLASS_MAPPINGS = {
-    "SaveWEBM": SaveWEBM,
-    "SaveVideo": SaveVideo,
-    "CreateVideo": CreateVideo,
-    "GetVideoComponents": GetVideoComponents,
-    "LoadVideo": LoadVideo,
-}
 
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "SaveVideo": "Save Video",
-    "CreateVideo": "Create Video",
-    "GetVideoComponents": "Get Video Components",
-    "LoadVideo": "Load Video",
-}
+class VideoExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            SaveWEBM,
+            SaveVideo,
+            CreateVideo,
+            GetVideoComponents,
+            LoadVideo,
+        ]
 
+async def comfy_entrypoint() -> VideoExtension:
+    return VideoExtension()

comfyui_version.py

@@ -1,3 +1,3 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
-__version__ = "0.3.56"
+__version__ = "0.3.57"

main.py

@@ -113,7 +113,6 @@ import gc
 
 if os.name == "nt":
     os.environ['MIMALLOC_PURGE_DELAY'] = '0'
-    logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
 
 if __name__ == "__main__":
     if args.default_device is not None:

nodes.py

@@ -2344,6 +2344,7 @@ async def init_builtin_api_nodes():
         "nodes_veo2.py",
         "nodes_kling.py",
         "nodes_bfl.py",
+        "nodes_bytedance.py",
         "nodes_luma.py",
         "nodes_recraft.py",
         "nodes_pixverse.py",

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.3.56"
+version = "0.3.57"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.9"

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.25.11
-comfyui-workflow-templates==0.1.70
+comfyui-workflow-templates==0.1.75
 comfyui-embedded-docs==0.2.6
 torch
 torchsde

script_examples/basic_api_example.py

@@ -3,11 +3,7 @@ from urllib import request
 
 #This is the ComfyUI api prompt format.
 
-#If you want it for a specific workflow you can "enable dev mode options"
-#in the settings of the UI (gear beside the "Queue Size: ") this will enable
-#a button on the UI to save workflows in api format.
-
-#keep in mind ComfyUI is pre alpha software so this format will change a bit.
+#If you want it for a specific workflow you can "File -> Export (API)" in the interface.
 
 #this is the one for the default workflow
 prompt_text = """

server.py

@@ -729,7 +729,34 @@ class PromptServer():
 
         @routes.post("/interrupt")
         async def post_interrupt(request):
-            nodes.interrupt_processing()
+            try:
+                json_data = await request.json()
+            except json.JSONDecodeError:
+                json_data = {}
+
+            # Check if a specific prompt_id was provided for targeted interruption
+            prompt_id = json_data.get('prompt_id')
+            if prompt_id:
+                currently_running, _ = self.prompt_queue.get_current_queue()
+
+                # Check if the prompt_id matches any currently running prompt
+                should_interrupt = False
+                for item in currently_running:
+                    # item structure: (number, prompt_id, prompt, extra_data, outputs_to_execute)
+                    if item[1] == prompt_id:
+                        logging.info(f"Interrupting prompt {prompt_id}")
+                        should_interrupt = True
+                        break
+
+                if should_interrupt:
+                    nodes.interrupt_processing()
+                else:
+                    logging.info(f"Prompt {prompt_id} is not currently running, skipping interrupt")
+            else:
+                # No prompt_id provided, do a global interrupt
+                logging.info("Global interrupt (no prompt_id specified)")
+                nodes.interrupt_processing()
+
             return web.Response(status=200)
 
         @routes.post("/free")