From 80cd31e061ace98307008a9da42ec0597cf03a93 Mon Sep 17 00:00:00 2001 From: logtd Date: Sun, 15 Sep 2024 17:15:08 -0600 Subject: [PATCH] updateref --- .gitignore | 2 + __init__.py | 11 + example_workflows/example_flux_reference.json | 1412 ++++++++++++++++ .../reflux_workflow_example_updated.json | 1445 +++++++++++++++++ flux/layers.py | 165 ++ flux/model.py | 147 ++ .../flip_sigmas_node.cpython-310.pyc | Bin 769 -> 769 bytes .../flux_deguidance_node.cpython-310.pyc | Bin 934 -> 934 bytes .../influx_model_pred_node.cpython-310.pyc | Bin 3695 -> 3683 bytes nodes/add_flow_flux_node.py | 29 + nodes/apply_rave_node.py | 41 + nodes/apply_ref_flux.py | 65 + utils/module_utils.py | 7 + utils/noise_utils.py | 26 + utils/rave_attention.py | 117 ++ utils/rave_rope_attention.py | 130 ++ utils/rave_utils.py | 57 + 17 files changed, 3654 insertions(+) create mode 100644 .gitignore create mode 100644 example_workflows/example_flux_reference.json create mode 100644 example_workflows/reflux_workflow_example_updated.json create mode 100644 flux/layers.py create mode 100644 flux/model.py create mode 100644 nodes/add_flow_flux_node.py create mode 100644 nodes/apply_rave_node.py create mode 100644 nodes/apply_ref_flux.py create mode 100644 utils/module_utils.py create mode 100644 utils/noise_utils.py create mode 100644 utils/rave_attention.py create mode 100644 utils/rave_rope_attention.py create mode 100644 utils/rave_utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bd1c29a --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +**/__pycache__ +.DS_Store diff --git a/__init__.py b/__init__.py index c72c750..f5f9e49 100644 --- a/__init__.py +++ b/__init__.py @@ -2,6 +2,9 @@ from .nodes.flip_sigmas_node import InFluxFlipSigmasNode from .nodes.influx_model_pred_node import InFluxModelSamplingPredNode, OutFluxModelSamplingPredNode from .nodes.flux_deguidance_node import FluxDeGuidance from .nodes.inverse_sampler_node import FluxInverseSamplerNode +from .nodes.apply_ref_flux import ApplyRefFluxNode, ConfigureRefFluxNode +from .nodes.add_flow_flux_node import AddFluxFlowNode +from .nodes.apply_rave_node import ApplyFluxRaveAttentionNode NODE_CLASS_MAPPINGS = { @@ -10,6 +13,10 @@ NODE_CLASS_MAPPINGS = { "OutFluxModelSamplingPred": OutFluxModelSamplingPredNode, "FluxDeGuidance": FluxDeGuidance, "FluxInverseSampler": FluxInverseSamplerNode, + "ApplyRefFlux": ApplyRefFluxNode, + "ConfigureRefFlux": ConfigureRefFluxNode, + # "AddFluxFlow": AddFluxFlowNode, + # "ApplyFluxRaveAttention": ApplyFluxRaveAttentionNode, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -18,4 +25,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "OutFluxModelSamplingPred": "Outverse Flux Model Pred", "FluxDeGuidance": "Flux DeGuidance", "FluxInverseSampler": "Flux Inverse Sampler", + "ApplyRefFlux": "Apply Ref Flux Model", + "ConfigureRefFlux": "Configure Ref for Flux", + # "AddFluxFlow": "Add Flux Flow", + # "ApplyFluxRaveAttention": "Apply Flux Rave Attn", } diff --git a/example_workflows/example_flux_reference.json b/example_workflows/example_flux_reference.json new file mode 100644 index 0000000..7535c3b --- /dev/null +++ b/example_workflows/example_flux_reference.json @@ -0,0 +1,1412 @@ +{ + "last_node_id": 50, + "last_link_id": 154, + "nodes": [ + { + "id": 11, + "type": "DualCLIPLoader", + "pos": { + "0": -860, + "1": 350 + }, + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 10 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader" + }, + "widgets_values": [ + "t5xxl_fp16.safetensors", + "clip_l.safetensors", + "flux" + ] + }, + { + "id": 10, + "type": "VAELoader", + "pos": { + "0": -860, + "1": 490 + }, + "size": { + "0": 311.81634521484375, + "1": 60.429901123046875 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 12, + 125 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "ae.safetensors" + ] + }, + { + "id": 40, + "type": "LoadImage", + "pos": { + "0": -836, + "1": 968 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 128 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "dog.jpg", + "image" + ] + }, + { + "id": 12, + "type": "UNETLoader", + "pos": { + "0": -738, + "1": 1366 + }, + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 117 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "UNETLoader" + }, + "widgets_values": [ + "flux1-dev-fp8.safetensors", + "fp8_e4m3fn" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 22, + "type": "BasicGuider", + "pos": { + "0": 500, + "1": 213 + }, + "size": { + "0": 222.3482666015625, + "1": 46 + }, + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 144, + "slot_index": 0 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 42, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 30 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicGuider" + } + }, + { + "id": 42, + "type": "VAEEncode", + "pos": { + "0": -437, + "1": 1167 + }, + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 129 + }, + { + "name": "vae", + "type": "VAE", + "link": 125 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 138 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 38, + "type": "ApplyRefFlux", + "pos": { + "0": -399, + "1": 1365 + }, + "size": [ + 168, + 26 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 117 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 134, + 150 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ApplyRefFlux" + } + }, + { + "id": 34, + "type": "PrimitiveNode", + "pos": { + "0": -741, + "1": 634 + }, + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 112, + 115, + 145 + ], + "slot_index": 0, + "widget": { + "name": "width" + } + } + ], + "title": "width", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 1024, + "fixed" + ], + "color": "#323", + "bgcolor": "#535" + }, + { + "id": 35, + "type": "PrimitiveNode", + "pos": { + "0": -740, + "1": 788 + }, + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 113, + 114, + 146 + ], + "slot_index": 0, + "widget": { + "name": "height" + } + } + ], + "title": "height", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 1024, + "fixed" + ], + "color": "#323", + "bgcolor": "#535" + }, + { + "id": 43, + "type": "ImageScale", + "pos": { + "0": -435, + "1": 934 + }, + "size": [ + 210, + 130 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 128 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 129, + 151 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "nearest-exact", + 1024, + 1024, + "center" + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": { + "0": 866, + "1": 367 + }, + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 24 + }, + { + "name": "vae", + "type": "VAE", + "link": 12 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9, + 152 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 46, + "type": "ModelSamplingFlux", + "pos": { + "0": 189, + "1": 927 + }, + "size": { + "0": 210, + "1": 122 + }, + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 143, + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "link": 145, + "slot_index": 1, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 146, + "slot_index": 2, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 144 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ModelSamplingFlux" + }, + "widgets_values": [ + 1.15, + 0.5, + 1024, + 1024 + ] + }, + { + "id": 17, + "type": "BasicScheduler", + "pos": { + "0": 187, + "1": 769 + }, + "size": [ + 210, + 106 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 141, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 20, + 139 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 20, + 1 + ] + }, + { + "id": 27, + "type": "EmptySD3LatentImage", + "pos": { + "0": 195, + "1": 632 + }, + "size": [ + 210, + 94.35025005634816 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "link": 112, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 113, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 116 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "EmptySD3LatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 16, + "type": "KSamplerSelect", + "pos": { + "0": 198, + "1": 510 + }, + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 19 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "euler" + ] + }, + { + "id": 26, + "type": "FluxGuidance", + "pos": { + "0": 182, + "1": 210 + }, + "size": [ + 211.60000610351562, + 58 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 41 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 42 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 3.5 + ], + "color": "#233", + "bgcolor": "#355" + }, + { + "id": 9, + "type": "SaveImage", + "pos": { + "0": 1282, + "1": 196 + }, + "size": { + "0": 985.3012084960938, + "1": 1060.3828125 + }, + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 30, + "type": "ModelSamplingFlux", + "pos": { + "0": -17, + "1": 796 + }, + "size": [ + 210, + 122 + ], + "flags": { + "collapsed": true + }, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 134, + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "link": 115, + "slot_index": 1, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 114, + "slot_index": 2, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 141 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ModelSamplingFlux" + }, + "widgets_values": [ + 1.15, + 0.5, + 1024, + 1024 + ] + }, + { + "id": 13, + "type": "SamplerCustomAdvanced", + "pos": { + "0": 864, + "1": 192 + }, + "size": [ + 236.8000030517578, + 106 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 37, + "slot_index": 0 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 30, + "slot_index": 1 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 19, + "slot_index": 2 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 20, + "slot_index": 3 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 116, + "slot_index": 4 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 24 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "SamplerCustomAdvanced" + } + }, + { + "id": 25, + "type": "RandomNoise", + "pos": { + "0": 889, + "1": 47 + }, + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 37 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 0, + "fixed" + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 47, + "type": "ImageConcanate", + "pos": { + "0": 1166, + "1": 1415 + }, + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "image1", + "type": "IMAGE", + "link": 153 + }, + { + "name": "image2", + "type": "IMAGE", + "link": 154 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 149 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageConcanate" + }, + "widgets_values": [ + "right", + false + ] + }, + { + "id": 49, + "type": "AddLabel", + "pos": { + "0": 763, + "1": 1438 + }, + "size": [ + 315, + 274 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 151 + }, + { + "name": "caption", + "type": "STRING", + "link": null, + "widget": { + "name": "caption" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 153 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "AddLabel" + }, + "widgets_values": [ + 10, + 2, + 48, + 32, + "white", + "black", + "FreeMono.ttf", + "reference", + "up", + "" + ] + }, + { + "id": 50, + "type": "AddLabel", + "pos": { + "0": 1278.6334228515625, + "1": 1718.9046630859375 + }, + "size": [ + 315, + 274 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 152 + }, + { + "name": "caption", + "type": "STRING", + "link": null, + "widget": { + "name": "caption" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 154 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "AddLabel" + }, + "widgets_values": [ + 10, + 2, + 48, + 32, + "white", + "black", + "FreeMono.ttf", + "output", + "up", + "" + ] + }, + { + "id": 48, + "type": "PreviewImage", + "pos": { + "0": 1624, + "1": 1411 + }, + "size": [ + 487.29357697753585, + 431.7405868100341 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 149 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": { + "0": -249, + "1": 195 + }, + "size": [ + 285.6000061035156, + 98.6700754880381 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 10 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 41 + ], + "slot_index": 0 + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "a dog wearing a top hat" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 45, + "type": "ConfigureRefFlux", + "pos": { + "0": 194, + "1": 1107 + }, + "size": { + "0": 210, + "1": 122 + }, + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 150 + }, + { + "name": "latent", + "type": "LATENT", + "link": 138 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 139 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 143 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ConfigureRefFlux" + }, + "widgets_values": [ + 0.1, + 0.2 + ] + } + ], + "links": [ + [ + 9, + 8, + 0, + 9, + 0, + "IMAGE" + ], + [ + 10, + 11, + 0, + 6, + 0, + "CLIP" + ], + [ + 12, + 10, + 0, + 8, + 1, + "VAE" + ], + [ + 19, + 16, + 0, + 13, + 2, + "SAMPLER" + ], + [ + 20, + 17, + 0, + 13, + 3, + "SIGMAS" + ], + [ + 24, + 13, + 0, + 8, + 0, + "LATENT" + ], + [ + 30, + 22, + 0, + 13, + 1, + "GUIDER" + ], + [ + 37, + 25, + 0, + 13, + 0, + "NOISE" + ], + [ + 41, + 6, + 0, + 26, + 0, + "CONDITIONING" + ], + [ + 42, + 26, + 0, + 22, + 1, + "CONDITIONING" + ], + [ + 112, + 34, + 0, + 27, + 0, + "INT" + ], + [ + 113, + 35, + 0, + 27, + 1, + "INT" + ], + [ + 114, + 35, + 0, + 30, + 2, + "INT" + ], + [ + 115, + 34, + 0, + 30, + 1, + "INT" + ], + [ + 116, + 27, + 0, + 13, + 4, + "LATENT" + ], + [ + 117, + 12, + 0, + 38, + 0, + "MODEL" + ], + [ + 125, + 10, + 0, + 42, + 1, + "VAE" + ], + [ + 128, + 40, + 0, + 43, + 0, + "IMAGE" + ], + [ + 129, + 43, + 0, + 42, + 0, + "IMAGE" + ], + [ + 134, + 38, + 0, + 30, + 0, + "MODEL" + ], + [ + 138, + 42, + 0, + 45, + 1, + "LATENT" + ], + [ + 139, + 17, + 0, + 45, + 2, + "SIGMAS" + ], + [ + 141, + 30, + 0, + 17, + 0, + "MODEL" + ], + [ + 143, + 45, + 0, + 46, + 0, + "MODEL" + ], + [ + 144, + 46, + 0, + 22, + 0, + "MODEL" + ], + [ + 145, + 34, + 0, + 46, + 1, + "INT" + ], + [ + 146, + 35, + 0, + 46, + 2, + "INT" + ], + [ + 149, + 47, + 0, + 48, + 0, + "IMAGE" + ], + [ + 150, + 38, + 0, + 45, + 0, + "MODEL" + ], + [ + 151, + 43, + 0, + 49, + 0, + "IMAGE" + ], + [ + 152, + 8, + 0, + 50, + 0, + "IMAGE" + ], + [ + 153, + 49, + 0, + 47, + 0, + "IMAGE" + ], + [ + 154, + 50, + 0, + 47, + 1, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.4641000000000006, + "offset": [ + -1382.392281013204, + -1398.2619898147707 + ] + }, + "groupNodes": {} + }, + "version": 0.4 +} \ No newline at end of file diff --git a/example_workflows/reflux_workflow_example_updated.json b/example_workflows/reflux_workflow_example_updated.json new file mode 100644 index 0000000..358cece --- /dev/null +++ b/example_workflows/reflux_workflow_example_updated.json @@ -0,0 +1,1445 @@ +{ + "last_node_id": 51, + "last_link_id": 158, + "nodes": [ + { + "id": 12, + "type": "UNETLoader", + "pos": { + "0": -738, + "1": 1366 + }, + "size": { + "0": 210, + "1": 82 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 117 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "UNETLoader" + }, + "widgets_values": [ + "flux1-dev-fp8.safetensors", + "fp8_e4m3fn" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 42, + "type": "VAEEncode", + "pos": { + "0": -437, + "1": 1167 + }, + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 129 + }, + { + "name": "vae", + "type": "VAE", + "link": 125 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 138 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 38, + "type": "ApplyRefFlux", + "pos": { + "0": -399, + "1": 1365 + }, + "size": { + "0": 168, + "1": 26 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 117 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 134, + 150 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ApplyRefFlux" + } + }, + { + "id": 8, + "type": "VAEDecode", + "pos": { + "0": 866, + "1": 367 + }, + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 24 + }, + { + "name": "vae", + "type": "VAE", + "link": 12 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9, + 152 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 9, + "type": "SaveImage", + "pos": { + "0": 1282, + "1": 196 + }, + "size": { + "0": 985.3012084960938, + "1": 1060.3828125 + }, + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 30, + "type": "ModelSamplingFlux", + "pos": { + "0": -17, + "1": 796 + }, + "size": { + "0": 210, + "1": 122 + }, + "flags": { + "collapsed": true + }, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 134, + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "link": 115, + "slot_index": 1, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 114, + "slot_index": 2, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 141 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ModelSamplingFlux" + }, + "widgets_values": [ + 1.15, + 0.5, + 1024, + 1024 + ] + }, + { + "id": 47, + "type": "ImageConcanate", + "pos": { + "0": 1166, + "1": 1415 + }, + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "image1", + "type": "IMAGE", + "link": 153 + }, + { + "name": "image2", + "type": "IMAGE", + "link": 154 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 149 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ImageConcanate" + }, + "widgets_values": [ + "right", + false + ] + }, + { + "id": 49, + "type": "AddLabel", + "pos": { + "0": 763, + "1": 1438 + }, + "size": { + "0": 315, + "1": 274 + }, + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 151 + }, + { + "name": "caption", + "type": "STRING", + "link": null, + "widget": { + "name": "caption" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 153 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "AddLabel" + }, + "widgets_values": [ + 10, + 2, + 48, + 32, + "white", + "black", + "FreeMono.ttf", + "reference", + "up", + "" + ] + }, + { + "id": 48, + "type": "PreviewImage", + "pos": { + "0": 1624, + "1": 1411 + }, + "size": { + "0": 487.2935791015625, + "1": 431.7406005859375 + }, + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 149 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 10, + "type": "VAELoader", + "pos": { + "0": -737, + "1": 511 + }, + "size": { + "0": 210, + "1": 58 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 12, + 125 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "ae.safetensors" + ] + }, + { + "id": 11, + "type": "DualCLIPLoader", + "pos": { + "0": -741, + "1": 346 + }, + "size": { + "0": 210, + "1": 106 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 10 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader" + }, + "widgets_values": [ + "t5xxl_fp16.safetensors", + "clip_l.safetensors", + "flux" + ] + }, + { + "id": 50, + "type": "AddLabel", + "pos": { + "0": 1172, + "1": 1594 + }, + "size": { + "0": 315, + "1": 274 + }, + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 152 + }, + { + "name": "caption", + "type": "STRING", + "link": null, + "widget": { + "name": "caption" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 154 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "AddLabel" + }, + "widgets_values": [ + 10, + 2, + 48, + 32, + "white", + "black", + "FreeMono.ttf", + "output", + "up", + "" + ] + }, + { + "id": 25, + "type": "RandomNoise", + "pos": { + "0": 889, + "1": 47 + }, + "size": { + "0": 210, + "1": 82 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 37 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 0, + "fixed" + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 13, + "type": "SamplerCustomAdvanced", + "pos": { + "0": 864, + "1": 192 + }, + "size": { + "0": 236.8000030517578, + "1": 106 + }, + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 37, + "slot_index": 0 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 30, + "slot_index": 1 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 19, + "slot_index": 2 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 20, + "slot_index": 3 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 116, + "slot_index": 4 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 24 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "SamplerCustomAdvanced" + } + }, + { + "id": 22, + "type": "BasicGuider", + "pos": { + "0": 500, + "1": 213 + }, + "size": { + "0": 161.1999969482422, + "1": 46 + }, + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 144, + "slot_index": 0 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 156, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 30 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicGuider" + } + }, + { + "id": 27, + "type": "EmptySD3LatentImage", + "pos": { + "0": 194, + "1": 632 + }, + "size": { + "0": 210, + "1": 78 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "link": 112, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 113, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 116 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "EmptySD3LatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 46, + "type": "ModelSamplingFlux", + "pos": { + "0": 189, + "1": 927 + }, + "size": { + "0": 210, + "1": 122 + }, + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 143, + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "link": 145, + "slot_index": 1, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 146, + "slot_index": 2, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 144 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ModelSamplingFlux" + }, + "widgets_values": [ + 1.1500000000000001, + 0.5, + 1024, + 1024 + ] + }, + { + "id": 16, + "type": "KSamplerSelect", + "pos": { + "0": 198, + "1": 510 + }, + "size": { + "0": 210, + "1": 58 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 19 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "heun" + ] + }, + { + "id": 43, + "type": "ImageScale", + "pos": { + "0": -435, + "1": 934 + }, + "size": [ + 210, + 130 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 128 + }, + { + "name": "width", + "type": "INT", + "link": 157, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 158, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 129, + 151 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "nearest-exact", + 1024, + 1024, + "center" + ] + }, + { + "id": 51, + "type": "FluxDeGuidance", + "pos": { + "0": 198, + "1": 363 + }, + "size": { + "0": 211.60000610351562, + "1": 58 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 155 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 156 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "FluxDeGuidance" + }, + "widgets_values": [ + 3.5 + ] + }, + { + "id": 17, + "type": "BasicScheduler", + "pos": { + "0": 187, + "1": 769 + }, + "size": { + "0": 210, + "1": 106 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 141, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 20, + 139 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 20, + 1 + ] + }, + { + "id": 40, + "type": "LoadImage", + "pos": { + "0": -836, + "1": 968 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 128 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "fortnite.jpeg", + "image" + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": { + "0": -270, + "1": 359 + }, + "size": { + "0": 285.6000061035156, + "1": 98.67007446289062 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 10 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 155 + ], + "slot_index": 0 + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "an anime character with blue hair" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 45, + "type": "ConfigureRefFlux", + "pos": { + "0": 190, + "1": 1111 + }, + "size": { + "0": 210, + "1": 146 + }, + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 150 + }, + { + "name": "latent", + "type": "LATENT", + "link": 138 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 139 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 143 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ConfigureRefFlux" + }, + "widgets_values": [ + 0, + 0.26, + 1 + ] + }, + { + "id": 34, + "type": "PrimitiveNode", + "pos": { + "0": -741, + "1": 634 + }, + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 112, + 115, + 145, + 157 + ], + "slot_index": 0, + "widget": { + "name": "width" + } + } + ], + "title": "width", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 1024, + "fixed" + ], + "color": "#323", + "bgcolor": "#535" + }, + { + "id": 35, + "type": "PrimitiveNode", + "pos": { + "0": -740, + "1": 788 + }, + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "links": [ + 113, + 114, + 146, + 158 + ], + "slot_index": 0, + "widget": { + "name": "height" + } + } + ], + "title": "height", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 1024, + "fixed" + ], + "color": "#323", + "bgcolor": "#535" + } + ], + "links": [ + [ + 9, + 8, + 0, + 9, + 0, + "IMAGE" + ], + [ + 10, + 11, + 0, + 6, + 0, + "CLIP" + ], + [ + 12, + 10, + 0, + 8, + 1, + "VAE" + ], + [ + 19, + 16, + 0, + 13, + 2, + "SAMPLER" + ], + [ + 20, + 17, + 0, + 13, + 3, + "SIGMAS" + ], + [ + 24, + 13, + 0, + 8, + 0, + "LATENT" + ], + [ + 30, + 22, + 0, + 13, + 1, + "GUIDER" + ], + [ + 37, + 25, + 0, + 13, + 0, + "NOISE" + ], + [ + 112, + 34, + 0, + 27, + 0, + "INT" + ], + [ + 113, + 35, + 0, + 27, + 1, + "INT" + ], + [ + 114, + 35, + 0, + 30, + 2, + "INT" + ], + [ + 115, + 34, + 0, + 30, + 1, + "INT" + ], + [ + 116, + 27, + 0, + 13, + 4, + "LATENT" + ], + [ + 117, + 12, + 0, + 38, + 0, + "MODEL" + ], + [ + 125, + 10, + 0, + 42, + 1, + "VAE" + ], + [ + 128, + 40, + 0, + 43, + 0, + "IMAGE" + ], + [ + 129, + 43, + 0, + 42, + 0, + "IMAGE" + ], + [ + 134, + 38, + 0, + 30, + 0, + "MODEL" + ], + [ + 138, + 42, + 0, + 45, + 1, + "LATENT" + ], + [ + 139, + 17, + 0, + 45, + 2, + "SIGMAS" + ], + [ + 141, + 30, + 0, + 17, + 0, + "MODEL" + ], + [ + 143, + 45, + 0, + 46, + 0, + "MODEL" + ], + [ + 144, + 46, + 0, + 22, + 0, + "MODEL" + ], + [ + 145, + 34, + 0, + 46, + 1, + "INT" + ], + [ + 146, + 35, + 0, + 46, + 2, + "INT" + ], + [ + 149, + 47, + 0, + 48, + 0, + "IMAGE" + ], + [ + 150, + 38, + 0, + 45, + 0, + "MODEL" + ], + [ + 151, + 43, + 0, + 49, + 0, + "IMAGE" + ], + [ + 152, + 8, + 0, + 50, + 0, + "IMAGE" + ], + [ + 153, + 49, + 0, + 47, + 0, + "IMAGE" + ], + [ + 154, + 50, + 0, + 47, + 1, + "IMAGE" + ], + [ + 155, + 6, + 0, + 51, + 0, + "CONDITIONING" + ], + [ + 156, + 51, + 0, + 22, + 1, + "CONDITIONING" + ], + [ + 157, + 34, + 0, + 43, + 1, + "INT" + ], + [ + 158, + 35, + 0, + 43, + 2, + "INT" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.0834705943388736, + "offset": [ + 83.6465795192072, + -710.630378328987 + ] + }, + "groupNodes": {} + }, + "version": 0.4 +} \ No newline at end of file diff --git a/flux/layers.py b/flux/layers.py new file mode 100644 index 0000000..4c61449 --- /dev/null +++ b/flux/layers.py @@ -0,0 +1,165 @@ +from typing import List +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor + +from comfy.ldm.flux.math import apply_rope +from comfy.ldm.flux.layers import SingleStreamBlock as OriginalSingleStreamBlock, DoubleStreamBlock as OriginalDoubleStreamBlock + + +import torch +from einops import rearrange +from torch import Tensor +from comfy.ldm.modules.attention import optimized_attention + +# from ..utils.rave_rope_attention import rave_rope_attention +# from ..utils.rave_attention import rave_attention +# from ..utils.joint_attention import joint_attention + + +def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, skip_rope: bool= False) -> Tensor: + if not skip_rope: + q, k = apply_rope(q, k, pe) + heads = q.shape[1] + x = optimized_attention(q, k, v, heads, skip_reshape=True) + return x + + +def ref_attention2(q: Tensor, k:Tensor, v:Tensor, pe: Tensor, ref_pes: List[Tensor], transformer_options, timesteps) -> Tensor: + # rope tgt + timesteps = timesteps.item() + q_tgt, k_tgt = apply_rope(q[[0]], k[[0]], pe) + strength = transformer_options['REF_OPTIONS']['strength'] + + # rope refs + q_ref = q[[-1],:,256:,:] + k_ref = k[[-1],:,256:,:] + + _, k_ref1 = apply_rope(q_ref, k_ref, ref_pes[0]) + # k_ref1 = k_ref1[[-1], :, 256:, :] * strength + + # if 0.85 <= timesteps: + + # k_ref1[:, :, :, 16+0 :16+0 +28+8-1] *= 0.1 + # k_ref1[:, :, :, 16+56:16+56+28+8-1] *= 0.1 + + # elif 0.6 <= timesteps < 0.85: + + # k_ref1[:, :, :, 16+0 :16+0 +28-20-1] *= 0.1 + # k_ref1[:, :, :, 16+56:16+56+28-20-1] *= 0.1 + + # _, k_ref2 = apply_rope(q_ref, k_ref, ref_pes[1]) + # k_ref2 = k_ref2[[-1], :, 256:, :] + v_ref = v[[-1],:, 256:, :] + + k_tgt = torch.cat([k_tgt, k_ref1], dim=2) + v_tgt = torch.cat([v[[0]], v_ref], dim=2) + tgt_out = attention(q_tgt, k_tgt, v_tgt, pe, skip_rope=True) + + ref_out = attention(q[[-1]], k[[-1]], v[[-1]], pe) + + return torch.cat([tgt_out, ref_out], dim=0) + + +def ref_attention(q: Tensor, k:Tensor, v:Tensor, pe: Tensor, ref_pes: List[Tensor], transformer_options, timesteps) -> Tensor: + # rope tgt + timesteps = timesteps.item() + strength = transformer_options['REF_OPTIONS']['strength'] + + # rope refs + q_ref1 = torch.cat([q[[0]], q[[-1], :,256:,:]], dim=2) + k_ref1 = torch.cat([k[[0]], k[[-1], :,256:,:]*strength], dim=2) + v_ref = torch.cat([v[[0]], v[[-1], :,256:,:]], dim=2) + + _, k_ref1 = apply_rope(q_ref1, k_ref1, ref_pes[1]) + + tgt_out = attention(q[[0]], k_ref1, v_ref, pe, skip_rope=True) + + ref_out = attention(q[[-1]], k[[-1]], v[[-1]], pe) + + return torch.cat([tgt_out, ref_out], dim=0) + + +class DoubleStreamBlock(OriginalDoubleStreamBlock): + def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, ref_pes: None | List[Tensor], timestep, transformer_options={}): + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + # prepare image for attention + img_modulated = self.img_norm1(img) + img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift + img_qkv = self.img_attn.qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) + + # prepare txt for attention + txt_modulated = self.txt_norm1(txt) + txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift + txt_qkv = self.txt_attn.qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) + + + # run actual attention + q = torch.cat((txt_q, img_q), dim=2) + k = torch.cat((txt_k, img_k), dim=2) + v = torch.cat((txt_v, img_v), dim=2) + ref_options = transformer_options.get('REF_OPTIONS', None) + if ref_options is not None and ref_pes is not None: + attn = ref_attention(q, k, v, pe, ref_pes, transformer_options, timestep[0]) + else: + attn = attention(q, k, v, pe=pe) + + txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] + txt_attn = txt_attn[0:1].repeat(img_attn.shape[0], 1, 1) + + # img_attn.shape [16, 2304, 3072] + + # calculate the img bloks + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + + # calculate the txt bloks + txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) + return img, txt + + +class SingleStreamBlock(OriginalSingleStreamBlock): + def forward(self, x: Tensor, vec: Tensor, pe: Tensor, ref_pes, timestep, transformer_options={}) -> Tensor: + mod, _ = self.modulation(vec) + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + + ref_options = transformer_options.get('REF_OPTIONS', None) + if ref_options is not None and ref_pes is not None: + attn = ref_attention(q, k, v, pe, ref_pes, transformer_options, timestep[0]) + else: + attn = attention(q, k, v, pe=pe) + # txt_attn, img_attn = attn[:, :256], attn[:, 256:] + + # txt_attn = temporal_attention(txt_attn, self.num_heads, transformer_options) + # attn[:, :256] = txt_attn + + # img_attn = interframe_attention(img_attn, self.num_heads, transformer_options) + # attn[:, 256:] = img_attn + + # compute activation in mlp stream, cat again and run second linear layer + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output + + +def inject_blocks(diffusion_model): + for i, block in enumerate(diffusion_model.double_blocks): + block.__class__ = DoubleStreamBlock + block.idx = i + + for i, block in enumerate(diffusion_model.single_blocks): + block.__class__ = SingleStreamBlock + block.idx = i + + return diffusion_model \ No newline at end of file diff --git a/flux/model.py b/flux/model.py new file mode 100644 index 0000000..00af850 --- /dev/null +++ b/flux/model.py @@ -0,0 +1,147 @@ +#Original code can be found on: https://github.com/black-forest-labs/flux +from typing import List +import torch +from torch import Tensor, nn + +from comfy.ldm.flux.layers import timestep_embedding +from comfy.ldm.flux.model import Flux as OriginalFlux + +from einops import rearrange, repeat +import comfy.ldm.common_dit + +from ..utils.noise_utils import add_noise_flux, add_noise + + +class Flux(OriginalFlux): + def forward_orig( + self, + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor = None, + control=None, + transformer_options = {}, + ref_img_ids: List[Tensor] | None = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + # running on sequences img + img = self.img_in(img) + vec = self.time_in(timestep_embedding(timesteps, 256).to(img.dtype)) + if self.params.guidance_embed: + if guidance is None: + raise ValueError("Didn't get guidance strength for guidance distilled model.") + vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype)) + + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + ref_pes = None + if ref_img_ids is not None: + ids = torch.cat((txt_ids[:-1], img_ids), dim=1) + pe = self.pe_embedder(ids) + ref_pe1 = self.pe_embedder(ref_img_ids[0]) + ref_id2 = torch.cat((txt_ids[-1:], img_ids, ref_img_ids[1]), dim=1) + ref_pe2 = self.pe_embedder(ref_id2) + ref_pes = [ref_pe1, ref_pe2] + else: + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + + for i, block in enumerate(self.double_blocks): + + img, txt = block(img=img, txt=txt, vec=vec, pe=pe, ref_pes=ref_pes, timestep=timesteps, transformer_options=transformer_options) + + if control is not None: # Controlnet + control_i = control.get("input") + if i < len(control_i): + add = control_i[i] + if add is not None: + img += add + + img = torch.cat((txt, img), 1) + for i, block in enumerate(self.single_blocks): + img = block(img, vec=vec, pe=pe, ref_pes=ref_pes, timestep=timesteps, transformer_options=transformer_options) + + if control is not None: # Controlnet + control_o = control.get("output") + if i < len(control_o): + add = control_o[i] + if add is not None: + img[:, txt.shape[1] :, ...] += add + + img = img[:, txt.shape[1] :, ...] + + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + return img + + def _get_img_ids(self, x, bs, h_len, w_len, h_start, h_end, w_start, w_end): + img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype) + img_ids[..., 1] = img_ids[..., 1] + torch.linspace(h_start, h_end - 1, steps=h_len, device=x.device, dtype=x.dtype)[:, None] + img_ids[..., 2] = img_ids[..., 2] + torch.linspace(w_start, w_end - 1, steps=w_len, device=x.device, dtype=x.dtype)[None, :] + img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) + return img_ids + + def forward(self, x, timestep, context, y, guidance, control=None, transformer_options={}, **kwargs): + bs, c, h, w = x.shape + transformer_options['original_shape'] = x.shape + patch_size = 2 + x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size)) + + h_len = ((h + (patch_size // 2)) // patch_size) + w_len = ((w + (patch_size // 2)) // patch_size) + + img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size) + + txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype) + img_ids_orig = self._get_img_ids(x, bs, h_len, w_len, 0, h_len, 0, w_len) + + ref_options = transformer_options.get('REF_OPTIONS', None) + + perform_ref = False + if ref_options is not None: + ref_start_percent = ref_options.get('start_percent', 0) + ref_end_percent = ref_options.get('end_percent', -1) + sigma_percents = ref_options.get('sigma_to_percent', {}) + step_percent = sigma_percents[timestep[0].item()] + perform_ref = ref_start_percent <= step_percent < ref_end_percent + sigma_to_step = ref_options.get('sigma_to_step', {}) + + ref_img_ids = None + if perform_ref: + # ref + ref_latent = ref_options['ref_latent'] + ref_latent = ref_latent.to(x.device) + sigma = ref_options.get('sigmas', [])[sigma_to_step[timestep[0].item()]].to(x.device) + ref_latent = add_noise_flux(ref_latent, torch.randn_like(ref_latent), sigma) + # noise = torch.randn_like(ref_latent) + # ref_latent = add_noise(ref_latent, noise, sigma) + ref_latent = rearrange(ref_latent, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size) + img = torch.cat([img, ref_latent], dim=0) # this won't work with cfg + # horizontal translation + diff = 2 + ref_img_id1 = self._get_img_ids(x, bs, h_len, w_len, diff+h_len, diff+2*h_len, w_len, diff+2*w_len) + ref_img_id2 = self._get_img_ids(x, bs, h_len, w_len, 0, h_len, 0, 2*w_len) + ref_img_ids = [ref_img_id1, ref_img_id2] + + timestep = timestep.repeat(2) + txt_ids = txt_ids.repeat(2, 1, 1) + context = context.repeat(2, 1, 1) + y = y.repeat(2, 1) + guidance = guidance.repeat(2) + + out = self.forward_orig(img, img_ids_orig, context, txt_ids, timestep, y, guidance, control, transformer_options=transformer_options, ref_img_ids=ref_img_ids) + + if perform_ref: + out = out[-1:] + return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w] + + +def inject_flux(diffusion_model: OriginalFlux): + diffusion_model.__class__ = Flux + diffusion_model.is_ref = True + return diffusion_model \ No newline at end of file diff --git a/nodes/__pycache__/flip_sigmas_node.cpython-310.pyc b/nodes/__pycache__/flip_sigmas_node.cpython-310.pyc index 4b573545d0252083d560767c2a71f86df44ee5f3..25ef45ee2aaa25438ed9528505e2979b2bb306fd 100644 GIT binary patch delta 19 YcmZoff=jG*M00OzDja=WE03&V$00000 delta 19 YcmZoff=jG*M00Nirja=WE03&+?`v3p{ diff --git a/nodes/__pycache__/flux_deguidance_node.cpython-310.pyc b/nodes/__pycache__/flux_deguidance_node.cpython-310.pyc index 194101473fa3855b7cac00e405c350512b2c74f2..cb3ad85f2b43e7f2ebf0686483bf3d4235702c9c 100644 GIT binary patch delta 19 ZcmZ3+zKoqKpO=@50SM%pHge5j1^_4z1M&a> delta 19 ZcmZ3+zKoqKpO=@50SIJrHge5j1^_3O1KI!p diff --git a/nodes/__pycache__/influx_model_pred_node.cpython-310.pyc b/nodes/__pycache__/influx_model_pred_node.cpython-310.pyc index d18af5afeffabc9c61e0f31a289b9006bf492d39..49c45984f3818a1e1e72241e556a49b55e637f2f 100644 GIT binary patch delta 595 zcmY+A%S%*o6vyY>duQf%ulL~tUoi!Vk1jq)!AA`_YFRpt*=ZWQnarJ1G^(kFla@Qg z3-TpDOcnW+FVMKSrZ_aNDV_i(fhpjTBB6R>osO!FC04zj&<-!jq3OvjprRrzL0fiK z-9Gtf-cU(8+7DVRqG2B@rU*GMvqHWMuDxToGg zge510ct;ObLNdz~=N>_21h^oT1Stx_g!H6Mn8qXL8^rj*eG7B=*Zl}}+{iEAZc)P& z#q~uj)1s~ljNqUA^uQ>cn=!^i#uUH9k4Ua?`h?IQ!^3=lqdA=)XD?>V;JZS_n(WN=n<67_1U+xdLA_u4e^MBlshh|eqr U|J=(d?Hygl*c-|Ycq2Fc3sePd{{R30 diff --git a/nodes/add_flow_flux_node.py b/nodes/add_flow_flux_node.py new file mode 100644 index 0000000..df944b4 --- /dev/null +++ b/nodes/add_flow_flux_node.py @@ -0,0 +1,29 @@ +import comfy.sd +import comfy.model_sampling +import comfy.latent_formats +import nodes + + +class AddFluxFlowNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + "flow": ("FLOW",), + }} + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "advanced/model" + + def patch(self, model, flow): + m = model.clone() + model_options = {**model.model_options} + model.model_options = model_options + transformer_options = {**model.model_options.get('transformer_options', {})} + model.model_options['transformer_options'] = transformer_options + + transformer_options['FLOW'] = flow + + return (m, ) + diff --git a/nodes/apply_rave_node.py b/nodes/apply_rave_node.py new file mode 100644 index 0000000..9a541c2 --- /dev/null +++ b/nodes/apply_rave_node.py @@ -0,0 +1,41 @@ + +default_attn = { + 'inputs': [True] * 10, + 'input_idxs': list(range(10)), + 'middle_0': True, + 'outputs': [True] * 12, + 'output_idxs': list(range(12)) +} + + +class ApplyFluxRaveAttentionNode: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "model": ("MODEL",), + "grid_size": ("INT", {"default": 3, "min": 1, "max": 10}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + "attn_override": ("ATTN_OVERRIDE",) + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + + CATEGORY = "attention" + + def apply(self, model, grid_size, seed, attn_override=default_attn): + model = model.clone() + + transformer_options = {**model.model_options.get('transformer_options', {})} + model.model_options['transformer_options'] = transformer_options + + transformer_options['RAVE'] = { + "grid_size": grid_size, + "seed": seed, + } + + return (model, ) \ No newline at end of file diff --git a/nodes/apply_ref_flux.py b/nodes/apply_ref_flux.py new file mode 100644 index 0000000..ece426f --- /dev/null +++ b/nodes/apply_ref_flux.py @@ -0,0 +1,65 @@ +from ..flux.model import inject_flux +from ..flux.layers import inject_blocks + + +class ApplyRefFluxNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + + CATEGORY = "reference" + + def apply(self, model): + # if hasattr(model.model.diffusion_model, 'is_ref') and model.model.diffusion_model.is_ref: + # return (model,) + inject_flux(model.model.diffusion_model) + inject_blocks(model.model.diffusion_model) + return (model,) + + +class ConfigureRefFluxNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "latent": ("LATENT",), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "end_percent": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + "strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}), + "sigmas": ("SIGMAS",) + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + + CATEGORY = "reference" + + def apply(self, + model, + latent, + start_percent, + end_percent, + strength, + sigmas): + model = model.clone() + sigma_to_percent = { sigma.item(): idx/len(sigmas) for idx, sigma in enumerate(sigmas)} + sigma_to_step = { sigma.item(): idx for idx, sigma in enumerate(sigmas)} + + transformer_options = model.model_options.get('transformer_options', {}) + transformer_options = { **transformer_options } + process_latent_in = model.get_model_object("process_latent_in") + transformer_options['REF_OPTIONS'] = { + 'ref_latent': process_latent_in(latent['samples']), + 'start_percent': start_percent, + 'end_percent': end_percent, + 'sigma_to_percent': sigma_to_percent, + 'sigma_to_step': sigma_to_step, + 'strength': strength, + 'sigmas': sigmas, + } + model.model_options['transformer_options'] = transformer_options + + return (model,) diff --git a/utils/module_utils.py b/utils/module_utils.py new file mode 100644 index 0000000..a22b0f9 --- /dev/null +++ b/utils/module_utils.py @@ -0,0 +1,7 @@ +def isinstance_str(x: object, cls_name: str): + for _cls in x.__class__.__mro__: + if _cls.__name__ == cls_name: + return True + + return False + diff --git a/utils/noise_utils.py b/utils/noise_utils.py new file mode 100644 index 0000000..75d1401 --- /dev/null +++ b/utils/noise_utils.py @@ -0,0 +1,26 @@ +import torch + + +def get_alphacumprod(sigma): + return 1 / ((sigma * sigma) + 1) + + +def add_noise(src_latent, noise, sigma): + alphas_cumprod = get_alphacumprod(sigma) + + sqrt_alpha_prod = alphas_cumprod ** 0.5 + sqrt_alpha_prod = sqrt_alpha_prod.flatten() + while len(sqrt_alpha_prod.shape) < len(src_latent.shape): + sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) + + sqrt_one_minus_alpha_prod = (1 - alphas_cumprod) ** 0.5 + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() + while len(sqrt_one_minus_alpha_prod.shape) < len(src_latent.shape): + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) + + noisy_samples = sqrt_alpha_prod * src_latent + sqrt_one_minus_alpha_prod * noise + return noisy_samples + + +def add_noise_flux(src_latent, noise, sigma): + return sigma * noise + (1.0 - sigma) * src_latent \ No newline at end of file diff --git a/utils/rave_attention.py b/utils/rave_attention.py new file mode 100644 index 0000000..4a8ee60 --- /dev/null +++ b/utils/rave_attention.py @@ -0,0 +1,117 @@ +import random +import torch +from einops import rearrange + +from comfy.ldm.modules.attention import optimized_attention + +from .rave_utils import grid_to_list, list_to_grid, shuffle_indices, shuffle_tensors2 + + +def padding_count(n_frames, grid_frame_count): + remainder = n_frames % grid_frame_count + if remainder == 0: + return 0 + else: + difference = grid_frame_count - remainder + return difference + + +def unpatchify(x, h, w, p=2): + x = rearrange(x, 'b (h w) (p q d) -> b (h p) (w q) d', h=h, p=p, q=p) + return rearrange(x, 'b h w d -> b (h w) d') + + +def patchify(x, h, w, p=2): + return rearrange(x, 'b (p h q w) d -> b (h w) (p q d)', p=p, q=p, h=h, w=w) + + +def rave_attention(q, k, v, extra_options, n_heads): + # get h,w + batch_size, sequence_length, dim = q.shape + shape = extra_options['original_shape'] + oh, ow = shape[-2:] + ratio = oh/ow + d = sequence_length + w = int((d/ratio)**(0.5)) + h = int(d/w) + + rave_opts = extra_options.get('RAVE', {}) + grid_size = rave_opts.get('grid_size', 2) + seed = rave_opts.get('seed', 1) + len_conds = len(extra_options['cond_or_uncond']) + n_frames = batch_size // len_conds + original_n_frames = n_frames + + grid_frame_count = grid_size * grid_size + n_padding_frames = padding_count(n_frames, grid_frame_count) + if n_padding_frames > 0: + random.seed(seed) + cond_qs = [] + cond_ks = [] + cond_vs = [] + padding_frames = [random.randint( + 0, n_frames-1) for _ in range(n_padding_frames)] + for cond_idx in range(len_conds): + start, end = cond_idx*n_frames, (cond_idx+1)*n_frames + cond_q = q[start:end] + cond_q = torch.cat([cond_q, cond_q[padding_frames]]) + cond_qs.append(cond_q) + cond_k = k[start:end] + cond_k = torch.cat([cond_k, cond_k[padding_frames]]) + cond_ks.append(cond_k) + cond_v = v[start:end] + cond_v = torch.cat([cond_v, cond_v[padding_frames]]) + cond_vs.append(cond_v) + + q = torch.cat(cond_qs) + k = torch.cat(cond_ks) + v = torch.cat(cond_vs) + + n_frames = n_frames + n_padding_frames + + q = rearrange(q, 'b (h w) c -> b h w c', h=h, w=w) + k = rearrange(k, 'b (h w) c -> b h w c', h=h, w=w) + v = rearrange(v, 'b (h w) c -> b h w c', h=h, w=w) + + target_indexes = shuffle_indices(n_frames, seed=seed) + + original_indexes = list(range(n_frames)) + qs = [] + ks = [] + vs = [] + + for i in range(len_conds): + start, end = i*n_frames, (i+1)*n_frames + q[start:end] = shuffle_tensors2( + q[start:end], original_indexes, target_indexes) + qs.append(list_to_grid(q[start:end], grid_size)) + k[start:end] = shuffle_tensors2( + k[start:end], original_indexes, target_indexes) + ks.append(list_to_grid(k[start:end], grid_size)) + v[start:end] = shuffle_tensors2( + v[start:end], original_indexes, target_indexes) + vs.append(list_to_grid(v[start:end], grid_size)) + + q = torch.cat(qs) + k = torch.cat(ks) + v = torch.cat(vs) + + q = rearrange(q, 'b h w c -> b (h w) c') + k = rearrange(k, 'b h w c -> b (h w) c') + v = rearrange(v, 'b h w c -> b (h w) c') + out = optimized_attention(q, k, v, n_heads, None) + + gh, gw = grid_size*h, grid_size*w + out = rearrange(out, 'b (h w) c -> b h w c', h=gh, w=gw) + out = grid_to_list(out, grid_size) + out = rearrange(out, 'b h w c -> b (h w) c') + + outs = [] + for i in range(len_conds): + start, end = i*n_frames, (i+1)*n_frames + cond_out = shuffle_tensors2( + out[start:end], target_indexes, original_indexes) + cond_out = cond_out[:original_n_frames] + outs.append(cond_out) + + return torch.cat(outs) diff --git a/utils/rave_rope_attention.py b/utils/rave_rope_attention.py new file mode 100644 index 0000000..e123ee1 --- /dev/null +++ b/utils/rave_rope_attention.py @@ -0,0 +1,130 @@ +import random +import torch +from torch import Tensor +from einops import rearrange + +from comfy.ldm.flux.math import apply_rope +from comfy.ldm.modules.attention import optimized_attention + +from .rave_utils import grid_to_list, list_to_grid, shuffle_indices, shuffle_tensors2 + + +def padding_count(n_frames, grid_frame_count): + remainder = n_frames % grid_frame_count + if remainder == 0: + return 0 + else: + difference = grid_frame_count - remainder + return difference + + +def attention(q: torch.Tensor, k: Tensor, v: Tensor, pe: Tensor,) -> Tensor: + q, k = apply_rope(q, k, pe) + heads = q.shape[1] + x = optimized_attention(q, k, v, heads, skip_reshape=True) + return x + + +def rave_rope_attention(q, k, v, txt_q, txt_k, txt_v, pe, extra_options, n_heads, txt_shape): + q = rearrange(q, 'b n r d -> b r (n d)') + k = rearrange(k, 'b n r d -> b r (n d)') + v = rearrange(v, 'b n r d -> b r (n d)') + rave_opts = extra_options.get('RAVE', {}) + grid_size = rave_opts.get('grid_size', 2) + seed = rave_opts.get('seed', 1) + batch_size, sequence_length, dim = q.shape + len_conds = len(extra_options['cond_or_uncond']) + n_frames = batch_size // len_conds + original_n_frames = n_frames + + grid_frame_count = grid_size * grid_size + n_padding_frames = padding_count(n_frames, grid_frame_count) + if n_padding_frames > 0: + random.seed(seed) + cond_qs = [] + cond_ks = [] + cond_vs = [] + padding_frames = [random.randint( + 0, n_frames-1) for _ in range(n_padding_frames)] + for cond_idx in range(len_conds): + start, end = cond_idx*n_frames, (cond_idx+1)*n_frames + cond_q = q[start:end] + cond_q = torch.cat([cond_q, cond_q[padding_frames]]) + cond_qs.append(cond_q) + cond_k = k[start:end] + cond_k = torch.cat([cond_k, cond_k[padding_frames]]) + cond_ks.append(cond_k) + cond_v = v[start:end] + cond_v = torch.cat([cond_v, cond_v[padding_frames]]) + cond_vs.append(cond_v) + + q = torch.cat(cond_qs) + k = torch.cat(cond_ks) + v = torch.cat(cond_vs) + + n_frames = n_frames + n_padding_frames + + # get h,w + shape = extra_options['original_shape'] + oh, ow = shape[-2:] + ratio = oh/ow + d = sequence_length + w = int((d/ratio)**(0.5)) + h = int(d/w) + + q = rearrange(q, 'b (h w) c -> b h w c', h=h, w=w) + k = rearrange(k, 'b (h w) c -> b h w c', h=h, w=w) + v = rearrange(v, 'b (h w) c -> b h w c', h=h, w=w) + + target_indexes = shuffle_indices(n_frames, seed=seed) + + original_indexes = list(range(n_frames)) + qs = [] + ks = [] + vs = [] + + for i in range(len_conds): + start, end = i*n_frames, (i+1)*n_frames + q[start:end] = shuffle_tensors2( + q[start:end], original_indexes, target_indexes) + qs.append(list_to_grid(q[start:end], grid_size)) + k[start:end] = shuffle_tensors2( + k[start:end], original_indexes, target_indexes) + ks.append(list_to_grid(k[start:end], grid_size)) + v[start:end] = shuffle_tensors2( + v[start:end], original_indexes, target_indexes) + vs.append(list_to_grid(v[start:end], grid_size)) + + q = torch.cat(qs) + k = torch.cat(ks) + v = torch.cat(vs) + + q = rearrange(q, 'b h w (n d) -> b n (h w) d', n=n_heads) + k = rearrange(k, 'b h w (n d) -> b n (h w) d', n=n_heads) + v = rearrange(v, 'b h w (n d) -> b n (h w) d', n=n_heads) + # out = optimized_attention(q, k, v, n_heads, None) + txt_q = txt_q[:len(q)] + txt_k = txt_k[:len(q)] + txt_v = txt_v[:len(q)] + q = torch.cat((txt_q, q), dim=2) + k = torch.cat((txt_k, k), dim=2) + v = torch.cat((txt_v, v), dim=2) + out = attention(q, k, v, pe) + + txt_attn, out = out[:, :txt_shape], out[:, txt_shape:] + txt_attn = txt_attn.repeat(grid_size*grid_size, 1, 1) + + gh, gw = grid_size*h, grid_size*w + out = rearrange(out, 'b (h w) c -> b h w c', h=gh, w=gw) + out = grid_to_list(out, grid_size) + out = rearrange(out, 'b h w c -> b (h w) c') + + outs = [] + for i in range(len_conds): + start, end = i*n_frames, (i+1)*n_frames + cond_out = shuffle_tensors2( + out[start:end], target_indexes, original_indexes) + cond_out = cond_out[:original_n_frames] + outs.append(cond_out) + + return txt_attn, torch.cat(outs) diff --git a/utils/rave_utils.py b/utils/rave_utils.py new file mode 100644 index 0000000..0a0197f --- /dev/null +++ b/utils/rave_utils.py @@ -0,0 +1,57 @@ +import torch +import random + + +def shuffle_indices(size, seed=None): + if seed is not None: + random.seed(seed) + indices = list(range(size)) + random.shuffle(indices) + return indices + + +def shuffle_tensors2(tensor, current_indices, target_indices): + tensor_dict = {current_idx: t for current_idx, + t in zip(current_indices, tensor)} + shuffled_tensors = [tensor_dict[current_idx] + for current_idx in target_indices] + return torch.stack(shuffled_tensors) + + +def grid_to_list(tensor, grid_size): + frame_count = len(tensor) * grid_size * grid_size + flattened_list = [flatten_grid(grid.unsqueeze( + 0), [grid_size, grid_size]) for grid in tensor] + list_tensor = torch.cat(flattened_list, dim=-2) + return torch.cat(torch.chunk(list_tensor, frame_count, dim=-2), dim=0) + + +def list_to_grid(tensor, grid_size): + grid_frame_count = grid_size * grid_size + grid_count = len(tensor) // grid_frame_count + flat_grids = [torch.cat([a for a in tensor[i * grid_frame_count:(i + 1) + * grid_frame_count]], dim=-2).unsqueeze(0) for i in range(grid_count)] + unflattened_grids = [unflatten_grid( + flat_grid, [grid_size, grid_size]) for flat_grid in flat_grids] + return torch.cat(unflattened_grids, dim=0) + + +def flatten_grid(x, grid_shape): + B, H, W, C = x.size() + hs, ws = grid_shape + img_h = H // hs + flattened = torch.cat(torch.split(x, img_h, dim=1), dim=2) + return flattened + + +def unflatten_grid(x, grid_shape): + ''' + x: B x C x H x W + ''' + B, H, W, C = x.size() + hs, ws = grid_shape + img_w = W // (ws) + + unflattened = torch.cat(torch.split(x, img_w, dim=2), dim=1) + + return unflattened \ No newline at end of file