working on adding tools

This commit is contained in:
Jason Schleifer
2024-10-30 09:48:00 +13:00
parent 7f45b55803
commit 7cc91cddfc
11 changed files with 169 additions and 21 deletions
+2
View File
@@ -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,
}
+1
View File
@@ -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;
+17 -1
View File
@@ -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")
}
}
});
+28
View File
@@ -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);
+5
View File
@@ -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
@@ -0,0 +1 @@
ciocore
+36
View File
@@ -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)))
+34
View File
@@ -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, "")}
+43 -19
View File
@@ -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)}
+1 -1
View File
@@ -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")
+1
View File
@@ -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"