From bb9e5f6911900e6e346a19ed7efdcbbf73dda729 Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Fri, 11 Oct 2024 05:17:39 +1300 Subject: [PATCH] added training task node update config node --- __init__.py | 2 + js/ConfigurationNodes.js | 22 +++++++- js/DisplayNodes.js | 23 ++++++-- nodes/lora/gtUILoRAConfig.py | 25 ++++++-- nodes/lora/gtUILoRATrainingTask.py | 88 +++++++++++++++++++++++++++++ nodes/tasks/gtUIImageCaptionTask.py | 16 +++++- 6 files changed, 162 insertions(+), 14 deletions(-) create mode 100644 nodes/lora/gtUILoRATrainingTask.py diff --git a/__init__.py b/__init__.py index 6fd91be..9518666 100644 --- a/__init__.py +++ b/__init__.py @@ -217,6 +217,7 @@ from .nodes.loaders.gtUILoadText import gtUILoadText # LORA from .nodes.lora.gtUILoRAConfig import gtUILoRAConfig +from .nodes.lora.gtUILoRATrainingTask import gtUILoRATrainingTask from .nodes.rag.gtUIFootnotePromptResponseRagModule import ( gtUIFootnotePromptResponseRagModule, ) @@ -470,6 +471,7 @@ NODE_CLASS_MAPPINGS = { # LORA "Griptape LoRA: Configuration": gtUILoRAConfig, "Griptape LoRA: Image Captions": gtUIImageCaptionTask, + "Griptape LoRA: Train using Conductor": gtUILoRATrainingTask, } diff --git a/js/ConfigurationNodes.js b/js/ConfigurationNodes.js index f307ddd..1fa32b2 100644 --- a/js/ConfigurationNodes.js +++ b/js/ConfigurationNodes.js @@ -1,5 +1,5 @@ import { getOllamaModels, getLMStudioModels, updatePromptModelList } from "./gtUIUtils.js"; -import { commentWidget } from "./utils.js"; +import { hideWidget, showWidget, commentWidget } from "./utils.js"; export function setupConfigurationNodes(nodeType, nodeData, app) { if ((nodeData.name.includes("Griptape Agent Config")) || @@ -7,6 +7,7 @@ export function setupConfigurationNodes(nodeType, nodeData, app) { (nodeData.name.includes("Griptape Util: Create Agent Modelfile")) || (nodeData.name.includes("Griptape Embedding Driver")) || (nodeData.name.includes("Griptape LoRA: Configuration")) || + (nodeData.name.includes("Griptape LoRA: Train")) || (nodeData.name.includes("Griptape Prompt Driver"))) { const onNodeCreated = nodeType.prototype.onNodeCreated; @@ -14,6 +15,25 @@ export function setupConfigurationNodes(nodeType, nodeData, app) { nodeType.prototype.onNodeCreated = async function () { onNodeCreated?.apply(this, arguments); + // Set lora display toggle + if (nodeData.name.includes("Griptape LoRA: Configuration")) { + const widget_toggle_lora = this.widgets.find(w => w.name === 'include_lora'); + const widget_base_lora = this.widgets.find(w=> w.name === 'base_lora'); + + // Hide the widget + if (widget_toggle_lora) { + widget_toggle_lora.callback = async() => { + hideWidget(this, widget_base_lora); + + switch (widget_toggle_lora.value) { + case true: + showWidget(widget_base_lora); + break; + } + } + setTimeout(() => { widget_toggle_lora.callback() }, 5); + } + } // Ensure this.widgets exists before trying to filter it if (this.widgets && Array.isArray(this.widgets)) { // Find widgets you want to turn into comments diff --git a/js/DisplayNodes.js b/js/DisplayNodes.js index 28ea673..0c88b7e 100644 --- a/js/DisplayNodes.js +++ b/js/DisplayNodes.js @@ -37,14 +37,25 @@ export function setupDisplayNodes(nodeType, nodeData, app) { return; } - if (message.sample_output && Array.isArray(message.sample_output)) { - const sampleOutput = message.sample_output.join(''); - - // Update the sample_output widget if it exists + if (message.sample_output) { + const allOutput = message.sample_output.join(''); + let index_widget = this.widgets.find(w => w.name === "sample_index"); let sampleOutputWidget = this.widgets.find(w => w.name === "sample_output"); - if (sampleOutputWidget) { - sampleOutputWidget.value = ("Example Output:\n" + sampleOutput); + if (index_widget) { + index_widget.callback = async() => { + let index = index_widget.value; + + // get the correct index of the output + const parsedOutput = JSON.parse(allOutput); + const sampleOutput = parsedOutput["captions"]; + if (sampleOutputWidget) { + const description = sampleOutput[index]?.description || "No description available"; + sampleOutputWidget.value = ("Example Output:\n" + description); + } + } + setTimeout(() => { index_widget.callback() }, 5); } + // Update the sample_output widget if it exists } else { console.error("Unexpected sample_output format:", message.sample_output); diff --git a/nodes/lora/gtUILoRAConfig.py b/nodes/lora/gtUILoRAConfig.py index 862cb6b..3b89003 100644 --- a/nodes/lora/gtUILoRAConfig.py +++ b/nodes/lora/gtUILoRAConfig.py @@ -1,3 +1,6 @@ +import folder_paths + + class gtUILoRAConfig: DESCRIPTION = "Griptape LoRA: Configuration" @@ -21,10 +24,22 @@ class gtUILoRAConfig: return { "required": {}, "optional": { - "model": ("MODEL", {}), + "model_settings_comment": ( + "STRING", + {"default": "Model Settings"}, + ), + "base_model": ( + folder_paths.get_filename_list("checkpoints"), + {"tooltip": "The name of the checkpoint (model) to load."}, + ), + "include_llora": ("BOOLEAN", {"default": False}), + "llora_optional": ( + folder_paths.get_filename_list("loras"), + {"tooltip": "(Optional) The name of the lora to apply."}, + ), "lora_settings_comment": ( "STRING", - {"default": "LoRA Settings"}, + {"default": "New LoRA Settings"}, ), "lora_name": ("STRING", {}), "test_image_settings_comment": ( @@ -61,12 +76,14 @@ class gtUILoRAConfig: RETURN_NAMES = ("LORA_CONFIG",) FUNCTION = "create" - + OUTPUT_NODE = True CATEGORY = "Griptape/LoRA" def create(self, **kwargs): config = { - "model": kwargs.get("model"), + "model": kwargs.get("base_model"), + "include_llora": kwargs.get("include_llora"), + "lora_optional": kwargs.get("llora_optional", None), "lora_name": kwargs.get("lora_name"), "VRAM": kwargs.get("VRAM"), "repeat_trains_per_image": kwargs.get("repeat_trains_per_image"), diff --git a/nodes/lora/gtUILoRATrainingTask.py b/nodes/lora/gtUILoRATrainingTask.py new file mode 100644 index 0000000..8b06e23 --- /dev/null +++ b/nodes/lora/gtUILoRATrainingTask.py @@ -0,0 +1,88 @@ +instance_types = ["h100", "a100"] + + +class gtUILoRATrainingTask: + DESCRIPTION = "Griptape LoRA: Training Task" + + @classmethod + def INPUT_TYPES(s): + return { + "required": {}, + "optional": { + "lora_config": ("LORA_CONFIG", {}), + "image": ("IMAGE", {}), + "captions": ("STRING", {"forceInput": True}), + "job_settings_comment": ( + "STRING", + {"default": "Job Settings"}, + ), + "job_title": ("STRING", {"tooltip": "The name of the job."}), + "project": ( + "STRING", + { + "tooltip": "The name of the project, if associated with one", + "default": "Train LoRA", + }, + ), + "instance_type": (instance_types, {"default": instance_types[0]}), + "status_comment": ( + "STRING", + {"default": "Status"}, + ), + "status": ("STRING", {"multiline": True}), + }, + } + + RETURN_TYPES = ("IMAGE", "MODEL", "STRING") + RETURN_NAMES = ("IMAGE", "MODEL", "LOGS") + + FUNCTION = "create" + + CATEGORY = "Griptape/LoRA" + + def create(self, **kwargs): + lora_config = kwargs.get("lora_config") + image = kwargs.get("image") + captions = kwargs.get("captions") + instance_type = kwargs.get("instance_type") + model = "" + status = "Not started.." + logs = {} + + print(lora_config) + payload = { + "job_title": "Blender example bmw_half_turn_low.blend", + "project": "default", + "instance_type": "n1-standard-4", + "software_package_ids": ["7be1b2410b3f93b2a2889f6ce191d4e1"], + "force": False, + "local_upload": True, + "preemptible": True, + "autoretry_policy": {"preempted": {"max_retries": 3}}, + "output_path": "projects/blender/renders", + "environment": { + "PATH": "/opt/blenderfoundation/blender/2/blender-2.93.0-glibc217", + "CONDUCTOR_PATHHELPER": "0", + }, + "upload_paths": [ + "/projects/blender/bmw_half_turn_low.blend", + "/projects/blender/sourceimages/img_0001.png", + ], + "scout_frames": "1,3", + "tasks_data": [ + { + "command": 'blender -b "/projects/blender/bmw_half_turn_low.blend" -E CYCLES --render-output "/projects/blender/renders/img_" --render-frame 1..1 ', + "frames": "1", + }, + { + "command": 'blender -b "/projects/blender/bmw_half_turn_low.blend" -E CYCLES --render-output "/projects/blender/renders/img_" --render-frame 2..2 ', + "frames": "2", + }, + { + "command": 'blender -b "/projects/blender/bmw_half_turn_low.blend" -E CYCLES --render-output "/projects/blender/renders/img_" --render-frame 3..3 ', + "frames": "3", + }, + ], + } + + return {"ui": {"status": status}, "result": (image, model, payload)} diff --git a/nodes/tasks/gtUIImageCaptionTask.py b/nodes/tasks/gtUIImageCaptionTask.py index 9f5d93f..e628621 100644 --- a/nodes/tasks/gtUIImageCaptionTask.py +++ b/nodes/tasks/gtUIImageCaptionTask.py @@ -72,6 +72,13 @@ class gtUIImageCaptionTask(gtUIBaseImageTask): inputs = super().INPUT_TYPES() inputs["optional"].update( { + "sample_index": ( + "INT", + { + "default": 0, + "tooltip": "Choose which image to display a sample caption for.", + }, + ), "sample_output": ("STRING", {"multiline": True}), } ) @@ -86,6 +93,7 @@ class gtUIImageCaptionTask(gtUIBaseImageTask): RETURN_NAMES = ("OUTPUT",) CATEGORY = "Griptape/LoRA" + OUTPUT_NODE = True def run(self, **kwargs): image = kwargs.get("image") @@ -94,6 +102,7 @@ class gtUIImageCaptionTask(gtUIBaseImageTask): unique_id = kwargs.get("unique_id") prompt = kwargs.get("prompt") extra_pnginfo = kwargs.get("extra_pnginfo") + sample_index = kwargs.get("sample_index") image_filenames = find_image_filenames(unique_id, prompt) images = convert_tensor_to_base_64(image) output_example = "No example available" @@ -101,7 +110,7 @@ class gtUIImageCaptionTask(gtUIBaseImageTask): if not agent: agent = Agent() prompt_driver = agent.drivers_config.prompt_driver - prompt_text = "Describe this image" + prompt_text = "Describe this image for LoRA training." image_artifacts = [] for base64Image in images: @@ -166,12 +175,13 @@ class gtUIImageCaptionTask(gtUIBaseImageTask): workflow = Workflow(tasks=[*tasks, end_task]) result = workflow.run() output = result.output_task.output.value - output_example = json.loads(output)["captions"][0]["description"] + index = min(sample_index, len(json.loads(output)["captions"]) - 1) + output_example = json.loads(output)["captions"][index]["description"] else: output = "No image provided" # return (output,) return { - "ui": {"sample_output": output_example}, + "ui": {"sample_output": output}, "result": (output,), }