From 7cc91cddfc0ca8ab4bac06ac5489f3bedaffe45d Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Wed, 30 Oct 2024 09:48:00 +1300 Subject: [PATCH] working on adding tools --- __init__.py | 2 + js/ConfigurationNodes.js | 1 + js/gtUINodes.js | 18 ++++++- js/gtUIUtils.js | 28 ++++++++++ nodes/lora/conductor_tool/manifest.yml | 5 ++ nodes/lora/conductor_tool/requirements.txt | 1 + nodes/lora/conductor_tool/tool.py | 36 +++++++++++++ nodes/lora/gtUILoRADownloadJob.py | 34 ++++++++++++ nodes/lora/gtUILoRATrainingTask.py | 62 +++++++++++++++------- nodes/tasks/gtUIImageCaptionTask.py | 2 +- pyproject.toml | 1 + 11 files changed, 169 insertions(+), 21 deletions(-) create mode 100644 nodes/lora/conductor_tool/manifest.yml create mode 100644 nodes/lora/conductor_tool/requirements.txt create mode 100644 nodes/lora/conductor_tool/tool.py create mode 100644 nodes/lora/gtUILoRADownloadJob.py diff --git a/__init__.py b/__init__.py index 650322a..0c33f62 100644 --- a/__init__.py +++ b/__init__.py @@ -223,6 +223,7 @@ from .nodes.loaders.gtUILoadText import gtUILoadText # LORA from .nodes.lora.gtUILoRAConfig import gtUILoRAConfig +from .nodes.lora.gtUILoRADownloadJob import gtUILoRADownloadJob from .nodes.lora.gtUILoRATrainingTask import gtUILoRATrainingTask from .nodes.rag.gtUIFootnotePromptResponseRagModule import ( gtUIFootnotePromptResponseRagModule, @@ -480,6 +481,7 @@ NODE_CLASS_MAPPINGS = { "Griptape LoRA: Configuration": gtUILoRAConfig, "Griptape LoRA: Image Captions": gtUIImageCaptionTask, "Griptape LoRA: Train using Conductor": gtUILoRATrainingTask, + "Griptape LoRA: Download Job": gtUILoRADownloadJob, } diff --git a/js/ConfigurationNodes.js b/js/ConfigurationNodes.js index 1fa32b2..54f9f30 100644 --- a/js/ConfigurationNodes.js +++ b/js/ConfigurationNodes.js @@ -8,6 +8,7 @@ export function setupConfigurationNodes(nodeType, nodeData, app) { (nodeData.name.includes("Griptape Embedding Driver")) || (nodeData.name.includes("Griptape LoRA: Configuration")) || (nodeData.name.includes("Griptape LoRA: Train")) || + (nodeData.name.includes("Griptape LoRA: Download")) || (nodeData.name.includes("Griptape Prompt Driver"))) { const onNodeCreated = nodeType.prototype.onNodeCreated; diff --git a/js/gtUINodes.js b/js/gtUINodes.js index 71f8c4b..8f29c39 100644 --- a/js/gtUINodes.js +++ b/js/gtUINodes.js @@ -7,7 +7,7 @@ import { setupDisplayNodes } from "./DisplayNodes.js"; import { setupCombineNodes } from "./CombineNodes.js"; import { setupExtractionNodes } from "./ExtractionNodes.js"; import { setupTextLoaderModuleNodes } from "./TextLoaderModuleNodes.js"; -import { gtUIAddUploadWidget } from "./gtUIUtils.js"; +import { gtUIAddUploadWidget, gtUIAddButtonWidget, gtUIAddUrlButtonWidget } from "./gtUIUtils.js"; import { setupMenuSeparator } from "./gtUIMenuSeparator.js"; // app.extensionManager.registerSidebarTab({ // id: "search", @@ -20,6 +20,12 @@ import { setupMenuSeparator } from "./gtUIMenuSeparator.js"; // }, // }); +const createUrlCallback = function(url) { + return function() { + window.open(url, "_blank"); // Opens the provided URL in a new tab + }; +} + app.registerExtension({ name: "comfy.gtUI", beforeConfigureGraph: (graphData, missingNodeTypes) => { @@ -71,5 +77,15 @@ app.registerExtension({ if (nodeData.name === "Griptape Load: Text") { gtUIAddUploadWidget(nodeType, nodeData, "text", "text") } + + // Add Conductor Dashboard Widgets + // add a button that will open a url in a new tab + + if (nodeData.name === "Griptape LoRA: Train using Conductor") { + gtUIAddButtonWidget(nodeType, "Open Conductor Dashboard", createUrlCallback("https://dashboard.conductortech.com/")) + } + if (nodeData.name === "Griptape LoRA: Download Job") { + gtUIAddUrlButtonWidget(nodeType, "Open Conductor Job Dashboard", "https://dashboard.conductortech.com/job", "job_id") + } } }); diff --git a/js/gtUIUtils.js b/js/gtUIUtils.js index 42a077b..ac56525 100644 --- a/js/gtUIUtils.js +++ b/js/gtUIUtils.js @@ -109,7 +109,35 @@ export async function uploadFile(file) { alert(error); } } +const createUrlCallbackWithWidget = function(baseurl, widgetName, node) { + return function() { + // Find the widget by its name when the button is clicked + const pathWidget = node.widgets.find((w) => w.name === widgetName); + console.log(pathWidget); + if (pathWidget) { + let value = pathWidget.value; // Get the current value of the widget + const url = baseurl + "/" + value; // Construct the URL + window.open(url, "_blank"); // Open the URL in a new tab + } else { + console.error(`Widget with name ${widgetName} not found`); + } + }; + } +export function gtUIAddUrlButtonWidget(nodeType, buttonText, baseurl, widgetName) { + chainCallback(nodeType.prototype, "onNodeCreated", function() { + console.log("Adding button widget"); + console.log(widgetName); + const buttonCallback = createUrlCallbackWithWidget(baseurl, widgetName, this); + let buttonWidget = this.addWidget("button", buttonText, null , buttonCallback); + } + ); +} +export function gtUIAddButtonWidget(nodeType, buttonText, buttonCallback) { + chainCallback(nodeType.prototype, "onNodeCreated", function() { + let buttonWidget = this.addWidget("button", buttonText, null , buttonCallback); + }); +} export function gtUIAddUploadWidget(nodeType, nodeData, widgetName, type="audio") { chainCallback(nodeType.prototype, "onNodeCreated", function() { const pathWidget = this.widgets.find((w) => w.name === widgetName); diff --git a/nodes/lora/conductor_tool/manifest.yml b/nodes/lora/conductor_tool/manifest.yml new file mode 100644 index 0000000..99c61e2 --- /dev/null +++ b/nodes/lora/conductor_tool/manifest.yml @@ -0,0 +1,5 @@ +version: "v1" +name: Conductor Tool +description: Tool for working Conductor +contact_email: jason@griptape.ai +legal_info_url: https://www.griptape.ai/legal \ No newline at end of file diff --git a/nodes/lora/conductor_tool/requirements.txt b/nodes/lora/conductor_tool/requirements.txt new file mode 100644 index 0000000..90ed1c7 --- /dev/null +++ b/nodes/lora/conductor_tool/requirements.txt @@ -0,0 +1 @@ +ciocore \ No newline at end of file diff --git a/nodes/lora/conductor_tool/tool.py b/nodes/lora/conductor_tool/tool.py new file mode 100644 index 0000000..01b3632 --- /dev/null +++ b/nodes/lora/conductor_tool/tool.py @@ -0,0 +1,36 @@ +import json +from copy import deepcopy + +from ciocore import conductor_submit +from griptape.artifacts import TextArtifact +from griptape.tools import BaseTool +from griptape.utils.decorators import activity +from schema import Literal, Schema + + +class ConductorTool(BaseTool): + @activity( + config={ + "description": "Can be used to execute a conductor submission", + "schema": Schema( + {Literal("submission", description="The submission dict"): dict} + ), + } + ) + def submit(self, params: dict) -> TextArtifact: + SUBMISSION = params["values"].get("submission") + # host_target = "blender 4.2.1.glibc217 linux" + # pt = package_tree.PackageTree(api_client.request_software_packages()) + # host_listing = pt.find_by_name(host_target) + + data = deepcopy(SUBMISSION) + # data["package_id"] = [host_listing["package_id"]] + # job_env = package_environment.PackageEnvironment( + # env_list=host_listing["environment"] + # ) + # data["environment"] = dict(job_env) + + submission = conductor_submit.Submit(data) + response, response_code = submission.main() + print(response_code) + return TextArtifact(str(json.dumps(response))) diff --git a/nodes/lora/gtUILoRADownloadJob.py b/nodes/lora/gtUILoRADownloadJob.py new file mode 100644 index 0000000..f31b2c1 --- /dev/null +++ b/nodes/lora/gtUILoRADownloadJob.py @@ -0,0 +1,34 @@ +# from .conductor_tool.tool import ConductorTool + + +class gtUILoRADownloadJob: + DESCRIPTION = "Griptape LoRA: Download Job" + + @classmethod + def INPUT_TYPES(s): + return { + "required": {}, + "optional": { + "job_id": ("STRING", {"default": "0000"}), + "status": ( + "STRING", + {"default": ""}, + ), + "links_comment": ( + "STRING", + {"default": "Links"}, + ), + }, + } + + RETURN_TYPES = ("IMAGE", "MODEL", "STRING") + RETURN_NAMES = ("IMAGE", "LORA", "LOGS") + + FUNCTION = "create" + + CATEGORY = "Griptape/LoRA" + + def create(self, **kwargs): + job_id = kwargs.get("job_id") + status = "Succeeded" + return {"ui": {"status": status}, "result": (None, None, "")} diff --git a/nodes/lora/gtUILoRATrainingTask.py b/nodes/lora/gtUILoRATrainingTask.py index e0a071b..9cff1b0 100644 --- a/nodes/lora/gtUILoRATrainingTask.py +++ b/nodes/lora/gtUILoRATrainingTask.py @@ -5,6 +5,11 @@ import re import shutil import folder_paths +from griptape.structures import Pipeline +from griptape.tasks import ToolTask +from griptape.tools import CalculatorTool + +# from .conductor_tool.tool import ConductorTool instance_types = ["A5000", "A100"] @@ -70,6 +75,14 @@ class gtUILoRATrainingTask: "lora_config": ("LORA_CONFIG", {}), "image": ("IMAGE", {}), "captions": ("STRING", {"forceInput": True}), + "job_setting_comment": ( + "STRING", + {"default": "Queue Settings"}, + ), + "start_job_or_get_status": ( + ["Start New Job", "Get Status"], + {"default": "Start New Job"}, + ), "job_settings_comment": ( "STRING", {"default": "Job Settings"}, @@ -91,16 +104,15 @@ class gtUILoRATrainingTask: "STRING", {"default": folder_paths.get_output_directory()}, ), - "status_comment": ( + "commands_comment": ( "STRING", - {"default": "Status"}, + {"default": "Links"}, ), - "status": ("STRING", {"multiline": True}), }, } - RETURN_TYPES = ("IMAGE", "MODEL", "STRING") - RETURN_NAMES = ("IMAGE", "MODEL", "LOGS") + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("JOB_ID", "LOGS") FUNCTION = "create" @@ -114,8 +126,9 @@ class gtUILoRATrainingTask: job_title = kwargs.get("job_title") project = kwargs.get("project") instance_type = kwargs.get("instance_type") + instance_type = "cw-epycmilan-4-rtx4000-1" # HACK: Hardcoded for now + output_path = kwargs.get("output_path") - status = "Not started.." # Copy the images to the training folder inputs = folder_paths.get_input_directory() @@ -172,35 +185,46 @@ class gtUILoRATrainingTask: --save_every_n_epochs=1 --network_module=networks.lora """ + SCENE_FILE_ROOT = os.path.dirname(os.path.abspath(__file__)) - payload = { + SUBMISSION = { "job_title": job_title, "project": project, "instance_type": instance_type, "software_package_ids": [ - "7be1b2410b3f93b2a2889f6ce191d4e1" + # "7be1b2410b3f93b2a2889f6ce191d4e1" ], # will need pytorch, flux, etc "force": False, "local_upload": True, "preemptible": True, - "autoretry_policy": {"preempted": {"max_retries": 3}}, + "autoretry_policy": None, "output_path": output_path, # local folder where whatever we're pulling is going to be - "environment": { - "PATH": "/opt/blenderfoundation/blender/2/blender-2.93.0-glibc217", - "CONDUCTOR_PATHHELPER": "0", - }, + "environment": {}, "upload_paths": [ - "/projects/training/images", - "/projects/training/checkpoints", - "/projects/training/loras", + # f"{SCENE_FILE_ROOT}/projects/training/images", + # f"{SCENE_FILE_ROOT}/projects/training/checkpoints", + # f"{SCENE_FILE_ROOT}/projects/training/loras", ], + "scout_frames": "1", "tasks_data": [ { - "command": 'blender -b "/projects/blender/bmw_half_turn_low.blend" -E CYCLES --render-output "/projects/blender/renders/img_" --render-frame 1..1 ', + "command": f'echo "Submitted a LORA training job!\n{lora_config}"', "frames": "1", } ], - "config": lora_config, + # "config": lora_config, } - return {"ui": {"status": status}, "result": (image, model, payload)} + pipeline = Pipeline() + conductor_task = ToolTask( + "Calculate 5*2", + # f"submit a render job: {SUBMISSION}", + # tool=ConductorTool(), + tool=CalculatorTool(), + ) + pipeline.add_task(conductor_task) + + response = pipeline.run() + print(response) + status = str(response.output_task.output) + return {"ui": {"status": status}, "result": (status, SUBMISSION)} diff --git a/nodes/tasks/gtUIImageCaptionTask.py b/nodes/tasks/gtUIImageCaptionTask.py index b46e4a7..5868544 100644 --- a/nodes/tasks/gtUIImageCaptionTask.py +++ b/nodes/tasks/gtUIImageCaptionTask.py @@ -93,7 +93,7 @@ class gtUIImageCaptionTask(gtUIBaseImageTask): RETURN_NAMES = ("OUTPUT",) CATEGORY = "Griptape/LoRA" - OUTPUT_NODE = True + # OUTPUT_NODE = True def run(self, **kwargs): image = kwargs.get("image") diff --git a/pyproject.toml b/pyproject.toml index 0dfb3c5..e3bdd19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ readme = "README.md" python = "^3.11" python-dotenv = "^1.0.1" griptape = { version = "^0.33.1", extras = ["all"]} +ciocore = "9.0.0b2" [tool.poetry.group.dev.dependencies] icecream = "^2.1.3"