From f5d5982606a286acb34457aafeea1fb7ed4738df Mon Sep 17 00:00:00 2001 From: Menno van Leeuwen Date: Fri, 21 Mar 2025 15:33:49 +0100 Subject: [PATCH] Refactor PromptBuilder to enhance node management and input resolution --- lib/src/prompt_builder.dart | 361 +++++++++++++++++++++++++++++++++--- 1 file changed, 333 insertions(+), 28 deletions(-) diff --git a/lib/src/prompt_builder.dart b/lib/src/prompt_builder.dart index 8e88a38..64aa228 100644 --- a/lib/src/prompt_builder.dart +++ b/lib/src/prompt_builder.dart @@ -7,24 +7,11 @@ class PromptBuilder { PromptBuilder({required this.clientId}); - /// Adds a node to the workflow - void addNode(String classType, Map inputs, + /// Adds a node to the workflow and returns its nodeId + String addNode(String classType, Map inputs, {String? title, Map? outputTags}) { final nodeId = _nodeIdCounter.toString(); - // Resolve dependencies for inputs - inputs.forEach((key, value) { - if (value is List && value.length == 2 && value[1] is int) { - final outputTag = value[0]; - if (_outputToNode.containsKey(outputTag)) { - inputs[key] = [_outputToNode[outputTag]!["nodeId"], value[1]]; - } else { - throw Exception( - "Unresolved dependency: No node provides output tag '$outputTag'"); - } - } - }); - _nodes[nodeId] = { "inputs": inputs, "class_type": classType, @@ -42,13 +29,92 @@ class PromptBuilder { } _nodeIdCounter++; + return nodeId; + } + + /// Provides access to the current nodes + Map get nodes => Map.unmodifiable(_nodes); + + /// Edits an existing node in the workflow + void editNode(String nodeId, + {Map? newInputs, String? newTitle}) { + if (!_nodes.containsKey(nodeId)) { + throw Exception("Node with ID $nodeId does not exist."); + } + + final node = _nodes[nodeId]; + if (newInputs != null) { + node["inputs"] = newInputs; // Inputs will be resolved in build() + } + + if (newTitle != null) { + node["_meta"]["title"] = newTitle; + } + } + + /// Removes a node from the workflow + void removeNode(String nodeId) { + if (!_nodes.containsKey(nodeId)) { + throw Exception("Node with ID $nodeId does not exist."); + } + + // Remove the node + _nodes.remove(nodeId); + + // Remove associated outputs + _outputToNode.removeWhere((_, value) => value["nodeId"] == nodeId); + } + + /// Reorders nodes in the workflow + void reorderNodes(List newOrder) { + if (newOrder.length != _nodes.keys.length || + !newOrder.every((id) => _nodes.containsKey(id))) { + throw Exception("Invalid new order: All node IDs must be included."); + } + + final reorderedNodes = {}; + for (var nodeId in newOrder) { + reorderedNodes[nodeId] = _nodes[nodeId]; + } + _nodes + ..clear() + ..addAll(reorderedNodes); } /// Generates the final workflow map Map build() { + final resolvedNodes = {}; + + _nodes.forEach((nodeId, node) { + final resolvedInputs = {}; + + // Resolve dependencies for inputs + node["inputs"].forEach((key, value) { + if (value is List && value.length == 2 && value[1] is int) { + final outputTag = value[0]; + if (_outputToNode.containsKey(outputTag)) { + resolvedInputs[key] = [ + _outputToNode[outputTag]!["nodeId"], + value[1] + ]; + } else { + throw Exception( + "Unresolved dependency: No node provides output tag '$outputTag'"); + } + } else { + resolvedInputs[key] = value; + } + }); + + resolvedNodes[nodeId] = { + ...node, + "inputs": resolvedInputs, + }; + }); + return { "client_id": clientId, - "prompt": _nodes, + "prompt": resolvedNodes, }; } @@ -72,18 +138,257 @@ class PromptBuilder { /// Helper method to get default outputs for a class type List _getDefaultOutputs(String classType) { - if (classType == "CheckpointLoaderSimple") { - return ["MODEL", "CLIP", "VAE"]; - } else if (classType == "CLIPTextEncode") { - return ["CONDITIONING"]; - } else if (classType == "EmptyLatentImage" || - classType == "KSampler" || - classType == "LatentUpscale" || - classType == "LatentUpscaleBy") { - return ["LATENT"]; - } else if (classType == "VAEDecode") { - return ["IMAGE"]; + switch (classType) { + case "KSampler": + return ["LATENT"]; + case "CheckpointLoaderSimple": + return ["MODEL", "CLIP", "VAE"]; + case "CLIPTextEncode": + return ["CONDITIONING"]; + case "CLIPSetLastLayer": + return ["CLIP"]; + case "VAEDecode": + return ["IMAGE"]; + case "VAEEncode": + return ["LATENT"]; + case "VAEEncodeForInpaint": + return ["LATENT"]; + case "VAELoader": + return ["VAE"]; + case "EmptyLatentImage": + return ["LATENT"]; + case "LatentUpscale": + return ["LATENT"]; + case "LatentUpscaleBy": + return ["LATENT"]; + case "LatentFromBatch": + return ["LATENT"]; + case "RepeatLatentBatch": + return ["LATENT"]; + case "SaveImage": + return []; + case "PreviewImage": + return []; + case "LoadImage": + return ["IMAGE", "MASK"]; + case "LoadImageMask": + return ["MASK"]; + case "LoadImageOutput": + return ["IMAGE", "MASK"]; + case "ImageScale": + return ["IMAGE"]; + case "ImageScaleBy": + return ["IMAGE"]; + case "ImageInvert": + return ["IMAGE"]; + case "ImageBatch": + return ["IMAGE"]; + case "ImagePadForOutpaint": + return ["IMAGE", "MASK"]; + case "EmptyImage": + return ["IMAGE"]; + case "ConditioningAverage": + return ["CONDITIONING"]; + case "ConditioningCombine": + return ["CONDITIONING"]; + case "ConditioningConcat": + return ["CONDITIONING"]; + case "ConditioningSetArea": + return ["CONDITIONING"]; + case "ConditioningSetAreaPercentage": + return ["CONDITIONING"]; + case "ConditioningSetAreaStrength": + return ["CONDITIONING"]; + case "ConditioningSetMask": + return ["CONDITIONING"]; + case "KSamplerAdvanced": + return ["LATENT"]; + case "SetLatentNoiseMask": + return ["LATENT"]; + case "LatentComposite": + return ["LATENT"]; + case "LatentBlend": + return ["LATENT"]; + case "LatentRotate": + return ["LATENT"]; + case "LatentFlip": + return ["LATENT"]; + case "LatentCrop": + return ["LATENT"]; + case "LoraLoader": + return ["MODEL", "CLIP"]; + case "CLIPLoader": + return ["CLIP"]; + case "UNETLoader": + return ["MODEL"]; + case "DualCLIPLoader": + return ["CLIP"]; + case "CLIPVisionEncode": + return ["CLIP_VISION_OUTPUT"]; + case "StyleModelApply": + return ["CONDITIONING"]; + case "StyleModelLoader": + return ["STYLE_MODEL"]; + case "CLIPVisionLoader": + return ["CLIP_VISION"]; + case "VAEDecodeTiled": + return ["IMAGE"]; + case "VAEEncodeTiled": + return ["LATENT"]; + case "unCLIPCheckpointLoader": + return ["MODEL", "CLIP", "VAE", "CLIP_VISION"]; + case "GLIGENLoader": + return ["GLIGEN"]; + case "GLIGENTextBoxApply": + return ["CONDITIONING"]; + case "InpaintModelConditioning": + return ["positive", "negative", "latent"]; + case "LatentAdd": + return ["LATENT"]; + case "LatentSubtract": + return ["LATENT"]; + case "LatentMultiply": + return ["LATENT"]; + case "LatentInterpolate": + return ["LATENT"]; + case "LatentBatch": + return ["LATENT"]; + case "LatentBatchSeedBehavior": + return ["LATENT"]; + case "LatentApplyOperation": + return ["LATENT"]; + case "LatentApplyOperationCFG": + return ["MODEL"]; + case "LatentOperationTonemapReinhard": + return ["LATENT_OPERATION"]; + case "LatentOperationSharpen": + return ["LATENT_OPERATION"]; + case "HypernetworkLoader": + return ["MODEL"]; + case "UpscaleModelLoader": + return ["UPSCALE_MODEL"]; + case "ImageUpscaleWithModel": + return ["IMAGE"]; + case "ImageBlend": + return ["IMAGE"]; + case "ImageBlur": + return ["IMAGE"]; + case "ImageQuantize": + return ["IMAGE"]; + case "ImageSharpen": + return ["IMAGE"]; + case "ImageScaleToTotalPixels": + return ["IMAGE"]; + case "LatentCompositeMasked": + return ["LATENT"]; + case "ImageCompositeMasked": + return ["IMAGE"]; + case "MaskToImage": + return ["IMAGE"]; + case "ImageToMask": + return ["MASK"]; + case "ImageColorToMask": + return ["MASK"]; + case "SolidMask": + return ["MASK"]; + case "InvertMask": + return ["MASK"]; + case "CropMask": + return ["MASK"]; + case "MaskComposite": + return ["MASK"]; + case "FeatherMask": + return ["MASK"]; + case "GrowMask": + return ["MASK"]; + case "ThresholdMask": + return ["MASK"]; + case "PorterDuffImageComposite": + return ["IMAGE", "MASK"]; + case "SplitImageWithAlpha": + return ["IMAGE", "MASK"]; + case "JoinImageWithAlpha": + return ["IMAGE"]; + case "RebatchLatents": + return ["LATENT"]; + case "RebatchImages": + return ["IMAGE"]; + case "ModelMergeSimple": + return ["MODEL"]; + case "ModelMergeBlocks": + return ["MODEL"]; + case "ModelMergeSubtract": + return ["MODEL"]; + case "ModelMergeAdd": + return ["MODEL"]; + case "CheckpointSave": + return []; + case "CLIPSave": + return []; + case "VAESave": + return []; + case "ModelSave": + return []; + case "TomePatchModel": + return ["MODEL"]; + case "CLIPTextEncodeSDXLRefiner": + return ["CONDITIONING"]; + case "CLIPTextEncodeSDXL": + return ["CONDITIONING"]; + case "Canny": + return ["IMAGE"]; + case "FreeU": + return ["MODEL"]; + case "FreeU_V2": + return ["MODEL"]; + case "SamplerCustom": + return ["LATENT", "LATENT"]; + case "BasicScheduler": + return ["SIGMAS"]; + case "KarrasScheduler": + return ["SIGMAS"]; + case "ExponentialScheduler": + return ["SIGMAS"]; + case "PolyexponentialScheduler": + return ["SIGMAS"]; + case "LaplaceScheduler": + return ["SIGMAS"]; + case "VPScheduler": + return ["SIGMAS"]; + case "BetaSamplingScheduler": + return ["SIGMAS"]; + case "SDTurboScheduler": + return ["SIGMAS"]; + case "KSamplerSelect": + return ["SAMPLER"]; + case "SamplerEulerAncestral": + return ["SAMPLER"]; + case "SamplerEulerAncestralCFGPP": + return ["SAMPLER"]; + case "SamplerLMS": + return ["SAMPLER"]; + case "SamplerDPMPP_3M_SDE": + return ["SAMPLER"]; + case "SamplerDPMPP_2M_SDE": + return ["SAMPLER"]; + case "SamplerDPMPP_SDE": + return ["SAMPLER"]; + case "SamplerDPMPP_2S_Ancestral": + return ["SAMPLER"]; + case "SamplerDPMAdaptative": + return ["SAMPLER"]; + case "SplitSigmas": + return ["SIGMAS", "SIGMAS"]; + case "SplitSigmasDenoise": + return ["SIGMAS", "SIGMAS"]; + case "FlipSigmas": + return ["SIGMAS"]; + case "SetFirstSigma": + return ["SIGMAS"]; + case "CFGGuider": + return ["GUIDER"]; + default: + return []; } - return []; } }