From f71877dfcd3391685bb3f341bee7cb96a6c2e621 Mon Sep 17 00:00:00 2001 From: logtd Date: Thu, 12 Dec 2024 23:34:52 -0700 Subject: [PATCH] add flow edit --- __init__.py | 10 +- example_workflows/example_flow_edit.json | 1559 ++++++++++++++++++++++ nodes/flowedit_nodes.py | 144 ++ 3 files changed, 1712 insertions(+), 1 deletion(-) create mode 100644 example_workflows/example_flow_edit.json create mode 100644 nodes/flowedit_nodes.py diff --git a/__init__.py b/__init__.py index 026f1c0..86b4e44 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,8 @@ from .nodes.configure_modified_flux_node import ConfigureModifiedFluxNode from .nodes.mix_noise_node import FluxNoiseMixerNode from .nodes.rectified_sampler_nodes import FluxForwardODESamplerNode, FluxReverseODESamplerNode from .nodes.rf_edit_sampler_nodes import FlowEditForwardSamplerNode, FlowEditReverseSamplerNode, PrepareAttnBankNode +## FlowEdit +from .nodes.flowedit_nodes import FlowEditGuiderNode, FlowEditSamplerNode ## Enhance from .nodes.apply_seg_node import SEGAttentionNode from .nodes.apply_pag_node import PAGAttentionNode @@ -40,6 +42,9 @@ NODE_CLASS_MAPPINGS = { "PrepareAttnBank": PrepareAttnBankNode, "RFSingleBlocksOverride": RFSingleBlocksOverrideNode, "RFDoubleBlocksOverride": RFDoubleBlocksOverrideNode, + ## Flow-Edit + "FlowEditGuider": FlowEditGuiderNode, + "FlowEditSampler": FlowEditSamplerNode, ## ATTN "FluxAttnOverride": FluxAttnOverrideNode, "SEGAttention": SEGAttentionNode, @@ -73,7 +78,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PrepareAttnBank": "RF-Edit Prep Attn Inj", "RFSingleBlocksOverride": "RF-Edit Single Layers Override", "RFDoubleBlocksOverride": "RF-Edit Double Layers Override", - ## Enhance + ## FlowEdit + "FlowEditGuider": "Flow Edit Guider", + "FlowEditSampler": "Flow Edit Sampler", + ## ATTN "FluxAttnOverride": "Flux Attention Override", "SEGAttention": "Apply Flux SEG Attention", "PAGAttention": "Apply Flux PAG Attention", diff --git a/example_workflows/example_flow_edit.json b/example_workflows/example_flow_edit.json new file mode 100644 index 0000000..4347860 --- /dev/null +++ b/example_workflows/example_flow_edit.json @@ -0,0 +1,1559 @@ +{ + "last_node_id": 362, + "last_link_id": 632, + "nodes": [ + { + "id": 238, + "type": "SamplerCustomAdvanced", + "pos": [ + -470, + 1800 + ], + "size": [ + 236.8000030517578, + 106 + ], + "flags": {}, + "order": 30, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 623, + "slot_index": 0 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 622, + "slot_index": 1 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 607, + "slot_index": 2 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 432, + "slot_index": 3 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 613, + "slot_index": 4 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 508 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "SamplerCustomAdvanced" + }, + "widgets_values": [] + }, + { + "id": 263, + "type": "SetNode", + "pos": [ + -1587.62939453125, + 2166.945556640625 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "INT", + "type": "INT", + "link": 463 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_WIDTH", + "properties": { + "previousName": "" + }, + "widgets_values": [ + "WIDTH" + ] + }, + { + "id": 347, + "type": "OutFluxModelSamplingPred", + "pos": [ + -462.64556884765625, + 2271.548583984375 + ], + "size": [ + 210, + 154 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 604 + }, + { + "name": "width", + "type": "INT", + "link": 602, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 603, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 605, + 606 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "OutFluxModelSamplingPred" + }, + "widgets_values": [ + 1.15, + 0.5, + 1024, + 1024, + true + ] + }, + { + "id": 261, + "type": "GetNode", + "pos": [ + -466.4717102050781, + 2471.769287109375 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 604 + ], + "slot_index": 0 + } + ], + "title": "Get_MODEL_FLUX", + "properties": {}, + "widgets_values": [ + "MODEL_FLUX" + ] + }, + { + "id": 340, + "type": "GetNode", + "pos": [ + -468.6756591796875, + 2523.3779296875 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 603 + ], + "slot_index": 0 + } + ], + "title": "Get_HEIGHT", + "properties": {}, + "widgets_values": [ + "HEIGHT" + ] + }, + { + "id": 339, + "type": "GetNode", + "pos": [ + -466.44476318359375, + 2572.911376953125 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 602 + ], + "slot_index": 0 + } + ], + "title": "Get_WIDTH", + "properties": {}, + "widgets_values": [ + "WIDTH" + ] + }, + { + "id": 276, + "type": "UNETLoader", + "pos": [ + -1977.3173828125, + 1654.719970703125 + ], + "size": [ + 336.3075866699219, + 82 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 473 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "UNETLoader" + }, + "widgets_values": [ + "flux1-dev-fp8-e4m3fn.safetensors", + "fp8_e4m3fn" + ] + }, + { + "id": 265, + "type": "SetNode", + "pos": [ + -1597.62939453125, + 2276.9453125 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "INT", + "type": "INT", + "link": 464 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_HEIGHT", + "properties": { + "previousName": "" + }, + "widgets_values": [ + "HEIGHT" + ] + }, + { + "id": 277, + "type": "DualCLIPLoader", + "pos": [ + -1965.8372802734375, + 1805.784423828125 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 462 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader" + }, + "widgets_values": [ + "t5xxl_fp16.safetensors", + "clip_l.safetensors", + "flux" + ] + }, + { + "id": 278, + "type": "VAELoader", + "pos": [ + -1962.0889892578125, + 1963.8216552734375 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 460 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "ae.safetensors" + ] + }, + { + "id": 254, + "type": "SetNode", + "pos": [ + -1617.237548828125, + 1966.9456787109375 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "VAE", + "type": "VAE", + "link": 460 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_VAE_FLUX", + "properties": { + "previousName": "" + }, + "widgets_values": [ + "VAE_FLUX" + ] + }, + { + "id": 256, + "type": "SetNode", + "pos": [ + -1625.9625244140625, + 1827.230224609375 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "CLIP", + "type": "CLIP", + "link": 462 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_CLIP_FLUX", + "properties": { + "previousName": "CLIP_FLUX" + }, + "widgets_values": [ + "CLIP_FLUX" + ] + }, + { + "id": 284, + "type": "ConfigureModifiedFlux", + "pos": [ + -1598.025634765625, + 1608.0743408203125 + ], + "size": [ + 193.1999969482422, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 473 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 592 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ConfigureModifiedFlux" + }, + "widgets_values": [] + }, + { + "id": 338, + "type": "SetNode", + "pos": [ + -1601.309814453125, + 1688.076904296875 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "MODEL", + "type": "MODEL", + "link": 592 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_MODEL_FLUX", + "properties": { + "previousName": "MODEL_FLUX" + }, + "widgets_values": [ + "MODEL_FLUX" + ] + }, + { + "id": 259, + "type": "GetNode", + "pos": [ + -1321.108154296875, + 1609.5753173828125 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": false + }, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 426, + 600 + ], + "slot_index": 0 + } + ], + "title": "Get_CLIP_FLUX", + "properties": {}, + "widgets_values": [ + "CLIP_FLUX" + ] + }, + { + "id": 344, + "type": "FlowEditGuider", + "pos": [ + -463.99652099609375, + 1696.387451171875 + ], + "size": [ + 229.20001220703125, + 66 + ], + "flags": {}, + "order": 29, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 606 + }, + { + "name": "source_cond", + "type": "CONDITIONING", + "link": 621 + }, + { + "name": "target_cond", + "type": "CONDITIONING", + "link": 619 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 622 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FlowEditGuider" + } + }, + { + "id": 356, + "type": "DisableNoise", + "pos": [ + -463.1148986816406, + 1651.5369873046875 + ], + "size": [ + 210, + 26 + ], + "flags": { + "collapsed": true + }, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 623 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "DisableNoise" + } + }, + { + "id": 351, + "type": "VAEEncode", + "pos": [ + -1175.1278076171875, + 2231.947509765625 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 27, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 612 + }, + { + "name": "vae", + "type": "VAE", + "link": 611 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 613 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 352, + "type": "GetNode", + "pos": [ + -1179.8497314453125, + 2325.322265625 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 611 + ], + "slot_index": 0 + } + ], + "title": "Get_VAE_FLUX", + "properties": {}, + "widgets_values": [ + "VAE_FLUX" + ] + }, + { + "id": 260, + "type": "GetNode", + "pos": [ + -202.75523376464844, + 1636.953369140625 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 9, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 509 + ], + "slot_index": 0 + } + ], + "title": "Get_VAE_FLUX", + "properties": {}, + "widgets_values": [ + "VAE_FLUX" + ] + }, + { + "id": 266, + "type": "INTConstant", + "pos": [ + -1887.62939453125, + 2266.9453125 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "value", + "type": "INT", + "links": [ + 464, + 609 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "INTConstant" + }, + "widgets_values": [ + 768 + ], + "color": "#1b4669", + "bgcolor": "#29699c" + }, + { + "id": 264, + "type": "INTConstant", + "pos": [ + -1887.62939453125, + 2156.945556640625 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "value", + "type": "INT", + "links": [ + 463, + 608 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "INTConstant" + }, + "widgets_values": [ + 1344 + ], + "color": "#1b4669", + "bgcolor": "#29699c" + }, + { + "id": 355, + "type": "FluxGuidance", + "pos": [ + -923.563232421875, + 1747.8154296875 + ], + "size": [ + 211.60000610351562, + 58 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 620 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 621 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 1.5 + ] + }, + { + "id": 354, + "type": "FluxGuidance", + "pos": [ + -922.0948486328125, + 1950.1080322265625 + ], + "size": [ + 211.60000610351562, + 58 + ], + "flags": {}, + "order": 26, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 618 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 619 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 5.5 + ] + }, + { + "id": 350, + "type": "ImageScale", + "pos": [ + -1584.5281982421875, + 2392.308349609375 + ], + "size": [ + 210, + 130 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 610 + }, + { + "name": "width", + "type": "INT", + "link": 608, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 609, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 612, + 624 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "nearest-exact", + 512, + 512, + "center" + ] + }, + { + "id": 359, + "type": "SetNode", + "pos": [ + -1597.1357421875, + 2586.158447265625 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 28, + "mode": 0, + "inputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "link": 624 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_IMG", + "properties": { + "previousName": "IMG" + }, + "widgets_values": [ + "IMG" + ] + }, + { + "id": 279, + "type": "PreviewImage", + "pos": [ + -179.13108825683594, + 1799.76904296875 + ], + "size": [ + 749.0811157226562, + 804.638916015625 + ], + "flags": {}, + "order": 32, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 510 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 301, + "type": "VAEDecode", + "pos": [ + -201.4112548828125, + 1683.3875732421875 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 31, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 508 + }, + { + "name": "vae", + "type": "VAE", + "link": 509 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 510 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 349, + "type": "LoadImage", + "pos": [ + -1936.162353515625, + 2395.129638671875 + ], + "size": [ + 315, + 314 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 610 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "love.jpeg", + "image" + ] + }, + { + "id": 236, + "type": "CLIPTextEncode", + "pos": [ + -1327.060791015625, + 1723.951904296875 + ], + "size": [ + 348.5190124511719, + 193.99783325195312 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 426 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 620 + ], + "slot_index": 0 + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "a billboard in a city with the text \"LOVE IS ALL YOU NEED\"" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 343, + "type": "CLIPTextEncode", + "pos": [ + -1325.83984375, + 1971.5255126953125 + ], + "size": [ + 348.5190124511719, + 193.99783325195312 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 600 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 618 + ], + "slot_index": 0 + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "a billboard in a city with the text \"FLUX IS ALL YOU NEED\"" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 239, + "type": "BasicScheduler", + "pos": [ + -459.8354797363281, + 1947.8602294921875 + ], + "size": [ + 210, + 106 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 605, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 432 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 28, + 1 + ] + }, + { + "id": 348, + "type": "FlowEditSampler", + "pos": [ + -460.8995361328125, + 2097.6875 + ], + "size": [ + 210, + 130 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 607 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FlowEditSampler" + }, + "widgets_values": [ + 4, + 0, + 1, + "randomize" + ] + } + ], + "links": [ + [ + 426, + 259, + 0, + 236, + 0, + "CLIP" + ], + [ + 432, + 239, + 0, + 238, + 3, + "SIGMAS" + ], + [ + 460, + 278, + 0, + 254, + 0, + "*" + ], + [ + 462, + 277, + 0, + 256, + 0, + "*" + ], + [ + 463, + 264, + 0, + 263, + 0, + "*" + ], + [ + 464, + 266, + 0, + 265, + 0, + "*" + ], + [ + 473, + 276, + 0, + 284, + 0, + "MODEL" + ], + [ + 508, + 238, + 0, + 301, + 0, + "LATENT" + ], + [ + 509, + 260, + 0, + 301, + 1, + "VAE" + ], + [ + 510, + 301, + 0, + 279, + 0, + "IMAGE" + ], + [ + 592, + 284, + 0, + 338, + 0, + "*" + ], + [ + 600, + 259, + 0, + 343, + 0, + "CLIP" + ], + [ + 602, + 339, + 0, + 347, + 1, + "INT" + ], + [ + 603, + 340, + 0, + 347, + 2, + "INT" + ], + [ + 604, + 261, + 0, + 347, + 0, + "MODEL" + ], + [ + 605, + 347, + 0, + 239, + 0, + "MODEL" + ], + [ + 606, + 347, + 0, + 344, + 0, + "MODEL" + ], + [ + 607, + 348, + 0, + 238, + 2, + "SAMPLER" + ], + [ + 608, + 264, + 0, + 350, + 1, + "INT" + ], + [ + 609, + 266, + 0, + 350, + 2, + "INT" + ], + [ + 610, + 349, + 0, + 350, + 0, + "IMAGE" + ], + [ + 611, + 352, + 0, + 351, + 1, + "VAE" + ], + [ + 612, + 350, + 0, + 351, + 0, + "IMAGE" + ], + [ + 613, + 351, + 0, + 238, + 4, + "LATENT" + ], + [ + 618, + 343, + 0, + 354, + 0, + "CONDITIONING" + ], + [ + 619, + 354, + 0, + 344, + 2, + "CONDITIONING" + ], + [ + 620, + 236, + 0, + 355, + 0, + "CONDITIONING" + ], + [ + 621, + 355, + 0, + 344, + 1, + "CONDITIONING" + ], + [ + 622, + 344, + 0, + 238, + 1, + "GUIDER" + ], + [ + 623, + 356, + 0, + 238, + 0, + "NOISE" + ], + [ + 624, + 350, + 0, + 359, + 0, + "*" + ] + ], + "groups": [ + { + "id": 1, + "title": "Setup", + "bounding": [ + -1987.3175048828125, + 1534.474365234375, + 622.7884521484375, + 1184.6556396484375 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 3, + "title": "Sampling", + "bounding": [ + -1337.060791015625, + 1535.975341796875, + 1905.8232421875, + 1075.565185546875 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 0.3452271214393729, + "offset": [ + 2956.9928229470947, + -539.5084793631383 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes/flowedit_nodes.py b/nodes/flowedit_nodes.py new file mode 100644 index 0000000..f5caf13 --- /dev/null +++ b/nodes/flowedit_nodes.py @@ -0,0 +1,144 @@ +import torch +from tqdm import trange + +from comfy.samplers import KSAMPLER, CFGGuider, sampling_function + + +class FlowEditGuider(CFGGuider): + def __init__(self, model_patcher): + super().__init__(model_patcher) + self.cfgs = {} + + def set_conds(self, **kwargs): + self.inner_set_conds(kwargs) + + def set_cfgs(self, **kwargs): + self.cfgs = {**kwargs} + + def predict_noise(self, x, timestep, model_options={}, seed=None): + latent_type = model_options['transformer_options']['latent_type'] + positive = self.conds.get(f'{latent_type}_positive', None) + negative = self.conds.get(f'{latent_type}_negative', None) + cfg = self.cfgs.get(latent_type, self.cfg) + return sampling_function(self.inner_model, x, timestep, negative, positive, cfg, model_options=model_options, seed=seed) + + +class FlowEditGuiderNode: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "model": ("MODEL",), + "source_cond": ("CONDITIONING", ), + "target_cond": ("CONDITIONING", ), + } + } + + RETURN_TYPES = ("GUIDER",) + + FUNCTION = "get_guider" + CATEGORY = "fluxtapoz" + + def get_guider(self, model, source_cond, target_cond): + guider = FlowEditGuider(model) + guider.set_conds(source_positive=source_cond, target_positive=target_cond) + guider.set_cfg(1.0) + return (guider,) + + + + +class FlowEditCFGGuiderNode: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "model": ("MODEL",), + "source_pos": ("CONDITIONING", ), + "source_neg": ("CONDITIONING", ), + "target_pos": ("CONDITIONING", ), + "target_neg": ("CONDITIONING", ), + "source_cfg": ("FLOAT", {"default": 3.5, "min": 0, "max": 0xffffffffffffffff, "step": 0.01 }), + "target_cfg": ("FLOAT", {"default": 13.5, "min": 0, "max": 0xffffffffffffffff, "step": 0.01 }), + } + } + + RETURN_TYPES = ("GUIDER",) + + FUNCTION = "get_guider" + CATEGORY = "fluxtapoz" + + def get_guider(self, model, source_pos, source_neg, target_pos, target_neg, source_cfg, target_cfg): + guider = FlowEditGuider(model) + guider.set_conds(source_positive=source_pos, source_negative=source_neg, target_positive=target_pos, target_negative=target_neg) + guider.set_cfgs(source=source_cfg, target=target_cfg) + return (guider,) + + +def get_flowedit_sample(skip_steps, refine_steps, seed): + generator = torch.manual_seed(seed) + @torch.no_grad() + def flowedit_sample(model, x_init, sigmas, extra_args=None, callback=None, disable=None): + extra_args = {} if extra_args is None else extra_args + + model_options = extra_args.get('model_options', {}) + transformer_options = model_options.get('transformer_options', {}) + transformer_options = {**transformer_options} + model_options['transformer_options'] = transformer_options + extra_args['model_options'] = model_options + + source_extra_args = {**extra_args, 'model_options': { 'transformer_options': { 'latent_type': 'source '} }} + + sigmas = sigmas[skip_steps:] + + x_tgt = x_init.clone() + N = len(sigmas)-1 + s_in = x_init.new_ones([x_init.shape[0]]) + + for i in trange(N, disable=disable): + sigma = sigmas[i] + noise = torch.randn(x_init.shape, generator=generator).to(x_init.device) + + zt_src = (1-sigma)*x_init + sigma*noise + + if i < N-refine_steps: + zt_tgt = x_tgt + zt_src - x_init + transformer_options['latent_type'] = 'source' + vt_src = model(zt_src, sigma*s_in, **source_extra_args) + else: + if i == N-refine_steps: + x_tgt = x_tgt + zt_src - x_init + zt_tgt = x_tgt + vt_src = 0 + + transformer_options['latent_type'] = 'target' + vt_tgt = model(zt_tgt, sigma*s_in, **extra_args) + + v_delta = vt_tgt - vt_src + x_tgt += (sigmas[i+1] - sigmas[i]) * v_delta + + if callback is not None: + callback({'x': x_tgt, 'denoised': x_tgt, 'i': i+skip_steps, 'sigma': sigmas[i], 'sigma_hat': sigmas[i]}) + + return x_tgt + + return flowedit_sample + + +class FlowEditSamplerNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "skip_steps": ("INT", {"default": 4, "min": 0, "max": 0xffffffffffffffff }), + "refine_steps": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff }), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff }), + }, "optional": { + }} + RETURN_TYPES = ("SAMPLER",) + FUNCTION = "build" + + CATEGORY = "fluxtapoz" + + def build(self, skip_steps, refine_steps, seed): + sampler = KSAMPLER(get_flowedit_sample(skip_steps, refine_steps, seed)) + return (sampler, )