diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7e99e36 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +*.pyc \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..e055aa3 --- /dev/null +++ b/README.md @@ -0,0 +1,9 @@ +# ComfyUI Prompt-Relay + +WORK IN PROGRESS + + + +Original project: + +https://gordonchen19.github.io/Prompt-Relay/ \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/example_workflows/prompt_relay_test.json b/example_workflows/prompt_relay_test.json new file mode 100644 index 0000000..1225fa3 --- /dev/null +++ b/example_workflows/prompt_relay_test.json @@ -0,0 +1,2115 @@ +{ + "id": "7f5e0c56-93b4-4937-b7f2-efd0f1853e33", + "revision": 0, + "last_node_id": 606, + "last_link_id": 1523, + "nodes": [ + { + "id": 550, + "type": "LTXVAudioVAEDecode", + "pos": [ + 407.59478603729974, + 4973.902967051035 + ], + "size": [ + 203.9, + 46 + ], + "flags": {}, + "order": 30, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 1416 + }, + { + "label": "Audio VAE", + "name": "audio_vae", + "type": "VAE", + "link": 1459 + } + ], + "outputs": [ + { + "name": "Audio", + "type": "AUDIO", + "links": [ + 1503 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVAudioVAEDecode", + "cnr_id": "comfy-core", + "ver": "0.16.0" + }, + "widgets_values": [] + }, + { + "id": 228, + "type": "SetNode", + "pos": [ + -1731.9566833205165, + 4899.953651561091 + ], + "size": [ + 210, + 60 + ], + "flags": { + "collapsed": true + }, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "VAE", + "type": "VAE", + "link": 1425 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": [] + } + ], + "title": "Set_video_vae", + "properties": { + "Node name for S&R": "SetNode", + "aux_id": "kijai/ComfyUI-KJNodes", + "previousName": "video_vae" + }, + "widgets_values": [ + "video_vae" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 575, + "type": "GetNode", + "pos": [ + 418.9799376524158, + 5068.570518529913 + ], + "size": [ + 210, + 60 + ], + "flags": { + "collapsed": true + }, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 1459 + ] + } + ], + "title": "Get_audio_vae", + "properties": { + "Node name for S&R": "GetNode", + "aux_id": "kijai/ComfyUI-KJNodes" + }, + "widgets_values": [ + "audio_vae" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 574, + "type": "SetNode", + "pos": [ + -1608.5190578719967, + 5067.579621493003 + ], + "size": [ + 210, + 60 + ], + "flags": { + "collapsed": true + }, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "VAE", + "type": "VAE", + "link": 1458 + } + ], + "outputs": [ + { + "name": "*", + "type": "*", + "links": null + } + ], + "title": "Set_audio_vae", + "properties": { + "Node name for S&R": "SetNode", + "aux_id": "kijai/ComfyUI-KJNodes", + "previousName": "audio_vae" + }, + "widgets_values": [ + "audio_vae" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 547, + "type": "LTXVEmptyLatentAudio", + "pos": [ + -286.91094566381753, + 5567.027711460725 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "label": "Audio VAE", + "name": "audio_vae", + "type": "VAE", + "link": 1460 + }, + { + "name": "frames_number", + "type": "INT", + "widget": { + "name": "frames_number" + }, + "link": 1463 + } + ], + "outputs": [ + { + "name": "Latent", + "type": "LATENT", + "links": [ + 1411 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVEmptyLatentAudio", + "cnr_id": "comfy-core", + "ver": "0.16.0" + }, + "widgets_values": [ + 489, + 25, + 1 + ] + }, + { + "id": 358, + "type": "VisualizeSigmasKJ", + "pos": [ + -246.25912055429907, + 5385.802247690478 + ], + "size": [ + 270, + 102 + ], + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "sigmas", + "type": "SIGMAS", + "link": 1364 + } + ], + "outputs": [ + { + "name": "sigmas_out", + "type": "SIGMAS", + "links": [ + 1445 + ] + }, + { + "name": "image", + "type": "IMAGE", + "links": [ + 1447 + ] + } + ], + "properties": { + "Node name for S&R": "VisualizeSigmasKJ", + "cnr_id": "comfyui-kjnodes", + "ver": "422e10c846768faf26197e8d2b44ceddec736df2" + }, + "widgets_values": [ + 0, + -1 + ] + }, + { + "id": 583, + "type": "GetNode", + "pos": [ + -313.15214395101117, + 5789.465597336383 + ], + "size": [ + 210, + 34 + ], + "flags": { + "collapsed": true + }, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 1465 + ] + } + ], + "title": "Get_video_vae", + "properties": { + "Node name for S&R": "GetNode", + "aux_id": "kijai/ComfyUI-KJNodes" + }, + "widgets_values": [ + "video_vae" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 584, + "type": "LoadImage", + "pos": [ + -1136.9553425517483, + 5887.5374449489145 + ], + "size": [ + 282.798828125, + 314 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 1470 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage", + "cnr_id": "comfy-core", + "ver": "0.16.2" + }, + "widgets_values": [ + "oldman_upscaled.png", + "image" + ] + }, + { + "id": 586, + "type": "LTXVPreprocess", + "pos": [ + -783.0470385021329, + 5851.695675069141 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 1470 + } + ], + "outputs": [ + { + "name": "output_image", + "type": "IMAGE", + "links": [ + 1473 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVPreprocess", + "cnr_id": "comfy-core", + "ver": "0.16.2" + }, + "widgets_values": [ + 30 + ] + }, + { + "id": 588, + "type": "LTX2SamplingPreviewOverride", + "pos": [ + 30.020029718971266, + 4826.277247838771 + ], + "size": [ + 316.7017578125, + 98 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1508 + }, + { + "name": "latent_upscale_model", + "shape": 7, + "type": "LATENT_UPSCALE_MODEL", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1475 + ] + } + ], + "properties": { + "Node name for S&R": "LTX2SamplingPreviewOverride", + "cnr_id": "comfyui-kjnodes", + "ver": "a9ab19d2a7933e03f72bc00103863f06b9483300" + }, + "widgets_values": [ + 8 + ] + }, + { + "id": 562, + "type": "PreviewImage", + "pos": [ + 457.0626012270293, + 5407.331924216174 + ], + "size": [ + 402.6036823594029, + 348.06865481820023 + ], + "flags": {}, + "order": 26, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1447 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage", + "cnr_id": "comfy-core", + "ver": "0.16.0" + }, + "widgets_values": [] + }, + { + "id": 549, + "type": "LTXVSeparateAVLatent", + "pos": [ + 410.29756241023017, + 4866.429072900162 + ], + "size": [ + 196.9, + 46 + ], + "flags": {}, + "order": 28, + "mode": 0, + "inputs": [ + { + "name": "av_latent", + "type": "LATENT", + "link": 1442 + } + ], + "outputs": [ + { + "name": "video_latent", + "type": "LATENT", + "links": [ + 1443 + ] + }, + { + "name": "audio_latent", + "type": "LATENT", + "links": [ + 1416 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVSeparateAVLatent", + "cnr_id": "comfy-core", + "ver": "0.16.0" + }, + "widgets_values": [] + }, + { + "id": 471, + "type": "VAEDecode", + "pos": [ + 418.5843206856501, + 4758.88746125511 + ], + "size": [ + 195.85402430293288, + 46 + ], + "flags": {}, + "order": 29, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 1443 + }, + { + "name": "vae", + "type": "VAE", + "link": 1288 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 1502 + ] + } + ], + "properties": { + "Node name for S&R": "VAEDecode", + "cnr_id": "comfy-core", + "ver": "0.9.1" + }, + "widgets_values": [] + }, + { + "id": 553, + "type": "VAELoaderKJ", + "pos": [ + -2075.071776414608, + 4872.248332282984 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 1425 + ] + } + ], + "properties": { + "Node name for S&R": "VAELoaderKJ", + "cnr_id": "comfyui-kjnodes", + "ver": "13101465f14df2b725694ad5e16ed29c838ec8f2" + }, + "widgets_values": [ + "LTX23\\vae\\LTX23_video_vae_bf16.safetensors", + "main_device", + "bf16" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 470, + "type": "VAELoaderKJ", + "pos": [ + -2069.145364496889, + 5036.235976414522 + ], + "size": [ + 422.5309084617047, + 106 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 1458 + ] + } + ], + "properties": { + "Node name for S&R": "VAELoaderKJ", + "cnr_id": "comfyui-kjnodes", + "ver": "47ca2b2c26358615974162059061e9dc78be6715" + }, + "widgets_values": [ + "LTX23\\vae\\LTX23_audio_vae_bf16.safetensors", + "main_device", + "bf16" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 416, + "type": "DualCLIPLoader", + "pos": [ + -2059.196643703852, + 5205.442554054825 + ], + "size": [ + 478.8229382589643, + 130 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 1506 + ] + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader", + "cnr_id": "comfy-core", + "ver": "0.8.2" + }, + "widgets_values": [ + "gemma_3_12B_it.safetensors", + "LTX23_text_projection_bf16.safetensors", + "ltxv", + "default" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 414, + "type": "UNETLoader", + "pos": [ + -2082.75718554929, + 4716.577345251733 + ], + "size": [ + 659.469041562375, + 82 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1510 + ] + } + ], + "properties": { + "Node name for S&R": "UNETLoader", + "cnr_id": "comfy-core", + "ver": "0.8.2" + }, + "widgets_values": [ + "LTX23\\ltx-2.3-22b-distilled-1.1_transformer_only_fp8_scaled.safetensors", + "default" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 604, + "type": "VHS_VideoCombine", + "pos": [ + 972.9907313321646, + 4759.096993398134 + ], + "size": [ + 809.9995141076606, + 1137.9995141076606 + ], + "flags": {}, + "order": 31, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1502 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": 1503 + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "title": "v2", + "properties": { + "Node name for S&R": "VHS_VideoCombine", + "cnr_id": "comfyui-videohelpersuite", + "ver": "3234937ff5f3ca19068aaba5042771514de2429d" + }, + "widgets_values": { + "frame_rate": 25, + "loop_count": 0, + "filename_prefix": "ltx23", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "ltx23_00242-audio.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 25, + "workflow": "ltx23_00242.png", + "fullpath": "S:\\comfy\\ComfyUI\\output\\ltx23_00242-audio.mp4" + } + } + } + }, + { + "id": 236, + "type": "GetNode", + "pos": [ + 395.8731619414867, + 4679.203601839582 + ], + "size": [ + 210, + 34 + ], + "flags": { + "collapsed": true + }, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 1288 + ] + } + ], + "title": "Get_video_vae", + "properties": { + "Node name for S&R": "GetNode", + "aux_id": "kijai/ComfyUI-KJNodes" + }, + "widgets_values": [ + "video_vae" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 582, + "type": "LTXVImgToVideoInplaceKJ", + "pos": [ + -137.03552859633817, + 5757.824606133932 + ], + "size": [ + 294.98046875, + 146 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "VAE", + "link": 1465 + }, + { + "name": "latent", + "type": "LATENT", + "link": 1466 + }, + { + "label": "image_1", + "name": "num_images.image_1", + "shape": 7, + "type": "IMAGE", + "link": 1473 + } + ], + "outputs": [ + { + "name": "latent", + "type": "LATENT", + "links": [ + 1467 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVImgToVideoInplaceKJ", + "cnr_id": "comfyui-kjnodes", + "ver": "a9ab19d2a7933e03f72bc00103863f06b9483300" + }, + "widgets_values": [ + "1", + 1, + 0 + ] + }, + { + "id": 577, + "type": "EmptyLTXVLatentVideo", + "pos": [ + -713.9999651729453, + 5357.060213306923 + ], + "size": [ + 270, + 130 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "length", + "type": "INT", + "widget": { + "name": "length" + }, + "link": 1462 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 1466, + 1507 + ] + } + ], + "properties": { + "Node name for S&R": "EmptyLTXVLatentVideo", + "cnr_id": "comfy-core", + "ver": "0.16.2" + }, + "widgets_values": [ + 640, + 640, + 489, + 1 + ] + }, + { + "id": 548, + "type": "LTXVConcatAVLatent", + "pos": [ + 64.87910868730921, + 5630.929291195288 + ], + "size": [ + 183.9, + 46 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "video_latent", + "type": "LATENT", + "link": 1467 + }, + { + "name": "audio_latent", + "type": "LATENT", + "link": 1411 + } + ], + "outputs": [ + { + "name": "latent", + "type": "LATENT", + "links": [ + 1446 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVConcatAVLatent", + "cnr_id": "comfy-core", + "ver": "0.16.0" + }, + "widgets_values": [] + }, + { + "id": 561, + "type": "SamplerCustom", + "pos": [ + 72.26226107618345, + 5056.932833069262 + ], + "size": [ + 271.125, + 495.125 + ], + "flags": {}, + "order": 27, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1475 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 1438 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 1439 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 1444 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 1445 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 1446 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 1442 + ] + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null + } + ], + "properties": { + "Node name for S&R": "SamplerCustom", + "cnr_id": "comfy-core", + "ver": "0.16.0" + }, + "widgets_values": [ + true, + 17, + "fixed", + 1 + ] + }, + { + "id": 164, + "type": "LTXVConditioning", + "pos": [ + -636.0686261237182, + 5079.243988964387 + ], + "size": [ + 270, + 86.66666666666667 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 1519 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 1522 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 1438 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 1439 + ] + } + ], + "properties": { + "Node name for S&R": "LTXVConditioning", + "cnr_id": "comfy-core", + "ver": "0.3.56", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + 25 + ] + }, + { + "id": 578, + "type": "PrimitiveNode", + "pos": [ + -990.3320059357658, + 5375.99584385698 + ], + "size": [ + 210, + 82 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "INT", + "type": "INT", + "widget": { + "name": "length" + }, + "links": [ + 1462, + 1463 + ] + } + ], + "title": "length", + "properties": { + "Run widget replace on values": false + }, + "widgets_values": [ + 489, + "fixed" + ] + }, + { + "id": 576, + "type": "GetNode", + "pos": [ + -442.13004655880457, + 5523.228231955801 + ], + "size": [ + 210, + 34 + ], + "flags": { + "collapsed": true + }, + "order": 9, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 1460 + ] + } + ], + "title": "Get_audio_vae", + "properties": { + "Node name for S&R": "GetNode", + "aux_id": "kijai/ComfyUI-KJNodes" + }, + "widgets_values": [ + "audio_vae" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 514, + "type": "BasicScheduler", + "pos": [ + -678.3456808389026, + 5608.670071476283 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1523 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 1364 + ] + } + ], + "properties": { + "Node name for S&R": "BasicScheduler", + "cnr_id": "comfy-core", + "ver": "0.10.0" + }, + "widgets_values": [ + "linear_quadratic", + 8, + 1 + ] + }, + { + "id": 154, + "type": "KSamplerSelect", + "pos": [ + -275.0811547364064, + 4958.11872694023 + ], + "size": [ + 270, + 68.88020833333334 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 1444 + ] + } + ], + "properties": { + "Node name for S&R": "KSamplerSelect", + "cnr_id": "comfy-core", + "ver": "0.3.56", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "euler_ancestral" + ] + }, + { + "id": 606, + "type": "LoraLoaderModelOnly", + "pos": [ + -1698.1787435019987, + 4560.154090823786 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 14, + "mode": 4, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1510 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1511 + ] + } + ], + "properties": { + "Node name for S&R": "LoraLoaderModelOnly" + }, + "widgets_values": [ + "LTXV\\ltx-2.3-22b-distilled-lora-384-1.1.safetensors", + 0.5 + ] + }, + { + "id": 542, + "type": "PathchSageAttentionKJ", + "pos": [ + -1356.9805324213662, + 4680.081199476252 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1511 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1505, + 1523 + ] + } + ], + "properties": { + "Node name for S&R": "PathchSageAttentionKJ", + "cnr_id": "comfyui-kjnodes", + "ver": "092a0fa72b725bb8cb544b5fe4ffdc61f2f84f7f" + }, + "widgets_values": [ + "auto", + false + ] + }, + { + "id": 605, + "type": "PromptRelayEncode", + "pos": [ + -992.6994628573964, + 4512.863666279692 + ], + "size": [ + 449.28124513207035, + 490.1880798983175 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1505 + }, + { + "name": "clip", + "type": "CLIP", + "link": 1506 + }, + { + "name": "latent", + "type": "LATENT", + "link": 1507 + } + ], + "outputs": [ + { + "name": "model", + "type": "MODEL", + "links": [ + 1508 + ] + }, + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 1519, + 1521 + ] + } + ], + "properties": { + "Node name for S&R": "PromptRelayEncode" + }, + "widgets_values": [ + "A single continuous cinematic shot of an old man and a red panda", + "An old man says \"Well well... what do we have here?\" and takes his hat off\n|\nA red panda jumps onto the old mans lap\n|\nThe old man feeds the red panda an ice cream cone\n|\nThe red panda grabs the ice cream cone and runs away, while the old man says \"oh, okay then... keep it\"", + "", + 0.001 + ] + }, + { + "id": 420, + "type": "ConditioningZeroOut", + "pos": [ + -876.8236196126293, + 5147.96788547305 + ], + "size": [ + 204.1750030517578, + 26 + ], + "flags": { + "collapsed": true + }, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 1521 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 1522 + ] + } + ], + "properties": { + "Node name for S&R": "ConditioningZeroOut", + "cnr_id": "comfy-core", + "ver": "0.8.2" + }, + "widgets_values": [] + } + ], + "links": [ + [ + 1288, + 236, + 0, + 471, + 1, + "VAE" + ], + [ + 1364, + 514, + 0, + 358, + 0, + "SIGMAS" + ], + [ + 1411, + 547, + 0, + 548, + 1, + "LATENT" + ], + [ + 1416, + 549, + 1, + 550, + 0, + "LATENT" + ], + [ + 1425, + 553, + 0, + 228, + 0, + "VAE" + ], + [ + 1438, + 164, + 0, + 561, + 1, + "CONDITIONING" + ], + [ + 1439, + 164, + 1, + 561, + 2, + "CONDITIONING" + ], + [ + 1442, + 561, + 0, + 549, + 0, + "LATENT" + ], + [ + 1443, + 549, + 0, + 471, + 0, + "LATENT" + ], + [ + 1444, + 154, + 0, + 561, + 3, + "SAMPLER" + ], + [ + 1445, + 358, + 0, + 561, + 4, + "SIGMAS" + ], + [ + 1446, + 548, + 0, + 561, + 5, + "LATENT" + ], + [ + 1447, + 358, + 1, + 562, + 0, + "IMAGE" + ], + [ + 1458, + 470, + 0, + 574, + 0, + "VAE" + ], + [ + 1459, + 575, + 0, + 550, + 1, + "VAE" + ], + [ + 1460, + 576, + 0, + 547, + 0, + "VAE" + ], + [ + 1462, + 578, + 0, + 577, + 0, + "INT" + ], + [ + 1463, + 578, + 0, + 547, + 1, + "INT" + ], + [ + 1465, + 583, + 0, + 582, + 0, + "VAE" + ], + [ + 1466, + 577, + 0, + 582, + 1, + "LATENT" + ], + [ + 1467, + 582, + 0, + 548, + 0, + "LATENT" + ], + [ + 1470, + 584, + 0, + 586, + 0, + "IMAGE" + ], + [ + 1473, + 586, + 0, + 582, + 2, + "IMAGE" + ], + [ + 1475, + 588, + 0, + 561, + 0, + "MODEL" + ], + [ + 1502, + 471, + 0, + 604, + 0, + "IMAGE" + ], + [ + 1503, + 550, + 0, + 604, + 1, + "AUDIO" + ], + [ + 1505, + 542, + 0, + 605, + 0, + "MODEL" + ], + [ + 1506, + 416, + 0, + 605, + 1, + "CLIP" + ], + [ + 1507, + 577, + 0, + 605, + 2, + "LATENT" + ], + [ + 1508, + 605, + 0, + 588, + 0, + "MODEL" + ], + [ + 1510, + 414, + 0, + 606, + 0, + "MODEL" + ], + [ + 1511, + 606, + 0, + 542, + 0, + "MODEL" + ], + [ + 1519, + 605, + 1, + 164, + 0, + "CONDITIONING" + ], + [ + 1521, + 605, + 1, + 420, + 0, + "CONDITIONING" + ], + [ + 1522, + 420, + 0, + 164, + 1, + "CONDITIONING" + ], + [ + 1523, + 542, + 0, + 514, + 0, + "MODEL" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.5337067413064844, + "offset": [ + 1747.2728363385809, + -4106.714187459462 + ] + }, + "frontendVersion": "1.43.1", + "workflowRendererVersion": "LG", + "prompt": { + "1": { + "inputs": { + "ckpt_name": "ltx-av-step-1751000_vocoder_24K.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "Load Checkpoint" + } + }, + "2": { + "inputs": { + "gemma_path": "gemma-3-12b-it-qat-q4_0-unquantized_readout_proj/model/model.safetensors", + "ltxv_path": "ltx-av-step-1751000_vocoder_24K.safetensors", + "max_length": 1024 + }, + "class_type": "LTXVGemmaCLIPModelLoader", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ Gemma 3 Model Loader" + } + }, + "3": { + "inputs": { + "text": "A medium close-up shot features a Caucasian man with a closely shaven head and face, wearing a black baseball cap with \"PNTR\" in white letters on the front, and a dark grey t-shirt with \"JUST DO IT\" visible across his chest. A small black microphone is clipped to his shirt collar. He is positioned slightly to the left of the frame, looking intently downwards and to his right, his eyes focused off-camera. His facial expression is one of deep concentration, with his brow slightly furrowed. As he looks down, a quick sniff sound is heard, and then he speaks with a deep male voice and a slightly frustrated tone, saying, \"I think it's so bad.\" The camera remains static throughout, maintaining a shallow depth of field, which keeps the man in sharp focus while the background is softly blurred, revealing a light-colored wall with white wooden shelving or trim, and a partially open white wooden door on the right. After a brief pause, another short, audible sniff is heard. The man then continues to speak, his voice maintaining the same quality, as he states, \"So bad. So bad.\" He elaborates further, emphasizing his point with a final statement, \"This got to be, it's got to be the worst tool I've ever seen.\"", + "clip": [ + "2", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "4": { + "inputs": { + "text": "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, wrong hand count, artifacts around text, unreadable text on shirt or hat, incorrect lettering on cap (β€œPNTR”), incorrect t-shirt slogan (β€œJUST DO IT”), missing microphone, misplaced microphone, inconsistent perspective, camera shake, incorrect depth of field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, smiling, laughing, exaggerated sadness, wrong gaze direction, eyes looking at camera, mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, missing sniff sounds, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, missing door or shelves, missing shallow depth of field, flat lighting, inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts.", + "clip": [ + "2", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "8": { + "inputs": { + "sampler_name": "euler" + }, + "class_type": "KSamplerSelect", + "_meta": { + "title": "KSamplerSelect" + } + }, + "9": { + "inputs": { + "steps": 20, + "max_shift": 2.05, + "base_shift": 0.95, + "stretch": true, + "terminal": 0.1, + "latent": [ + "28", + 0 + ] + }, + "class_type": "LTXVScheduler", + "_meta": { + "title": "LTXVScheduler" + } + }, + "11": { + "inputs": { + "noise_seed": 10 + }, + "class_type": "RandomNoise", + "_meta": { + "title": "RandomNoise" + } + }, + "12": { + "inputs": { + "samples": [ + "29", + 0 + ], + "vae": [ + "1", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "13": { + "inputs": { + "ckpt_name": "ltx-av-step-1751000_vocoder_24K.safetensors" + }, + "class_type": "LTXVAudioVAELoader", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ LTXV Audio VAE Loader" + } + }, + "14": { + "inputs": { + "samples": [ + "29", + 1 + ], + "audio_vae": [ + "13", + 0 + ] + }, + "class_type": "LTXVAudioVAEDecode", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ LTXV Audio VAE Decode" + } + }, + "15": { + "inputs": { + "frame_rate": [ + "23", + 0 + ], + "loop_count": 0, + "filename_prefix": "AnimateDiff", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "images": [ + "12", + 0 + ], + "audio": [ + "14", + 0 + ] + }, + "class_type": "VHS_VideoCombine", + "_meta": { + "title": "Video Combine πŸŽ₯πŸ…₯πŸ…—πŸ…’" + } + }, + "17": { + "inputs": { + "skip_blocks": "29", + "model": [ + "28", + 1 + ], + "positive": [ + "22", + 0 + ], + "negative": [ + "22", + 1 + ], + "parameters": [ + "18", + 0 + ] + }, + "class_type": "MultimodalGuider", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ Multimodal Guider" + } + }, + "18": { + "inputs": { + "modality": "VIDEO", + "cfg": 3, + "stg": 0, + "rescale": 0, + "modality_scale": 3, + "parameters": [ + "19", + 0 + ] + }, + "class_type": "GuiderParameters", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ Guider Parameters" + } + }, + "19": { + "inputs": { + "modality": "AUDIO", + "cfg": 7, + "stg": 0, + "rescale": 0, + "modality_scale": 3 + }, + "class_type": "GuiderParameters", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ Guider Parameters" + } + }, + "21": { + "inputs": { + "audioUI": "", + "audio": [ + "14", + 0 + ] + }, + "class_type": "PreviewAudio", + "_meta": { + "title": "PreviewAudio" + } + }, + "22": { + "inputs": { + "frame_rate": [ + "23", + 0 + ], + "positive": [ + "3", + 0 + ], + "negative": [ + "4", + 0 + ] + }, + "class_type": "LTXVConditioning", + "_meta": { + "title": "LTXVConditioning" + } + }, + "23": { + "inputs": { + "value": 25 + }, + "class_type": "FloatConstant", + "_meta": { + "title": "Float Constant" + } + }, + "26": { + "inputs": { + "frames_number": [ + "27", + 0 + ], + "frame_rate": [ + "42", + 0 + ], + "batch_size": 1 + }, + "class_type": "LTXVEmptyLatentAudio", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ LTXV Empty Latent Audio" + } + }, + "27": { + "inputs": { + "value": 105 + }, + "class_type": "INTConstant", + "_meta": { + "title": "INT Constant" + } + }, + "28": { + "inputs": { + "video_latent": [ + "43", + 0 + ], + "audio_latent": [ + "26", + 0 + ], + "model": [ + "44", + 0 + ] + }, + "class_type": "LTXVConcatAVLatent", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ LTXV Concat AV Latent" + } + }, + "29": { + "inputs": { + "av_latent": [ + "41", + 0 + ], + "model": [ + "28", + 1 + ] + }, + "class_type": "LTXVSeparateAVLatent", + "_meta": { + "title": "πŸ…›πŸ…£πŸ…§ LTXV Separate AV Latent" + } + }, + "41": { + "inputs": { + "noise": [ + "11", + 0 + ], + "guider": [ + "17", + 0 + ], + "sampler": [ + "8", + 0 + ], + "sigmas": [ + "9", + 0 + ], + "latent_image": [ + "28", + 0 + ] + }, + "class_type": "SamplerCustomAdvanced", + "_meta": { + "title": "SamplerCustomAdvanced" + } + }, + "42": { + "inputs": { + "a": [ + "23", + 0 + ] + }, + "class_type": "CM_FloatToInt", + "_meta": { + "title": "FloatToInt" + } + }, + "43": { + "inputs": { + "width": 768, + "height": 512, + "length": [ + "27", + 0 + ], + "batch_size": 1 + }, + "class_type": "EmptyLTXVLatentVideo", + "_meta": { + "title": "EmptyLTXVLatentVideo" + } + }, + "44": { + "inputs": { + "torch_compile": true, + "disable_backup": false, + "model": [ + "1", + 0 + ] + }, + "class_type": "LTXVSequenceParallelMultiGPUPatcher", + "_meta": { + "title": "LTXVSequenceParallelMultiGPUPatcher" + } + }, + "45": { + "inputs": { + "frame_idx": 0, + "strength": 1 + }, + "class_type": "LTXVAddGuide", + "_meta": { + "title": "LTXVAddGuide" + } + } + }, + "comfy_fork_version": "feature/av_inference@a6994ed1", + "VHS_latentpreview": true, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..e44d048 --- /dev/null +++ b/nodes.py @@ -0,0 +1,107 @@ +import logging + +from comfy_api.latest import io + +from .prompt_relay import ( + get_raw_tokenizer, + map_token_indices, + build_segments, + create_mask_fn, + distribute_segment_lengths, +) + +from .patches import detect_model_type, apply_patches + +log = logging.getLogger(__name__) + + +class PromptRelayEncode(io.ComfyNode): + """Encodes temporal local prompts and patches the model for Prompt Relay.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PromptRelayEncode", + display_name="Prompt Relay Encode", + category="conditioning/prompt_relay", + description=( + "Encodes a global prompt combined with temporal local prompts and patches the model " + "for Prompt Relay temporal control. Local prompts are separated by |. " + "Use a standard CLIPTextEncode for the negative prompt." + ), + inputs=[ + io.Model.Input("model"), + io.Clip.Input("clip"), + io.Latent.Input("latent", tooltip="Empty latent video β€” dimensions are read from its shape."), + io.String.Input( + "global_prompt", multiline=True, default="", + tooltip="Conditions the entire video. Anchors persistent characters, objects, and scene context.", + ), + io.String.Input( + "local_prompts", multiline=True, default="", + tooltip="Ordered prompts for each temporal segment, separated by |", + ), + io.String.Input( + "segment_lengths", default="", + tooltip="Comma-separated pixel space frame counts per segment. Leave empty to auto-distribute evenly.", + ), + io.Float.Input( + "epsilon", default=1e-3, min=1e-6, max=0.99, step=1e-4, + tooltip="Penalty decay parameter. Values below ~0.1 all produce sharp boundaries (paper default 0.001). For softer transitions, try 0.5 or higher.", + ), + ], + outputs=[ + io.Model.Output(display_name="model"), + io.Conditioning.Output(display_name="positive"), + ], + ) + + @classmethod + def execute(cls, model, clip, latent, global_prompt, local_prompts, segment_lengths, epsilon) -> io.NodeOutput: + locals_list = [p.strip() for p in local_prompts.split("|") if p.strip()] + if not locals_list: + raise ValueError("At least one local prompt is required (separate with |)") + + arch, patch_size, temporal_stride = detect_model_type(model) + + parsed_lengths = None + if segment_lengths.strip(): + pixel_lengths = [int(x.strip()) for x in segment_lengths.split(",") if x.strip()] + parsed_lengths = [max(1, round(p / temporal_stride)) for p in pixel_lengths] + + raw_tokenizer = get_raw_tokenizer(clip) + full_prompt, token_ranges = map_token_indices(raw_tokenizer, global_prompt, locals_list) + + log.info("[PromptRelay] Global: tokens [0:%d] (%d tokens)", token_ranges[0][0], token_ranges[0][0]) + for i, (s, e) in enumerate(token_ranges): + log.info("[PromptRelay] Segment %d: tokens [%d:%d] (%d tokens)", i, s, e, e - s) + + conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(full_prompt)) + + samples = latent["samples"] + latent_frames = samples.shape[2] + tokens_per_frame = (samples.shape[3] // patch_size[1]) * (samples.shape[4] // patch_size[2]) + + effective_lengths = distribute_segment_lengths(len(locals_list), latent_frames, parsed_lengths) + + log.info( + "[PromptRelay] Latent: %d frames, %d tokens/frame, segments: %s", + latent_frames, tokens_per_frame, effective_lengths, + ) + + q_token_idx = build_segments(token_ranges, effective_lengths, epsilon) + mask_fn = create_mask_fn(q_token_idx, tokens_per_frame, latent_frames) + + patched = model.clone() + apply_patches(patched, arch, mask_fn) + + return io.NodeOutput(patched, conditioning) + + +NODE_CLASS_MAPPINGS = { + "PromptRelayEncode": PromptRelayEncode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PromptRelayEncode": "Prompt Relay Encode", +} diff --git a/patches.py b/patches.py new file mode 100644 index 0000000..1404146 --- /dev/null +++ b/patches.py @@ -0,0 +1,167 @@ +import types +import torch +import comfy.ldm.modules.attention + + +def _masked_attention(q, k, v, heads, mask, transformer_options={}, **kwargs): + # Bypass wrap_attn (sage/etc may ignore masks) by calling attention_pytorch directly. + return comfy.ldm.modules.attention.attention_pytorch( + q, k, v, heads, mask=mask, + _inside_attn_wrapper=True, + transformer_options=transformer_options, + **kwargs, + ) + + +def _wan_t2v_forward(self, mask_fn, x, context, transformer_options={}, **kwargs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(context)) + v = self.v(context) + + mask = mask_fn(q, k, transformer_options) + if mask is not None: + x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask, + transformer_options=transformer_options) + else: + x = comfy.ldm.modules.attention.optimized_attention( + q, k, v, heads=self.num_heads, transformer_options=transformer_options, + ) + return self.o(x) + + +def _wan_i2v_forward(self, mask_fn, x, context, context_img_len, transformer_options={}, **kwargs): + context_img = context[:, :context_img_len] + context_text = context[:, context_img_len:] + + q = self.norm_q(self.q(x)) + + k_img = self.norm_k_img(self.k_img(context_img)) + v_img = self.v_img(context_img) + img_x = comfy.ldm.modules.attention.optimized_attention( + q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options, + ) + + k = self.norm_k(self.k(context_text)) + v = self.v(context_text) + + mask = mask_fn(q, k, transformer_options) + if mask is not None: + x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask, + transformer_options=transformer_options) + else: + x = comfy.ldm.modules.attention.optimized_attention( + q, k, v, heads=self.num_heads, transformer_options=transformer_options, + ) + + return self.o(x + img_x) + + +def _ltx_forward(self, mask_fn, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}): + from comfy.ldm.lightricks.model import apply_rotary_emb + + is_self_attn = context is None + context = x if is_self_attn else context + + q = self.q_norm(self.to_q(x)) + k = self.k_norm(self.to_k(context)) + v = self.to_v(context) + + if pe is not None: + q = apply_rotary_emb(q, pe) + k = apply_rotary_emb(k, pe if k_pe is None else k_pe) + + if not is_self_attn: + temporal_mask = mask_fn(q, k, transformer_options) + if temporal_mask is not None: + mask = temporal_mask if mask is None else mask + temporal_mask + + if mask is None: + out = comfy.ldm.modules.attention.optimized_attention( + q, k, v, self.heads, attn_precision=self.attn_precision, + transformer_options=transformer_options, + ) + else: + out = _masked_attention(q, k, v, self.heads, mask=mask, + attn_precision=self.attn_precision, + transformer_options=transformer_options) + + if self.to_gate_logits is not None: + gate_logits = self.to_gate_logits(x) + b, t, _ = out.shape + out = out.view(b, t, self.heads, self.dim_head) + out = out * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1) + out = out.view(b, t, self.heads * self.dim_head) + + return self.to_out(out) + + +class _CrossAttnPatch: + """Descriptor that binds (impl, mask_fn) as a method onto a cross-attn module.""" + + def __init__(self, impl, mask_fn): + self.impl = impl + self.mask_fn = mask_fn + + def __get__(self, obj, objtype=None): + impl, mask_fn = self.impl, self.mask_fn + + def wrapped(self_module, *args, **kwargs): + return impl(self_module, mask_fn, *args, **kwargs) + + return types.MethodType(wrapped, obj) + + +def detect_model_type(model): + """Return (arch, patch_size, temporal_stride) for latent geometry. + + temporal_stride is the VAE's pixelβ†’latent temporal compression factor, + used to convert user-facing pixel frame counts to latent frames. + """ + diff_model = model.model.diffusion_model + + if hasattr(diff_model, "patch_size") and not hasattr(diff_model, "patchifier"): + return "wan", tuple(diff_model.patch_size), 4 + + if hasattr(diff_model, "patchifier"): + return "ltx", (1, 1, 1), int(diff_model.vae_scale_factors[0]) + + raise ValueError( + f"Unsupported model type: {type(diff_model).__name__}. " + f"Currently supports Wan and LTX models." + ) + + +def _check_unpatched(model_clone, key): + if key in getattr(model_clone, "object_patches", {}): + raise RuntimeError( + f"PromptRelay: cross-attention forward at '{key}' is already patched by " + "another node (e.g. KJNodes NAG). Stacking is not supported β€” remove the " + "conflicting node." + ) + + +def apply_patches(model_clone, arch, mask_fn): + diffusion_model = model_clone.get_model_object("diffusion_model") + + if arch == "wan": + from comfy.ldm.wan.model import WanI2VCrossAttention + for idx, block in enumerate(diffusion_model.blocks): + key = f"diffusion_model.blocks.{idx}.cross_attn.forward" + _check_unpatched(model_clone, key) + cross_attn = block.cross_attn + impl = _wan_i2v_forward if isinstance(cross_attn, WanI2VCrossAttention) else _wan_t2v_forward + model_clone.add_object_patch(key, _CrossAttnPatch(impl, mask_fn).__get__(cross_attn, cross_attn.__class__)) + return + + if arch == "ltx": + for idx, block in enumerate(diffusion_model.transformer_blocks): + for attr in ("attn2", "audio_attn2"): + module = getattr(block, attr, None) + if module is None: + continue + key = f"diffusion_model.transformer_blocks.{idx}.{attr}.forward" + _check_unpatched(model_clone, key) + model_clone.add_object_patch(key, _CrossAttnPatch(_ltx_forward, mask_fn).__get__(module, module.__class__)) + return + + raise ValueError(f"Unknown model arch: {arch}") diff --git a/prompt_relay.py b/prompt_relay.py new file mode 100644 index 0000000..a3d2b2a --- /dev/null +++ b/prompt_relay.py @@ -0,0 +1,165 @@ +import logging +import math +import torch + +log = logging.getLogger(__name__) + + +def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame): + """Gaussian penalty matrix [Lq, Lk] for video cross-attention (integer frame indexing).""" + offset = torch.zeros(Lq, Lk, device=device, dtype=dtype) + query_frames = torch.arange(Lq, device=device, dtype=torch.long) // tokens_per_frame + + for seg in q_token_idx: + local = seg["local_token_idx"].to(device=device) + d = (query_frames.float()[:, None] - seg["midpoint"]).abs() + cost = (torch.relu(d - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2) + offset[:, local] = cost.to(offset.dtype) + + return offset + + +def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames): + """Penalty matrix for queries that don't map to integer frames (e.g. LTXAV audio tokens).""" + offset = torch.zeros(Lq, Lk, device=device, dtype=dtype) + query_frames = torch.arange(Lq, device=device, dtype=torch.float32) * latent_frames / Lq + + for seg in q_token_idx: + local = seg["local_token_idx"].to(device=device) + d = (query_frames[:, None] - seg["midpoint"]).abs() + cost = (torch.relu(d - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2) + offset[:, local] = cost.to(offset.dtype) + + return offset + + +def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames): + """Closure: mask_fn(q, k, transformer_options) -> additive mask or None.""" + cache = {} + max_token_idx = max(int(seg["local_token_idx"].max().item()) for seg in q_token_idx) + 1 + + def mask_fn(q, k, transformer_options): + Lq, Lk = q.shape[1], k.shape[1] + + if Lq == Lk: + return None + + # Only apply on conditional pass β€” not unconditional (negative prompt) + cond_or_uncond = transformer_options.get("cond_or_uncond", []) + if 1 in cond_or_uncond and 0 not in cond_or_uncond: + return None + + grid_sizes = transformer_options.get("grid_sizes", None) + video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) if grid_sizes is not None else fallback_tokens_per_frame + video_lq = latent_frames * video_tpf + + # Skip cross-modal attention β€” text keys are padded to a fixed length β‰₯ max_token_idx and != video_lq + if Lk == video_lq or Lk < max_token_idx: + return None + + mode = "video" if Lq == video_lq else "scaled" + + key = (Lq, Lk, mode, q.device) + if key not in cache: + if mode == "video": + cost = build_temporal_cost(q_token_idx, Lq, Lk, q.device, q.dtype, video_tpf) + else: + cost = build_temporal_cost_scaled(q_token_idx, Lq, Lk, q.device, q.dtype, latent_frames) + log.info( + "[PromptRelay] Built penalty matrix (%s): Lq=%d, Lk=%d, nonzero=%d/%d", + mode, Lq, Lk, (cost > 0).sum().item(), cost.numel(), + ) + cache[key] = -cost + + return cache[key].to(q.dtype) + + return mask_fn + + +def build_segments(token_ranges, segment_lengths, epsilon=1e-3): + """Per-segment metadata (local_token_idx, midpoint, window, sigma) for the temporal penalty.""" + # Paper uses constant sigma = 1/ln(1/epsilon) regardless of segment length + sigma = 1.0 / math.log(1.0 / epsilon) if 0 < epsilon < 1 else 0.1448 + + q_token_idx = [] + frame_cursor = 0 + + for (tok_start, tok_end), L in zip(token_ranges, segment_lengths): + if L <= 0: + frame_cursor += L + continue + midpoint = (2 * frame_cursor + L) // 2 + window = max(L // 2 - 2, 0) + q_token_idx.append({ + "local_token_idx": torch.arange(tok_start, tok_end), + "midpoint": midpoint, + "window": window, + "sigma": sigma, + }) + frame_cursor += L + + return q_token_idx + + +def get_raw_tokenizer(clip): + """Extract the raw SPiece/HF tokenizer from a ComfyUI CLIP object.""" + tokenizer_wrapper = clip.tokenizer + for attr_name in dir(tokenizer_wrapper): + if attr_name.startswith("_"): + continue + inner = getattr(tokenizer_wrapper, attr_name, None) + if inner is not None and hasattr(inner, "tokenizer"): + return inner.tokenizer + + raise RuntimeError( + f"Could not find raw tokenizer on CLIP object. " + f"Known attributes: {[a for a in dir(tokenizer_wrapper) if not a.startswith('_')]}" + ) + + +def map_token_indices(raw_tokenizer, global_prompt, local_prompts): + """Tokenize global + space-prefixed locals; return (full_prompt, per-local token ranges). + + Uses incremental tokenization to avoid SentencePiece context-dependency issues. + """ + prefixed_locals = [" " + lp for lp in local_prompts] + full_prompt = global_prompt + "".join(prefixed_locals) + has_eos = getattr(raw_tokenizer, "add_eos", False) + eos_adj = 1 if has_eos else 0 + + prev_len = len(raw_tokenizer(global_prompt)["input_ids"]) - eos_adj + token_ranges = [] + built = global_prompt + + for plp in prefixed_locals: + built += plp + cur_len = len(raw_tokenizer(built)["input_ids"]) - eos_adj + if cur_len <= prev_len: + raise ValueError(f"Local prompt produced no tokens: '{plp.strip()}'") + token_ranges.append((prev_len, cur_len)) + prev_len = cur_len + + return full_prompt, token_ranges + + +def distribute_segment_lengths(num_segments, latent_frames, specified_lengths=None): + """Validate or auto-distribute segment frame counts, capped to fit within latent_frames.""" + if specified_lengths: + if len(specified_lengths) != num_segments: + raise ValueError( + f"Number of segment_lengths ({len(specified_lengths)}) " + f"must match number of local prompts ({num_segments})" + ) + lengths = specified_lengths + else: + # ceil division β€” matches reference implementation + step = -(-latent_frames // num_segments) + lengths = [step] * num_segments + + effective = [] + cursor = 0 + for L in lengths: + end = min(cursor + L, latent_frames) + effective.append(max(end - cursor, 0)) + cursor = end + return effective