nodes/griptape-ai-ComfyUI-Griptape
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added training task node

Browse Source 
update config node
This commit is contained in:
Jason Schleifer
2024-10-11 05:17:39 +13:00
parent 8a1ab6ed1c
commit bb9e5f6911
6 changed files with 162 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2
View File
@@ -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,
}
js/ConfigurationNodes.js
+21 -1
View File
@@ -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
js/DisplayNodes.js
+17 -6
View File
@@ -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);
nodes/lora/gtUILoRAConfig.py
+21 -4
View File
@@ -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"),
nodes/lora/gtUILoRATrainingTask.py
+88
View File
@@ -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)}
nodes/tasks/gtUIImageCaptionTask.py
+13 -3
View File
@@ -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,),
}