Skip to content

Commit

Permalink
Add Batch Loader node #13
Browse files Browse the repository at this point in the history
  • Loading branch information
receyuki committed Nov 23, 2023
1 parent 4f42729 commit f56c1e0
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 7 deletions.
36 changes: 36 additions & 0 deletions js/loaderDisplay.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
import {app} from "../../scripts/app.js";
import {ComfyWidgets} from "../../scripts/widgets.js";

// Create a read-only string widget with opacity set
function createWidget(app, node, widgetName) {
const widget = ComfyWidgets["STRING"](node, widgetName, ["STRING", {multiline: true}], app).widget;
widget.inputEl.readOnly = true;
widget.inputEl.style.opacity = 0.7;
return widget;
}

// Displays file list on the node
app.registerExtension({
name: "sd_prompt_reader.loaderDisplay",
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeData.name === "SDBatchLoader") {
const onNodeCreated = nodeType.prototype.onNodeCreated;

nodeType.prototype.onNodeCreated = function () {
const result = onNodeCreated?.apply(this, arguments);

// Create prompt and setting widgets
const fileList = createWidget(app, this, "fileList");
return result;
};

// Update widgets
const onExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
onExecuted?.apply(this, arguments);
this.widgets.find(obj => obj.name === "fileList").value = message.text[0];

};
}
},
});
103 changes: 96 additions & 7 deletions nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,18 +53,34 @@ def output_to_terminal(text: str):
output_to_terminal("Core version: " + CORE_VERSION)


class AnyType(str):
"""A special type that can be connected to any other types. Credit to pythongosssss"""

def __ne__(self, __value: object) -> bool:
return False


any_type = AnyType("*")


class SDPromptReader:
files = []

@classmethod
def INPUT_TYPES(s):
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))
]
SDPromptReader.files = sorted(
[
f
for f in os.listdir(input_dir)
if os.path.isfile(os.path.join(input_dir, f))
]
)
return {
"required": {
"image": (sorted(files), {"image_upload": True}),
"image": (SDPromptReader.files, {"image_upload": True}),
},
"optional": {
"parameter_index": (
"INT",
{"default": 0, "min": 0, "max": 255, "step": 1},
Expand Down Expand Up @@ -104,7 +120,10 @@ def INPUT_TYPES(s):
OUTPUT_NODE = True

def load_image(self, image, parameter_index):
image_path = folder_paths.get_annotated_filepath(image)
if image in SDPromptReader.files:
image_path = folder_paths.get_annotated_filepath(image)
else:
image_path = image
i = Image.open(image_path)
i = ImageOps.exif_transpose(i)
image = i.convert("RGB")
Expand Down Expand Up @@ -778,12 +797,81 @@ def convert_string(
)


class SDBatchLoader:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"path": ("STRING", {"default": "./input/"}),
},
"optional": {
"image_load_limit": ("INT", {"default": 0, "min": 0, "step": 1}),
"start_index": ("INT", {"default": 0, "min": 0, "step": 1}),
},
}

RETURN_TYPES = (any_type,)

RETURN_NAMES = ("IMAGE",)
OUTPUT_IS_LIST = (True,)
FUNCTION = "load_path"
CATEGORY = "SD Prompt Reader"

def load_path(
self,
path: str = "./input/",
image_load_limit: int = 0,
start_index: int = 0,
):
if not Path(path).is_dir():
raise FileNotFoundError(f"Invalid directory: {path}")

files = list(
filter(lambda file: file.suffix in SUPPORTED_FORMATS, Path(path).iterdir())
)

files = (
sorted(files)[start_index : start_index + image_load_limit]
if image_load_limit > 0
else sorted(files)[start_index:]
)

files_str = list(map(str, files))
return {
"ui": {
"text": ("\n".join(files_str),),
},
"result": (files_str,),
}

@classmethod
def IS_CHANGED(
s,
path,
image_load_limit,
start_index,
):
return os.listdir(path)

@classmethod
def VALIDATE_INPUTS(
s,
path,
image_load_limit,
start_index,
):
if not Path(path).is_dir():
return f"Invalid directory: {path}"
return True


NODE_CLASS_MAPPINGS = {
"SDPromptReader": SDPromptReader,
"SDPromptSaver": SDPromptSaver,
"SDParameterGenerator": SDParameterGenerator,
"SDPromptMerger": SDPromptMerger,
"SDTypeConverter": SDTypeConverter,
"SDBatchLoader": SDBatchLoader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
Expand All @@ -792,4 +880,5 @@ def convert_string(
"SDParameterGenerator": "SD Parameter Generator",
"SDPromptMerger": "SD Prompt Merger",
"SDTypeConverter": "SD Type Converter",
"SDBatchLoader": "SD Batch Loader",
}

0 comments on commit f56c1e0

Please sign in to comment.