Developing Custom Nodes in ComfyUI - Part 2
In the previous post, we set up the workspace. Now it is time to implement the actual logic. We will create a functional node that extracts specific frames from a video; depending on the user, this will be an essential tool for the workflow—simple yet important.
__init__.py
First, open the __init__.py file we created. Think of this as the "ID card" for your node. Without this, ComfyUI won't even realize your folder exists. We need to map our Python class to a node name that will appear in the UI.
from .frame_extract import FrameExtractor
#.frame_extract is the file name, and FrameExtractor is the class name of that file.
NODE_CLASS_MAPPINGS = {
#Name to be displayed on the node
"FrameExtractor": FrameExtractor
}
__all__ = ['NODE_CLASS_MAPPINGS']
frame_extract.py
Now, let's open frame_extract.py. This is where the heavy lifting happens. Since we are dealing with video frames, we need to define the inputs (the video) and the output (the specific image frame).
Below is the core structure. We define a class, set the INPUT_TYPES, and write a function to return the frame at a specific index.
import torch
import numpy as np
class FrameExtractor:
@classmethod
def INPUT_TYPES(s):
#Setting variables and types
return {
"required": {
"video_frames": ("IMAGE",),
"frame_index": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "extract"
DESCRIPTION = "Extract selected frame from video."
CATEGORY = "MyCustom/Video"
def extract(self, video_frames, frame_index):
#Logic implementation
total_frames = video_frames.shape[0]
actual_index = min(frame_index, total_frames - 1)
out_frame = video_frames[actual_index:actual_index + 1]
return (out_frame,)
Verification: Does it work?
Save the code and restart ComfyUI. If everything went well, you should be able to find "Video Frame Extractor" in the menu. It will be a truly rewarding moment when the program you created appears on the screen for the first time!
I set up a quick workflow to test it. I loaded a video, connected it to my new node, and hit 'Queue Prompt'. As you can see below, it successfully grabbed the exact frame I wanted.
Wrapping
To be honest, if you are not a developer, this stage might feel a bit daunting. There is a huge difference between using tools and building things from scratch. I wasn't a Python expert either, so I struggled quite a bit at first. However, once I got it working, I had an epiphany. I began to understand how not only my node but all the other nodes actually function.
It was a truly powerful experience. Now, instead of waiting for someone else to create a solution, I am able to solve my own problems directly.
Comments
Post a Comment