diff --git a/custom_nodes/10S-Comfy-nodes/workflows/10Eros_10SNodes_I2V_v3_TiledSampler.json b/custom_nodes/10S-Comfy-nodes/workflows/10Eros_10SNodes_I2V_v3_TiledSampler.json new file mode 100644 index 0000000000000000000000000000000000000000..dd5283f9485f41e8f0fc057119a9f10b6e506354 --- /dev/null +++ b/custom_nodes/10S-Comfy-nodes/workflows/10Eros_10SNodes_I2V_v3_TiledSampler.json @@ -0,0 +1 @@ +{"id":"ee686d30-b3a6-4a60-991e-b8bba9c03fd5","revision":0,"last_node_id":798,"last_link_id":2221,"nodes":[{"id":560,"type":"GetNode","pos":[207.38018918223796,3356.709215779959],"size":[210,58],"flags":{"collapsed":true},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"FLOAT","type":"FLOAT","links":[1320]}],"title":"Get_fps_float","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":546,"type":"GetNode","pos":[-148.71863287884347,4545.6150248054455],"size":[210,50],"flags":{"collapsed":true},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","links":[1297]}],"title":"Get_length_0","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["length_0"],"color":"#1b4669","bgcolor":"#29699c"},{"id":630,"type":"SolidMask","pos":[1238.7248700088776,2874.9357428509056],"size":[270,106],"flags":{},"order":64,"mode":0,"inputs":[{"localized_name":"value","name":"value","type":"FLOAT","widget":{"name":"value"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":1566},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":1567}],"outputs":[{"localized_name":"MASK","name":"MASK","type":"MASK","links":[1457]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"SolidMask","ue_properties":{"widget_ue_connectable":{"value":true,"width":true,"height":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0,1024,1024]},{"id":279,"type":"GetNode","pos":[1254.5088394005425,3043.669001545407],"size":[210,58],"flags":{"collapsed":true},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[586]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":275,"type":"SetLatentNoiseMask","pos":[1491.3415734193702,3061.244993543124],"size":[236.13729359019888,46],"flags":{},"order":81,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":583},{"localized_name":"mask","name":"mask","type":"MASK","link":1457}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"SetLatentNoiseMask","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[]},{"id":525,"type":"RandomNoise","pos":[-137.47058558484093,3643.9027586613515],"size":[270,82],"flags":{"collapsed":true},"order":46,"mode":0,"inputs":[{"localized_name":"noise_seed","name":"noise_seed","type":"INT","widget":{"name":"noise_seed"},"link":1273}],"outputs":[{"localized_name":"NOISE","name":"NOISE","type":"NOISE","links":[1272]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"RandomNoise","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[94345222367891,"randomize"]},{"id":526,"type":"SetNode","pos":[-143.90757904721147,3583.5426406626075],"size":[210,58],"flags":{"collapsed":true},"order":54,"mode":0,"inputs":[{"name":"NOISE","type":"NOISE","link":1272}],"outputs":[{"name":"NOISE","type":"NOISE","links":[1274]}],"title":"Set_noise_seed","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"noise_seed","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["noise_seed"],"color":"#2e2e2e","bgcolor":"#242121"},{"id":214,"type":"SetNode","pos":[-1141.6115413299842,3468.352336401272],"size":[210,60],"flags":{"collapsed":true},"order":42,"mode":0,"inputs":[{"name":"VAE","type":"VAE","link":1548}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_vid_vae","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"vid_vae","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":207,"type":"SetNode","pos":[-1156.2777177683586,3334.0638385561215],"size":[236.85766571079353,58],"flags":{"collapsed":true},"order":43,"mode":0,"inputs":[{"name":"VAE","type":"VAE","link":1423}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_aud_vae","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"aud_vae","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":552,"type":"VAEDecode","pos":[743.6726958514785,3808.798438402876],"size":[254.98011089086322,46],"flags":{},"order":75,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":1905},{"localized_name":"vae","name":"vae","type":"VAE","link":1307}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[1303]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"VAEDecode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.5.2"}},"widgets_values":[]},{"id":652,"type":"Sigmas Easing","pos":[715.491530232057,4192.883029204863],"size":[447.6995362297288,154],"flags":{},"order":38,"mode":0,"inputs":[{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":1568},{"localized_name":"easing_type","name":"easing_type","type":"COMBO","widget":{"name":"easing_type"},"link":null},{"localized_name":"easing_mode","name":"easing_mode","type":"COMBO","widget":{"name":"easing_mode"},"link":null},{"localized_name":"normalize_input","name":"normalize_input","type":"BOOLEAN","widget":{"name":"normalize_input"},"link":null},{"localized_name":"normalize_output","name":"normalize_output","type":"BOOLEAN","widget":{"name":"normalize_output"},"link":null},{"localized_name":"strength","name":"strength","type":"FLOAT","widget":{"name":"strength"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[1569]}],"title":"Sigmas Easing (cubic, sine, quad) don't mess with this tbh","properties":{"cnr_id":"RES4LYF","ver":"0dc91c00c4c3fb38e7874fcd7a2a327765e8882c","Node name for S&R":"Sigmas Easing","ue_properties":{"widget_ue_connectable":{"easing_type":true,"easing_mode":true,"normalize_input":true,"normalize_output":true,"strength":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["cubic","in_out",true,true,0.7],"color":"#233","bgcolor":"#355"},{"id":551,"type":"AudioAdjustVolume","pos":[1102.859909498982,3854.587820408097],"size":[270,58],"flags":{},"order":77,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":1306},{"localized_name":"volume","name":"volume","type":"INT","widget":{"name":"volume"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[1714]}],"properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"AudioAdjustVolume","ue_properties":{"widget_ue_connectable":{"volume":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[-3],"color":"#222","bgcolor":"#000"},{"id":535,"type":"LTXVEmptyLatentAudio","pos":[-147.71312885760986,4596.138881089034],"size":[270,120],"flags":{"collapsed":false},"order":37,"mode":0,"inputs":[{"localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":1318},{"localized_name":"frames_number","name":"frames_number","type":"INT","widget":{"name":"frames_number"},"link":1295},{"localized_name":"frame_rate","name":"frame_rate","type":"INT","widget":{"name":"frame_rate"},"link":1296},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"localized_name":"Latent","name":"Latent","type":"LATENT","links":[1609]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"LTXVEmptyLatentAudio","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{"frames_number":true,"frame_rate":true},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[1,24,1]},{"id":545,"type":"GetNode","pos":[-139.3191399780596,4787.418245887232],"size":[210,58],"flags":{"collapsed":true},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","links":[1296]}],"title":"Get_fps_int","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_int"],"color":"#1b4669","bgcolor":"#29699c"},{"id":543,"type":"GetNode","pos":[-136.65609065824208,4835.279017619837],"size":[210,50],"flags":{"collapsed":true},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","links":[1295]}],"title":"Get_length_0","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["length_0"],"color":"#1b4669","bgcolor":"#29699c"},{"id":558,"type":"GetNode","pos":[-135.49266770428545,4892.598685074088],"size":[210,50],"flags":{"collapsed":true},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1318]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":510,"type":"SamplerCustomAdvanced","pos":[744.2161311585023,3451.0839655295053],"size":[212.3638671875,106],"flags":{},"order":70,"mode":0,"inputs":[{"localized_name":"noise","name":"noise","type":"NOISE","link":1274},{"localized_name":"guider","name":"guider","type":"GUIDER","link":2043},{"localized_name":"sampler","name":"sampler","type":"SAMPLER","link":1276},{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":1569},{"localized_name":"latent_image","name":"latent_image","type":"LATENT","link":1531}],"outputs":[{"localized_name":"output","name":"output","type":"LATENT","links":[]},{"localized_name":"denoised_output","name":"denoised_output","type":"LATENT","links":[1546]}],"title":"First Pass","properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"SamplerCustomAdvanced","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#323","bgcolor":"#535"},{"id":635,"type":"GetNode","pos":[3675.6779403363385,3944.3050892612646],"size":[210,58],"flags":{"collapsed":true},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"FLOAT","type":"FLOAT","links":[1481]}],"title":"Get_fps_float","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":714,"type":"GetNode","pos":[2813.360751805555,3608.779911647415],"size":[210,60],"flags":{"collapsed":false},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[2107]}],"title":"Get_resized_image","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["resized_image"],"color":"#2a363b","bgcolor":"#3f5159"},{"id":589,"type":"GetNode","pos":[2811.237769683004,3354.9545348795436],"size":[210,58],"flags":{"collapsed":false},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1957,2109]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":548,"type":"LTXVConcatAVLatent","pos":[521.5234886410276,4088.4653662852597],"size":[270,46],"flags":{},"order":67,"mode":0,"inputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","link":2120},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","link":1609}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[1531]}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"LTXVConcatAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":553,"type":"GetNode","pos":[742.4573892813876,3895.2825936892127],"size":[210,50],"flags":{"collapsed":true},"order":9,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1307]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":554,"type":"GetNode","pos":[742.707928003558,4040.7565523013463],"size":[210,34],"flags":{"collapsed":true},"order":10,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1305]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":331,"type":"SetNode","pos":[-383.2042149286154,4791.8832983604525],"size":[210,60],"flags":{"collapsed":true},"order":58,"mode":0,"inputs":[{"name":"IMAGE","type":"IMAGE","link":2106}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_resized_image","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"resized_image","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["resized_image"],"color":"#2a363b","bgcolor":"#3f5159"},{"id":769,"type":"PreviewImage","pos":[-621.1482182304868,4877.562166931689],"size":[294.6088617022692,468.8946724679554],"flags":{},"order":57,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":2103}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"PreviewImage","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#222","bgcolor":"#000"},{"id":653,"type":"STGGuiderAdvanced","pos":[489.975564844525,4419.025736243821],"size":[684.9663455547113,358],"flags":{"collapsed":false},"order":69,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1889},{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2144},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2145},{"localized_name":"preset","name":"preset","shape":7,"type":"STG_ADVANCED_PRESET","link":null},{"localized_name":"skip_steps_sigma_threshold","name":"skip_steps_sigma_threshold","type":"FLOAT","widget":{"name":"skip_steps_sigma_threshold"},"link":null},{"localized_name":"cfg_star_rescale","name":"cfg_star_rescale","type":"BOOLEAN","widget":{"name":"cfg_star_rescale"},"link":null},{"localized_name":"sigmas","name":"sigmas","type":"STRING","widget":{"name":"sigmas"},"link":null},{"localized_name":"cfg_values","name":"cfg_values","type":"STRING","widget":{"name":"cfg_values"},"link":null},{"localized_name":"stg_scale_values","name":"stg_scale_values","type":"STRING","widget":{"name":"stg_scale_values"},"link":null},{"localized_name":"stg_rescale_values","name":"stg_rescale_values","type":"STRING","widget":{"name":"stg_rescale_values"},"link":null},{"localized_name":"stg_layers_indices","name":"stg_layers_indices","type":"STRING","widget":{"name":"stg_layers_indices"},"link":null},{"localized_name":"apply_apg","name":"apply_apg","shape":7,"type":"BOOLEAN","widget":{"name":"apply_apg"},"link":null},{"localized_name":"apg_cfg_scale","name":"apg_cfg_scale","shape":7,"type":"FLOAT","widget":{"name":"apg_cfg_scale"},"link":null},{"localized_name":"eta","name":"eta","shape":7,"type":"FLOAT","widget":{"name":"eta"},"link":null},{"localized_name":"norm_threshold","name":"norm_threshold","shape":7,"type":"FLOAT","widget":{"name":"norm_threshold"},"link":null}],"outputs":[{"localized_name":"GUIDER","name":"GUIDER","type":"GUIDER","links":[2043]}],"properties":{"cnr_id":"ComfyUI-LTXVideo","ver":"82bd963cdeb66d023bed8c99324a307020907ef8","Node name for S&R":"STGGuiderAdvanced","ue_properties":{"widget_ue_connectable":{"skip_steps_sigma_threshold":true,"cfg_star_rescale":true,"sigmas":true,"cfg_values":true,"stg_scale_values":true,"stg_rescale_values":true,"stg_layers_indices":true,"apply_apg":true,"apg_cfg_scale":true,"eta":true,"norm_threshold":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[1,false,"1.0,0.99375,0.9875,0.98125,0.9550,0.8925,0.8120,0.7150,0.6030,0.4824,0.3618,0.2412,0.1206,0.0","2,1.5,1,1,1,1,1,1,1,1,1,1,1","2,1.5,1,1,1,1,1,1,1,1,1,1,1","1,1,1,1,1,1,1,1,1,1,1,1,1","[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999]",false,1,1,1],"color":"#233","bgcolor":"#355"},{"id":534,"type":"EmptyLTXVLatentVideo","pos":[-147.71695198347226,4316.913665791615],"size":[270,146.66666666666669],"flags":{},"order":63,"mode":0,"inputs":[{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":1289},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":1290},{"localized_name":"length","name":"length","type":"INT","widget":{"name":"length"},"link":1297},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[2117,2141]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.60","Node name for S&R":"EmptyLTXVLatentVideo","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{"width":true,"height":true,"length":true},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[704,64,97,1]},{"id":527,"type":"ManualSigmas","pos":[170.64914217692814,3754.348551183275],"size":[279.41750632544984,59.3366153498655],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"sigmas","name":"sigmas","type":"STRING","widget":{"name":"sigmas"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[1568,1888]}],"title":"Sigmas First Pass","properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"ManualSigmas","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.1","input_ue_unconnectable":{}}},"widgets_values":["1.0,0.99375,0.9875,0.98125,0.9550,0.8925,0.8120,0.7150,0.6030,0.4824,0.3618,0.2412,0.1206,0.0"],"color":"#323","bgcolor":"#535"},{"id":520,"type":"KSamplerSelect","pos":[181.6643151770866,3623.4807225016107],"size":[261.46963267468345,59.13737920167023],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null}],"outputs":[{"localized_name":"SAMPLER","name":"SAMPLER","type":"SAMPLER","links":[1276]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"KSamplerSelect","ue_properties":{"widget_ue_connectable":{"sampler_name":true},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":["euler_ancestral"],"color":"#233","bgcolor":"#355"},{"id":523,"type":"LTXVConditioning","pos":[179.95106180825883,3488.4504562735087],"size":[259.4071611759441,79.50641394128222],"flags":{"collapsed":false},"order":50,"mode":0,"inputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2142},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2143},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":1320}],"outputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","links":[2148]},{"localized_name":"negative","name":"negative","type":"CONDITIONING","links":[2149]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"LTXVConditioning","ue_properties":{"widget_ue_connectable":{"frame_rate":true},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":[24],"color":"#233","bgcolor":"#355"},{"id":582,"type":"ManualSigmas","pos":[3512.1741421202964,3254.6564277665666],"size":[322.9496561279293,69.61597045898407],"flags":{},"order":13,"mode":0,"inputs":[{"localized_name":"sigmas","name":"sigmas","type":"STRING","widget":{"name":"sigmas"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[2209]}],"title":"Sigmas Upscale","properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"ManualSigmas","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["0.85, 0.7250, 0.4219, 0.0"],"color":"#323","bgcolor":"#535"},{"id":585,"type":"KSamplerSelect","pos":[2813.2370757937415,3478.790375760474],"size":[270,58],"flags":{"collapsed":false},"order":14,"mode":0,"inputs":[{"localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null}],"outputs":[{"localized_name":"SAMPLER","name":"SAMPLER","type":"SAMPLER","links":[2208]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"KSamplerSelect","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{"sampler_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["euler_ancestral_cfg_pp"],"color":"#322","bgcolor":"#533"},{"id":583,"type":"CFGGuider","pos":[3542.6365852942336,3663.4475380786234],"size":[270,98],"flags":{},"order":74,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2018},{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2137},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2138},{"localized_name":"cfg","name":"cfg","type":"FLOAT","widget":{"name":"cfg"},"link":null}],"outputs":[{"localized_name":"GUIDER","name":"GUIDER","type":"GUIDER","links":[2207]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.71","Node name for S&R":"CFGGuider","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[1]},{"id":584,"type":"GetNode","pos":[3658.1883945443683,3604.7452337560367],"size":[210,50],"flags":{"collapsed":true},"order":15,"mode":0,"inputs":[],"outputs":[{"name":"NOISE","type":"NOISE","links":[2206]}],"title":"Get_noise_seed","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["noise_seed"],"color":"#2e2e2e","bgcolor":"#242121"},{"id":550,"type":"LTXVAudioVAEDecode","pos":[736.2139899777825,3948.5308046878276],"size":[240,46],"flags":{},"order":73,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":1608},{"label":"Audio VAE","localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":1305}],"outputs":[{"localized_name":"Audio","name":"Audio","type":"AUDIO","links":[1306,1602]}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"LTXVAudioVAEDecode","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":[]},{"id":596,"type":"LTXVSeparateAVLatent","pos":[4187.725994744947,3295.4824999445127],"size":[200.43247236282696,46],"flags":{},"order":85,"mode":0,"inputs":[{"localized_name":"av_latent","name":"av_latent","type":"LATENT","link":2211}],"outputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","links":[2172]},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","links":[1611]}],"properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"LTXVSeparateAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":591,"type":"LTXVConcatAVLatent","pos":[3146.6765223307043,3745.3069037036903],"size":[270,46],"flags":{},"order":83,"mode":0,"inputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","link":2110},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","link":2214}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[2210]}],"properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"LTXVConcatAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[]},{"id":770,"type":"LTXVImgToVideoInplaceKJ","pos":[3148.455220746667,3526.853708945063],"size":[294.98046875,146],"flags":{},"order":79,"mode":0,"inputs":[{"localized_name":"vae","name":"vae","type":"VAE","link":2109},{"localized_name":"latent","name":"latent","type":"LATENT","link":2122},{"localized_name":"num_images","name":"num_images","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"num_images"},"link":null},{"localized_name":"num_images.strength_1","name":"num_images.strength_1","type":"FLOAT","widget":{"name":"num_images.strength_1"},"link":null},{"label":"image_1","localized_name":"num_images.image_1","name":"num_images.image_1","shape":7,"type":"IMAGE","link":2107},{"localized_name":"num_images.index_1","name":"num_images.index_1","shape":7,"type":"INT","widget":{"name":"num_images.index_1"},"link":null}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[2110]}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTXVImgToVideoInplaceKJ","ue_properties":{"widget_ue_connectable":{"num_images":true,"num_images.strength_1":true,"num_images.index_1":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["1",1,0],"color":"#233","bgcolor":"#355"},{"id":593,"type":"LTXVAudioVAEDecode","pos":[4186.736631469633,3668.359225683233],"size":[240,46],"flags":{},"order":87,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":1611},{"label":"Audio VAE","localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":1382}],"outputs":[{"localized_name":"Audio","name":"Audio","type":"AUDIO","links":[1387,2179]}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"LTXVAudioVAEDecode","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":[]},{"id":783,"type":"easy loraNames","pos":[-947.4148190736797,3246.3912051532157],"size":[573.4679999999979,64.3888243713377],"flags":{},"order":16,"mode":0,"inputs":[{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":null}],"outputs":[{"localized_name":"lora_name","name":"lora_name","type":"*","links":[2184,2185,2186,2187]}],"title":"Distilled Lora","properties":{"cnr_id":"comfyui-easy-use","ver":"1.3.5","Node name for S&R":"easy loraNames","ue_properties":{"widget_ue_connectable":{"lora_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors"],"color":"#222","bgcolor":"#000"},{"id":780,"type":"JWStringToFloat","pos":[188.4736608306006,2933.8859567773025],"size":[270,58],"flags":{},"order":17,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2181]}],"title":"First Pass Audio Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["0.9"],"color":"#222","bgcolor":"#000"},{"id":782,"type":"JWStringToFloat","pos":[516.2577673613606,2935.444411519735],"size":[270,58],"flags":{},"order":18,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2183]}],"title":"Upscale Video Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["0.58"],"color":"#222","bgcolor":"#000"},{"id":594,"type":"GetNode","pos":[4017.8094594305244,3645.5361672191157],"size":[210,50],"flags":{"collapsed":true},"order":19,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1942]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":592,"type":"GetNode","pos":[3998.8089536312877,3815.052219814488],"size":[210,50],"flags":{"collapsed":true},"order":20,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1382]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":598,"type":"AudioAdjustVolume","pos":[3543.226714786005,3823.3257733842042],"size":[270,58],"flags":{},"order":89,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":1387},{"localized_name":"volume","name":"volume","type":"INT","widget":{"name":"volume"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[]}],"properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"AudioAdjustVolume","ue_properties":{"widget_ue_connectable":{"volume":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0],"color":"#222","bgcolor":"#000"},{"id":781,"type":"JWStringToFloat","pos":[838.2050706633592,2936.915769378011],"size":[270,58],"flags":{"collapsed":false},"order":21,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2182]}],"title":"Upscale Audio Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["1"],"color":"#222","bgcolor":"#000"},{"id":722,"type":"LTX2LoraLoaderAdvanced","pos":[-117.5080221723439,3066.4864860588896],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":52,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1859},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2184},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2180},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[1860]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,1,0,0,0,1]},{"id":723,"type":"LTX2LoraLoaderAdvanced","pos":[201.42706107783408,3065.5746013380663],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":60,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1860},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2185},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2181},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[2088]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,0,1,1,1,0]},{"id":719,"type":"LTX2LoraLoaderAdvanced","pos":[533.6158560203178,3067.449408361733],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":51,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1852},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2186},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2183},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[1853]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,1,0,0,0,1]},{"id":718,"type":"LTX2LoraLoaderAdvanced","pos":[864.9976953461493,3064.0054448447786],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":59,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1853},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2187},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2182},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[2017]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,0,1,1,1,0]},{"id":743,"type":"LatentUpscaleModelLoader","pos":[2742.801612317037,3240.087891289923],"size":[353.96248779296855,58],"flags":{},"order":22,"mode":0,"inputs":[{"localized_name":"model_name","name":"model_name","type":"COMBO","widget":{"name":"model_name"},"link":null}],"outputs":[{"localized_name":"LATENT_UPSCALE_MODEL","name":"LATENT_UPSCALE_MODEL","type":"LATENT_UPSCALE_MODEL","links":[1956]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LatentUpscaleModelLoader","ue_properties":{"widget_ue_connectable":{"model_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx-2.3-spatial-upscaler-x2-1.1.safetensors"],"color":"#323","bgcolor":"#535"},{"id":537,"type":"CLIPTextEncode","pos":[-636.294950541902,4074.23542703082],"size":[397.87115230398985,118.93963289899375],"flags":{"collapsed":false},"order":40,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":1420},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","links":[2143]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"CLIPTextEncode","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["captions, music, transition, VR, bad quality, subtitles, text, watermark, overlay effects, pc game, yelling, console game, video game, cartoon, childish, ugly, text, blur, logo, wordmark, static, low quality, noise, white noise, bleep, censoring, censor, bleeping, beep, beeping, newscast, interview, podcast, mutant, horror, 70's, film grain, cinematic, comedy, stand-up "],"color":"#322","bgcolor":"#533"},{"id":559,"type":"GetNode","pos":[-459.24751170881257,4260.730291995303],"size":[210,58],"flags":{"collapsed":true},"order":23,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1886,2075,2118]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":755,"type":"RTXVideoSuperResolution","pos":[5042.307315358794,3574.972988633893],"size":[285.81015625,106],"flags":{},"order":88,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":2025},{"localized_name":"resize_type","name":"resize_type","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"resize_type"},"link":null},{"localized_name":"resize_type.scale","name":"resize_type.scale","type":"FLOAT","widget":{"name":"resize_type.scale"},"link":null},{"localized_name":"quality","name":"quality","type":"COMBO","widget":{"name":"quality"},"link":null}],"outputs":[{"localized_name":"upscaled_images","name":"upscaled_images","type":"IMAGE","links":[2026]}],"properties":{"cnr_id":"comfyui_nvidia_rtx_nodes","ver":"892515e3eb9a4920a131a502a047e47adca9eb0d","Node name for S&R":"RTXVideoSuperResolution"},"widgets_values":["scale by multiplier",1.3,"HIGH"]},{"id":274,"type":"LTXVAudioVAEEncode","pos":[1547.6286442388596,2949.1333592283954],"size":[201.4126953125,46],"flags":{},"order":78,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":1602},{"label":"Audio VAE","localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":586}],"outputs":[{"localized_name":"Audio Latent","name":"Audio Latent","type":"LATENT","links":[583,2212]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"LTXVAudioVAEEncode","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[]},{"id":556,"type":"LTXVSeparateAVLatent","pos":[743.8788106962104,3691.3292693154235],"size":[244.39530248670872,46],"flags":{},"order":71,"mode":0,"inputs":[{"localized_name":"av_latent","name":"av_latent","type":"LATENT","link":1546}],"outputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","links":[1900]},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","links":[1608,2213]}],"properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"LTXVSeparateAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":789,"type":"TwoWaySwitch","pos":[2100.2422289843566,3308.3526900407155],"size":[474.2561723693839,115.75201477050769],"flags":{"collapsed":false},"order":82,"mode":0,"inputs":[{"localized_name":"input_1","name":"input_1","shape":7,"type":"*","link":2212},{"localized_name":"input_2","name":"input_2","shape":7,"type":"*","link":2213},{"localized_name":"selection_setting","name":"selection_setting","type":"INT","widget":{"name":"selection_setting"},"link":null}],"outputs":[{"localized_name":"output","name":"output","type":"*","links":[2214]}],"title":"Resample Audio (2) or Force First Pass Audio (1)","properties":{"cnr_id":"controlaltai-nodes","ver":"721492b66c9cede8ae23ae10615462ad80cfd061","Node name for S&R":"TwoWaySwitch","ue_properties":{"widget_ue_connectable":{"selection_setting":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[2],"color":"#222","bgcolor":"#000"},{"id":753,"type":"LTXTextAttentionAmplifier","pos":[3146.764295574803,3858.251464179244],"size":[307.0767578125,154],"flags":{},"order":62,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2017},{"localized_name":"text_amplification","name":"text_amplification","shape":7,"type":"FLOAT","widget":{"name":"text_amplification"},"link":null},{"localized_name":"spatial_focus","name":"spatial_focus","shape":7,"type":"FLOAT","widget":{"name":"spatial_focus"},"link":null},{"localized_name":"block_index_filter","name":"block_index_filter","shape":7,"type":"STRING","widget":{"name":"block_index_filter"},"link":null},{"localized_name":"bypass","name":"bypass","shape":7,"type":"BOOLEAN","widget":{"name":"bypass"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[2018]}],"properties":{"aux_id":"TenStrip/10S-Comfy-nodes","ver":"47852934037fbf96ed6b6572dded3c1da37bf36d","Node name for S&R":"LTXTextAttentionAmplifier"},"widgets_values":[1.3,0.15000000000000002,"36-48",false,false],"color":"#223","bgcolor":"#335"},{"id":616,"type":"LTXAVTextEncoderLoader","pos":[-1484.9909480082615,3537.073274708929],"size":[318.444921875,106],"flags":{},"order":24,"mode":0,"inputs":[{"localized_name":"text_encoder","name":"text_encoder","type":"COMBO","widget":{"name":"text_encoder"},"link":null},{"localized_name":"ckpt_name","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null}],"outputs":[{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":[1418,1419,1420]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LTXAVTextEncoderLoader","ue_properties":{"widget_ue_connectable":{"text_encoder":true,"ckpt_name":true,"device":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["gemma_3_12B_it_fp8_e4m3fn.safetensors","10Eros_v1-fp8mixed_learned.safetensors","default"],"color":"#222","bgcolor":"#000"},{"id":646,"type":"CheckpointLoaderSimple","pos":[-1482.6914004943567,3359.3040731716023],"size":[270,98],"flags":{},"order":25,"mode":0,"inputs":[{"localized_name":"ckpt_name","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null}],"outputs":[{"localized_name":"MODEL","name":"MODEL","type":"MODEL","links":[1547]},{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":null},{"localized_name":"VAE","name":"VAE","type":"VAE","links":[1548]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"CheckpointLoaderSimple","ue_properties":{"widget_ue_connectable":{"ckpt_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["10Eros_v1-fp8mixed_learned.safetensors"],"color":"#432","bgcolor":"#653"},{"id":617,"type":"LTXVAudioVAELoader","pos":[-1480.7914338617873,3235.5379033147087],"size":[270,58],"flags":{},"order":26,"mode":0,"inputs":[{"localized_name":"ckpt_name","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null}],"outputs":[{"localized_name":"Audio VAE","name":"Audio VAE","type":"VAE","links":[1423]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LTXVAudioVAELoader","ue_properties":{"widget_ue_connectable":{"ckpt_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["10Eros_v1-fp8mixed_learned.safetensors"],"color":"#332922","bgcolor":"#593930"},{"id":531,"type":"ImageResizeKJv2","pos":[-623.0137075112468,4404.1052439928335],"size":[288.23754794013735,336.00000000000045],"flags":{"collapsed":true},"order":49,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":2150},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":2215},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":2216},{"localized_name":"upscale_method","name":"upscale_method","type":"COMBO","widget":{"name":"upscale_method"},"link":null},{"localized_name":"keep_proportion","name":"keep_proportion","type":"COMBO","widget":{"name":"keep_proportion"},"link":null},{"localized_name":"pad_color","name":"pad_color","type":"STRING","widget":{"name":"pad_color"},"link":null},{"localized_name":"crop_position","name":"crop_position","type":"COMBO","widget":{"name":"crop_position"},"link":null},{"localized_name":"divisible_by","name":"divisible_by","type":"INT","widget":{"name":"divisible_by"},"link":null},{"localized_name":"device","name":"device","shape":7,"type":"COMBO","widget":{"name":"device"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[1281,2103,2106,2111,2139,2140]},{"localized_name":"width","name":"width","type":"INT","links":null},{"localized_name":"height","name":"height","type":"INT","links":null},{"localized_name":"mask","name":"mask","type":"MASK","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"4dfb85dcc52e4315c33170d97bb987baa46d128b","Node name for S&R":"ImageResizeKJv2","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[0,0,"area","crop","0, 0, 0","center",32,"cpu"],"color":"#232","bgcolor":"#353"},{"id":532,"type":"ResizeImageMaskNode","pos":[-627.975883287854,4495.453864775762],"size":[270,106],"flags":{},"order":56,"mode":0,"inputs":[{"localized_name":"input","name":"input","type":"IMAGE,MASK","link":1281},{"localized_name":"resize_type","name":"resize_type","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"resize_type"},"link":null},{"localized_name":"resize_type.multiplier","name":"resize_type.multiplier","type":"FLOAT","widget":{"name":"resize_type.multiplier"},"link":null},{"localized_name":"scale_method","name":"scale_method","type":"COMBO","widget":{"name":"scale_method"},"link":null}],"outputs":[{"localized_name":"resized","name":"resized","type":"IMAGE","links":[1282]}],"properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"ResizeImageMaskNode","ue_properties":{"widget_ue_connectable":{"resize_type":true,"resize_type.multiplier":true,"scale_method":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["scale by multiplier",0.5,"area"]},{"id":533,"type":"GetImageSize","pos":[-625.7563863146385,4658.162256891575],"size":[210,136],"flags":{"collapsed":true},"order":61,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":1282}],"outputs":[{"localized_name":"width","name":"width","type":"INT","links":[1289,1566]},{"localized_name":"height","name":"height","type":"INT","links":[1290,1567]},{"localized_name":"batch_size","name":"batch_size","type":"INT","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"GetImageSize","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#232","bgcolor":"#353"},{"id":621,"type":"SetNode","pos":[-982.6473011877084,4117.627386747715],"size":[210,58],"flags":{"collapsed":true},"order":47,"mode":0,"inputs":[{"name":"INT","type":"INT","link":1426}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_seed","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"seed","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["seed"],"color":"#1b4669","bgcolor":"#29699c"},{"id":539,"type":"SetNode","pos":[-1350.4852103468006,3976.131691041333],"size":[210,50],"flags":{"collapsed":true},"order":53,"mode":0,"inputs":[{"name":"INT","type":"INT","link":1291}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_fps_int","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"fps_int","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["fps_int"],"color":"#1b4669","bgcolor":"#29699c"},{"id":540,"type":"CM_FloatToInt","pos":[-1353.334898876357,3979.7395006294246],"size":[269.31637824001507,60.0508913584531],"flags":{"collapsed":true},"order":44,"mode":0,"inputs":[{"localized_name":"a","name":"a","type":"FLOAT","widget":{"name":"a"},"link":1292}],"outputs":[{"localized_name":"INT","name":"INT","type":"INT","links":[1291]}],"properties":{"cnr_id":"ComfyMath","ver":"c01177221c31b8e5fbc062778fc8254aeb541638","Node name for S&R":"CM_FloatToInt","ue_properties":{"widget_ue_connectable":{"a":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0]},{"id":541,"type":"SetNode","pos":[-1349.0391107707837,3972.7422031848832],"size":[210,58],"flags":{"collapsed":true},"order":45,"mode":0,"inputs":[{"name":"FLOAT","type":"FLOAT","link":1293}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_fps_float","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"fps_float","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":597,"type":"VHS_VideoCombine","pos":[3989.618639202434,3886.685868295037],"size":[1682.5545248780247,2108.35184987085],"flags":{"collapsed":false},"order":90,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":2026},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":2179},{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":1481},{"localized_name":"loop_count","name":"loop_count","type":"INT","widget":{"name":"loop_count"},"link":null},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"pingpong","name":"pingpong","type":"BOOLEAN","widget":{"name":"pingpong"},"link":null},{"localized_name":"save_output","name":"save_output","type":"BOOLEAN","widget":{"name":"save_output"},"link":null},{"name":"pix_fmt","type":["yuv420p","yuv420p10le"],"widget":{"name":"pix_fmt"},"link":null},{"name":"crf","type":"INT","widget":{"name":"crf"},"link":null},{"name":"save_metadata","type":"BOOLEAN","widget":{"name":"save_metadata"},"link":null},{"name":"trim_to_audio","type":"BOOLEAN","widget":{"name":"trim_to_audio"},"link":null}],"outputs":[{"localized_name":"Filenames","name":"Filenames","type":"VHS_FILENAMES","links":null}],"properties":{"cnr_id":"comfyui-videohelpersuite","ver":"1.7.9","Node name for S&R":"VHS_VideoCombine","ue_properties":{"widget_ue_connectable":{"frame_rate":true,"loop_count":true,"filename_prefix":true,"format":true,"pingpong":true,"save_output":true,"pix_fmt":true,"crf":true,"save_metadata":true,"trim_to_audio":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":{"frame_rate":24,"loop_count":0,"filename_prefix":"Eros/I2V","format":"video/h264-mp4","pix_fmt":"yuv420p","crf":10,"save_metadata":true,"trim_to_audio":false,"pingpong":false,"save_output":true,"videopreview":{"hidden":false,"paused":false,"params":{"filename":"I2V_00013-audio.mp4","subfolder":"Eros","type":"output","format":"video/h264-mp4","frame_rate":25,"workflow":"I2V_00013.png","fullpath":"H:\\New_AI_hub\\matrix2\\Data\\Packages\\ComfyUI2\\output\\Eros\\I2V_00013-audio.mp4"}}},"color":"#222","bgcolor":"#000"},{"id":740,"type":"VAEDecode","pos":[4193.035737094888,3532.5951130775347],"size":[254.98011089086322,46],"flags":{},"order":86,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":2172},{"localized_name":"vae","name":"vae","type":"VAE","link":1942}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[2025]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"VAEDecode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.5.2"}},"widgets_values":[]},{"id":549,"type":"VHS_VideoCombine","pos":[1381.9752964545542,4109.515776559602],"size":[805.7524473679769,1428.4096151309682],"flags":{"collapsed":false},"order":80,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":1303},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":1714},{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":1308},{"localized_name":"loop_count","name":"loop_count","type":"INT","widget":{"name":"loop_count"},"link":null},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"pingpong","name":"pingpong","type":"BOOLEAN","widget":{"name":"pingpong"},"link":null},{"localized_name":"save_output","name":"save_output","type":"BOOLEAN","widget":{"name":"save_output"},"link":null},{"name":"pix_fmt","type":["yuv420p","yuv420p10le"],"widget":{"name":"pix_fmt"},"link":null},{"name":"crf","type":"INT","widget":{"name":"crf"},"link":null},{"name":"save_metadata","type":"BOOLEAN","widget":{"name":"save_metadata"},"link":null},{"name":"trim_to_audio","type":"BOOLEAN","widget":{"name":"trim_to_audio"},"link":null}],"outputs":[{"localized_name":"Filenames","name":"Filenames","type":"VHS_FILENAMES","links":null}],"properties":{"cnr_id":"comfyui-videohelpersuite","ver":"1.7.9","Node name for S&R":"VHS_VideoCombine","ue_properties":{"widget_ue_connectable":{"frame_rate":true,"loop_count":true,"filename_prefix":true,"format":true,"pingpong":true,"save_output":true,"pix_fmt":true,"crf":true,"save_metadata":true,"trim_to_audio":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":{"frame_rate":1,"loop_count":0,"filename_prefix":"10E_firstpass","format":"video/h264-mp4","pix_fmt":"yuv420p","crf":8,"save_metadata":true,"trim_to_audio":false,"pingpong":false,"save_output":false,"videopreview":{"hidden":false,"paused":false,"params":{"filename":"10E_firstpass_00018-audio.mp4","subfolder":"","type":"temp","format":"video/h264-mp4","frame_rate":25,"workflow":"10E_firstpass_00018.png","fullpath":"H:\\New_AI_hub\\matrix2\\Data\\Packages\\ComfyUI2\\temp\\10E_firstpass_00018-audio.mp4"}}},"color":"#222","bgcolor":"#000"},{"id":744,"type":"LTXVLatentUpsamplerTiled","pos":[3144.4283806312183,3251.3863845050637],"size":[315.929296875,194],"flags":{},"order":76,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":2091},{"localized_name":"upscale_model","name":"upscale_model","type":"LATENT_UPSCALE_MODEL","link":1956},{"localized_name":"vae","name":"vae","type":"VAE","link":1957},{"localized_name":"tile_size","name":"tile_size","shape":7,"type":"INT","widget":{"name":"tile_size"},"link":null},{"localized_name":"overlap","name":"overlap","shape":7,"type":"INT","widget":{"name":"overlap"},"link":null},{"localized_name":"max_size_for_no_tile","name":"max_size_for_no_tile","shape":7,"type":"INT","widget":{"name":"max_size_for_no_tile"},"link":null},{"localized_name":"rotate_for_landscape","name":"rotate_for_landscape","shape":7,"type":"BOOLEAN","widget":{"name":"rotate_for_landscape"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[2122]}],"properties":{"aux_id":"TenStrip/10S-Comfy-nodes","ver":"47852934037fbf96ed6b6572dded3c1da37bf36d","Node name for S&R":"LTXVLatentUpsamplerTiled"},"widgets_values":[11,6,22,false,false],"color":"#223","bgcolor":"#335"},{"id":555,"type":"GetNode","pos":[1132.0161214120128,4089.5184760428774],"size":[210,58],"flags":{"collapsed":true},"order":27,"mode":0,"inputs":[],"outputs":[{"name":"FLOAT","type":"FLOAT","links":[1308]}],"title":"Get_fps_float","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":734,"type":"LTXVCropGuides","pos":[1115.4955143216073,3969.1883218852795],"size":[152.7134765625,66],"flags":{},"order":72,"mode":0,"inputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2146},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2147},{"localized_name":"latent","name":"latent","type":"LATENT","link":1900}],"outputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","links":[2137]},{"localized_name":"negative","name":"negative","type":"CONDITIONING","links":[2138]},{"localized_name":"latent","name":"latent","type":"LATENT","links":[1905,2091]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LTXVCropGuides","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[]},{"id":779,"type":"JWStringToFloat","pos":[-131.4793794648068,2931.8241735329775],"size":[270,58],"flags":{},"order":28,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2180]}],"title":"First Pass Video Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["0.88"],"color":"#222","bgcolor":"#000"},{"id":772,"type":"LTXVImgToVideoInplaceKJ","pos":[176.24606920810962,4085.756030025743],"size":[294.98046875,146],"flags":{},"order":65,"mode":0,"inputs":[{"localized_name":"vae","name":"vae","type":"VAE","link":2118},{"localized_name":"latent","name":"latent","type":"LATENT","link":2117},{"localized_name":"num_images","name":"num_images","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"num_images"},"link":null},{"localized_name":"num_images.strength_1","name":"num_images.strength_1","type":"FLOAT","widget":{"name":"num_images.strength_1"},"link":2219},{"label":"image_1","localized_name":"num_images.image_1","name":"num_images.image_1","shape":7,"type":"IMAGE","link":2140},{"localized_name":"num_images.index_1","name":"num_images.index_1","shape":7,"type":"INT","widget":{"name":"num_images.index_1"},"link":null}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[2120,2123]}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTXVImgToVideoInplaceKJ","ue_properties":{"widget_ue_connectable":{"num_images":true,"num_images.strength_1":true,"num_images.index_1":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["1",0.7499999999999999,0],"color":"#233","bgcolor":"#355"},{"id":767,"type":"LTXVAddGuide","pos":[177.01976308882737,3864.7547451448645],"size":[270,162],"flags":{},"order":68,"mode":0,"inputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2148},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2149},{"localized_name":"vae","name":"vae","type":"VAE","link":2075},{"localized_name":"latent","name":"latent","type":"LATENT","link":2123},{"localized_name":"image","name":"image","type":"IMAGE","link":2139},{"localized_name":"frame_idx","name":"frame_idx","type":"INT","widget":{"name":"frame_idx"},"link":null},{"localized_name":"strength","name":"strength","type":"FLOAT","widget":{"name":"strength"},"link":2218}],"outputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","links":[2144,2146]},{"localized_name":"negative","name":"negative","type":"CONDITIONING","links":[2145,2147]},{"localized_name":"latent","name":"latent","type":"LATENT","links":[]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LTXVAddGuide","ue_properties":{"widget_ue_connectable":{"frame_idx":true,"strength":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0,0.7499999999999999],"color":"#233","bgcolor":"#355"},{"id":542,"type":"PrimitiveFloat","pos":[-1425.8431702584294,3927.973226571728],"size":[263.24,58],"flags":{},"order":29,"mode":0,"inputs":[{"localized_name":"value","name":"value","type":"FLOAT","widget":{"name":"value"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[1292,1293]}],"title":"Base Frame Rate (24 Default)","properties":{"cnr_id":"comfy-core","ver":"0.8.0","Node name for S&R":"PrimitiveFloat","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[25],"color":"#222","bgcolor":"#000"},{"id":788,"type":"LTXTiledSampler","pos":[3875.1742002272445,3304.8312860366577],"size":[279.60546875,282],"flags":{},"order":84,"mode":0,"inputs":[{"localized_name":"noise","name":"noise","type":"NOISE","link":2206},{"localized_name":"guider","name":"guider","type":"GUIDER","link":2207},{"localized_name":"sampler","name":"sampler","type":"SAMPLER","link":2208},{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":2209},{"localized_name":"latent_image","name":"latent_image","type":"LATENT","link":2210},{"localized_name":"tile_axis","name":"tile_axis","shape":7,"type":"COMBO","widget":{"name":"tile_axis"},"link":null},{"localized_name":"n_tiles","name":"n_tiles","shape":7,"type":"INT","widget":{"name":"n_tiles"},"link":null},{"localized_name":"tile_overlap","name":"tile_overlap","shape":7,"type":"INT","widget":{"name":"tile_overlap"},"link":null},{"localized_name":"max_size_for_no_tile","name":"max_size_for_no_tile","shape":7,"type":"INT","widget":{"name":"max_size_for_no_tile"},"link":null},{"localized_name":"audio_pass","name":"audio_pass","shape":7,"type":"COMBO","widget":{"name":"audio_pass"},"link":null},{"localized_name":"audio_carrier_tile","name":"audio_carrier_tile","shape":7,"type":"COMBO","widget":{"name":"audio_carrier_tile"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"output","name":"output","type":"LATENT","links":null},{"localized_name":"denoised_output","name":"denoised_output","type":"LATENT","links":[2211]}],"properties":{"aux_id":"TenStrip/10S-Comfy-nodes","ver":"47852934037fbf96ed6b6572dded3c1da37bf36d","Node name for S&R":"LTXTiledSampler"},"widgets_values":["auto",2,4,38,"tile_carrying","first",false],"color":"#223","bgcolor":"#335"},{"id":557,"type":"Power Lora Loader (rgthree)","pos":[-796.6583203205506,3371.2238770967456],"size":[428.9642467678492,322.75598817305604],"flags":{},"order":41,"mode":0,"inputs":[{"dir":3,"name":"model","type":"MODEL","link":1547},{"dir":3,"name":"clip","type":"CLIP","link":1418}],"outputs":[{"dir":4,"name":"MODEL","shape":3,"type":"MODEL","links":[1852,1859]},{"dir":4,"name":"CLIP","shape":3,"type":"CLIP","links":null}],"properties":{"cnr_id":"rgthree-comfy","ver":"1.0.2512112053","Show Strengths":"Single Strength","Match":"","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[{},{"type":"PowerLoraLoaderHeaderWidget"},{},""],"color":"#222","bgcolor":"#000"},{"id":797,"type":"mxSlider","pos":[-1291.7110040938053,4233.142465026225],"size":[329.41939288068033,30],"flags":{"collapsed":false},"order":31,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"FLOAT","links":[2218,2219]}],"title":"Conditioning Strength","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":0.8,"min":0.4,"max":1,"step":0.05,"decimals":2,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[0,0.8,1],"color":"#222","bgcolor":"#000"},{"id":538,"type":"SetNode","pos":[-1040.509645734978,3950.804887433771],"size":[210,50],"flags":{"collapsed":true},"order":55,"mode":0,"inputs":[{"name":"INT","type":"INT","link":2221}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_length","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["length_0"],"color":"#1b4669","bgcolor":"#29699c"},{"id":798,"type":"ComfyMathExpression","pos":[-1053.6205311505564,3975.8234801246163],"size":[400,200],"flags":{"collapsed":true},"order":48,"mode":0,"inputs":[{"label":"a","localized_name":"values.a","name":"values.a","type":"FLOAT,INT","link":2220},{"label":"b","localized_name":"values.b","name":"values.b","shape":7,"type":"FLOAT,INT","link":null},{"localized_name":"expression","name":"expression","type":"STRING","widget":{"name":"expression"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":null},{"localized_name":"INT","name":"INT","type":"INT","links":[2221]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"ComfyMathExpression"},"widgets_values":["a + 1"]},{"id":791,"type":"mxSlider","pos":[-1432.9324238695663,4053.537250181444],"size":[280.438609128239,30],"flags":{"collapsed":false},"order":33,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"INT","links":[2215]}],"title":"Video Width","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":1024,"min":480,"max":2400,"step":32,"decimals":0,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[1024,1024,0],"color":"#222","bgcolor":"#000"},{"id":773,"type":"LoadImage","pos":[-1429.4456334516165,4356.670409703162],"size":[646.0397431410997,996.8304004758456],"flags":{},"order":34,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[2150]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LoadImage","ue_properties":{"widget_ue_connectable":{"image":true,"upload":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["pasted/image (97).png","image"],"color":"#222","bgcolor":"#000"},{"id":731,"type":"LTXLatentAnchorAware","pos":[-155.21856858203535,3792.1885062578012],"size":[290.362890625,450],"flags":{},"order":66,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2088},{"localized_name":"reference_image","name":"reference_image","shape":7,"type":"IMAGE","link":2111},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":1886},{"localized_name":"energy_latent","name":"energy_latent","shape":7,"type":"LATENT","link":2141},{"localized_name":"sigmas","name":"sigmas","shape":7,"type":"SIGMAS","link":1888},{"localized_name":"strength","name":"strength","shape":7,"type":"FLOAT","widget":{"name":"strength"},"link":null},{"localized_name":"cache_at_step","name":"cache_at_step","shape":7,"type":"INT","widget":{"name":"cache_at_step"},"link":null},{"localized_name":"similarity_threshold","name":"similarity_threshold","shape":7,"type":"FLOAT","widget":{"name":"similarity_threshold"},"link":null},{"localized_name":"decay_with_distance","name":"decay_with_distance","shape":7,"type":"FLOAT","widget":{"name":"decay_with_distance"},"link":null},{"localized_name":"energy_threshold","name":"energy_threshold","shape":7,"type":"FLOAT","widget":{"name":"energy_threshold"},"link":null},{"localized_name":"bypass","name":"bypass","shape":7,"type":"BOOLEAN","widget":{"name":"bypass"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null},{"localized_name":"advanced_mode","name":"advanced_mode","shape":7,"type":"BOOLEAN","widget":{"name":"advanced_mode"},"link":null},{"localized_name":"cache_mode","name":"cache_mode","shape":7,"type":"COMBO","widget":{"name":"cache_mode"},"link":null},{"localized_name":"forwards_per_step","name":"forwards_per_step","shape":7,"type":"INT","widget":{"name":"forwards_per_step"},"link":null},{"localized_name":"cache_warmup","name":"cache_warmup","shape":7,"type":"INT","widget":{"name":"cache_warmup"},"link":null},{"localized_name":"anchor_frame","name":"anchor_frame","shape":7,"type":"INT","widget":{"name":"anchor_frame"},"link":null},{"localized_name":"depth_curve","name":"depth_curve","shape":7,"type":"COMBO","widget":{"name":"depth_curve"},"link":null},{"localized_name":"block_index_filter","name":"block_index_filter","shape":7,"type":"STRING","widget":{"name":"block_index_filter"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[1889]}],"properties":{"aux_id":"TenStrip/10S-Comfy-nodes","ver":"47852934037fbf96ed6b6572dded3c1da37bf36d","Node name for S&R":"LTXLatentAnchorAware"},"widgets_values":[0.11000000000000001,6,0.5,0.15000000000000002,0.3,false,false,true,"schedule",1,432,0,"flat","10-30"],"color":"#223","bgcolor":"#335"},{"id":795,"type":"MarkdownNote","pos":[-2021.1608937808342,3197.3792734479475],"size":[521.3413669628835,1161.0100683552218],"flags":{},"order":30,"mode":0,"inputs":[],"outputs":[],"properties":{"ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["# 10S Workflows\n\nThese are advanced quality-focused conditioned workflows for 10Eros/Sulphur/Base LTX 2.3. They make some compromises on speed for quality. \nSteal these or study 10S Node intergation if you are a workflow maker or want to customize and release them as your own. \n**They are NOT FAST.** I find with LTX, fast ≠ good.\n\n---\n\n## 📦 Downloads\n\n| Resource | Link |\n|----------|------|\n| 🎬 Workflows | [LTX2.3-10Eros Workflows](https://huggingface.co/TenStrip/LTX2.3-10Eros_Workflows) |\n| 🧠 Models (10Eros/Sulphur) | [huggingface.co/TenStrip/LTX2.3-10Eros](https://huggingface.co/TenStrip/LTX2.3-10Eros) |\n| 🔧 10S Comfy Nodes | [TenStrip/10S-Comfy-nodes](https://github.com/TenStrip/10S-Comfy-nodes) |\n| ⚗️ Distilled LoRAs | [LTX2.3 Distilled LoRA 1.1 Experiments](https://huggingface.co/TenStrip/LTX2.3_Distilled_Lora_1.1_Experiments/tree/main) |\n---\n\n## ⚙️ How It Works\n\nThese workflows rely on **[10S-Comfy-nodes]**\n\nGuided prompt conditioning from the guide is leveraged against the i2v conditioned latent with Latent anchor. Increasing strength on Latent anchor as well as **depth_curve** gives you a more manageable control over less/more movement and gives you 3 points of strength to control the i2v freedom. \n\nLatent Anchor node needs paper-level study in what it does, but for most calculations turn debug on and follow this formula to set it up correctly: \n*target_call_count = cache_at_step × forwards_per_step*\n\nIt is currently configured for 13 steps at step 6 cache with 1 CFG, although forwards per step must be increased to 3 to account for CFG,STG,audioCFG-but isn't impactful on just first 2 CFG steps.\n\nUpscale distortion is removed and detail is increased with tiled sampler, this is a default-level improvement to quality. At it's current setting it will only trigger when distortion inducing resolution is reached for vertical ratios, if working in landscape increase max_tile_size accordingly to bypass it. Audio-sync issues can arise if the **tile_carrying** mode does not target the tile with talking heads or sound-causing actors. Generally you will select first for vertical resolutions and middle for huge landscape passes.\nBecause of how tiled sampler works you will want to avoid hand motion and movement near the middle of the composition.\n\nDistillation is handled at full break down. Lower pass video distillation will increase motion. Upscale pass video distillation will add detail, but also add artifacts and unaesthetic blur artifacts. Audio weight must be kept high for audio quality.\n\nIn some cases first pass audio may be more desirable, that can be toggled to be hard encoded instead of passed and resampled fully.\n\nThe first pass output is a full **motion** preview, if it contains mutation, distortion, bad motion, bad audio, etc... The upscale pass will not fix it. Only allow it to run to upscale when first pass motion looks good. **Visual quality and style will be restored** at upscale.\n\n\n---\n\n## ✍️ Prompt Enhancement\n\nThese use STG Guider and have CFG for initial steps which is where most broad motion and audio form. **Your negative prompt is impactful, so use it.**\n\nFor best positive prompt results, use this foreword in **Grok** or an **uncensored LLM**:\n\n---\n\n*Generate a video scene script with a description based on the attached image for an LLM that has a tokenizer that uses interleaved attention to support long-context understanding that is fed into a multimodal video model. Strict specification, follow up to the word: No timestamps. No unnecessary embellishment. Output only plain English text and make it a copy box.*\n\n*First, describe the image initial scene in concise natural language; subject(s), subject(s) appearance, subject(s) composition and pose, background, and context.*\n\n*Next, formulate a naturally evolving scenario that would take place describing every moving body part, composition change, and manipulation from the uploaded initial frame that would be reflected in the video models post-latent evolution output. If the image is explicit or sexual in nature, use full anatomical terminology and spice it up slightly with visually representable erotic themes.*\n\n*Center the prompt around this basic idea: **[ concept ]***\n\n*Interweave this dialogue or sound concept into the scene with descriptions of voice tone followed by the lines delivered in quotations, in a temporal sequence between or during motions. Dialogue should be concise and non-rambling as it will take away from video quality: **[ dialogue ]***\n\n*Inside that prompt describe only notable audio and audio cues, both normal and explicit; background noise as well as foley and natural sounds. In a temporal sequence paired with coinciding motions. In the case of absent dialogue or soundscapes and only if background music is fitting; describe a fitting genre and melodic tone with matching mood.*\n\n*Output only text following above instruction. Follow-up suggestions should be on the topic of expanding or changing motion or dialogue from the output text.*\n"],"color":"#432","bgcolor":"#653"},{"id":536,"type":"CLIPTextEncode","pos":[-641.7618745880603,3760.1840791968625],"size":[394.7031737883649,255.53167788014707],"flags":{},"order":39,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":1419},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","links":[2142]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"CLIPTextEncode","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[""],"color":"#232","bgcolor":"#353"},{"id":796,"type":"mxSlider","pos":[-1140.5320997743233,3944.5899083019813],"size":[329.41939288068033,30],"flags":{"collapsed":false},"order":35,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"INT","links":[2220]}],"title":"Length (Frame Count X/24 = Seconds)","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":264,"min":48,"max":960,"step":24,"decimals":0,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[264,264,0],"color":"#222","bgcolor":"#000"},{"id":524,"type":"Seed (rgthree)","pos":[-1122.836985011658,4031.7241836848057],"size":[288.55371900826447,130],"flags":{},"order":32,"mode":0,"inputs":[],"outputs":[{"dir":4,"name":"SEED","shape":3,"type":"INT","links":[1273,1426]}],"properties":{"cnr_id":"rgthree-comfy","ver":"1.0.2511270846","randomMax":1125899906842624,"randomMin":0,"ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[-1,"","",""],"color":"#222","bgcolor":"#000"},{"id":792,"type":"mxSlider","pos":[-1434.642031574845,4150.46662371214],"size":[280.438609128239,30],"flags":{"collapsed":false},"order":36,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"INT","links":[2216]}],"title":"Video Height","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":1344,"min":480,"max":2400,"step":32,"decimals":0,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[1344,1344,0],"color":"#222","bgcolor":"#000"}],"links":[[583,274,0,275,0,"LATENT"],[586,279,0,274,1,"VAE"],[1272,525,0,526,0,"NOISE"],[1273,524,0,525,0,"INT"],[1274,526,0,510,0,"NOISE"],[1276,520,0,510,2,"SAMPLER"],[1281,531,0,532,0,"IMAGE"],[1282,532,0,533,0,"IMAGE"],[1289,533,0,534,0,"INT"],[1290,533,1,534,1,"INT"],[1291,540,0,539,0,"INT"],[1292,542,0,540,0,"FLOAT"],[1293,542,0,541,0,"FLOAT"],[1295,543,0,535,1,"INT"],[1296,545,0,535,2,"INT"],[1297,546,0,534,2,"INT"],[1303,552,0,549,0,"IMAGE"],[1305,554,0,550,1,"VAE"],[1306,550,0,551,0,"AUDIO"],[1307,553,0,552,1,"VAE"],[1308,555,0,549,4,"FLOAT"],[1318,558,0,535,0,"VAE"],[1320,560,0,523,2,"FLOAT"],[1382,592,0,593,1,"VAE"],[1387,593,0,598,0,"AUDIO"],[1418,616,0,557,1,"CLIP"],[1419,616,0,536,0,"CLIP"],[1420,616,0,537,0,"CLIP"],[1423,617,0,207,0,"VAE"],[1426,524,0,621,0,"INT"],[1457,630,0,275,1,"MASK"],[1481,635,0,597,4,"FLOAT"],[1531,548,0,510,4,"LATENT"],[1546,510,1,556,0,"LATENT"],[1547,646,0,557,0,"MODEL"],[1548,646,2,214,0,"VAE"],[1566,533,0,630,1,"INT"],[1567,533,1,630,2,"INT"],[1568,527,0,652,0,"SIGMAS"],[1569,652,0,510,3,"SIGMAS"],[1602,550,0,274,0,"AUDIO"],[1608,556,1,550,0,"LATENT"],[1609,535,0,548,1,"LATENT"],[1611,596,1,593,0,"LATENT"],[1714,551,0,549,1,"AUDIO"],[1852,557,0,719,0,"MODEL"],[1853,719,0,718,0,"MODEL"],[1859,557,0,722,0,"MODEL"],[1860,722,0,723,0,"MODEL"],[1886,559,0,731,2,"VAE"],[1888,527,0,731,4,"SIGMAS"],[1889,731,0,653,0,"MODEL"],[1900,556,0,734,2,"LATENT"],[1905,734,2,552,0,"LATENT"],[1942,594,0,740,1,"VAE"],[1956,743,0,744,1,"LATENT_UPSCALE_MODEL"],[1957,589,0,744,2,"VAE"],[2017,718,0,753,0,"MODEL"],[2018,753,0,583,0,"MODEL"],[2025,740,0,755,0,"IMAGE"],[2026,755,0,597,0,"IMAGE"],[2043,653,0,510,1,"GUIDER"],[2075,559,0,767,2,"VAE"],[2088,723,0,731,0,"MODEL"],[2091,734,2,744,0,"LATENT"],[2103,531,0,769,0,"IMAGE"],[2106,531,0,331,0,"IMAGE"],[2107,714,0,770,4,"IMAGE"],[2109,589,0,770,0,"VAE"],[2110,770,0,591,0,"LATENT"],[2111,531,0,731,1,"IMAGE"],[2117,534,0,772,1,"LATENT"],[2118,559,0,772,0,"VAE"],[2120,772,0,548,0,"LATENT"],[2122,744,0,770,1,"LATENT"],[2123,772,0,767,3,"LATENT"],[2137,734,0,583,1,"CONDITIONING"],[2138,734,1,583,2,"CONDITIONING"],[2139,531,0,767,4,"IMAGE"],[2140,531,0,772,4,"IMAGE"],[2141,534,0,731,3,"LATENT"],[2142,536,0,523,0,"CONDITIONING"],[2143,537,0,523,1,"CONDITIONING"],[2144,767,0,653,1,"CONDITIONING"],[2145,767,1,653,2,"CONDITIONING"],[2146,767,0,734,0,"CONDITIONING"],[2147,767,1,734,1,"CONDITIONING"],[2148,523,0,767,0,"CONDITIONING"],[2149,523,1,767,1,"CONDITIONING"],[2150,773,0,531,0,"IMAGE"],[2172,596,0,740,0,"LATENT"],[2179,593,0,597,1,"AUDIO"],[2180,779,0,722,4,"FLOAT"],[2181,780,0,723,4,"FLOAT"],[2182,781,0,718,4,"FLOAT"],[2183,782,0,719,4,"FLOAT"],[2184,783,0,722,3,"COMBO"],[2185,783,0,723,3,"COMBO"],[2186,783,0,719,3,"COMBO"],[2187,783,0,718,3,"COMBO"],[2206,584,0,788,0,"NOISE"],[2207,583,0,788,1,"GUIDER"],[2208,585,0,788,2,"SAMPLER"],[2209,582,0,788,3,"SIGMAS"],[2210,591,0,788,4,"LATENT"],[2211,788,1,596,0,"LATENT"],[2212,274,0,789,0,"LATENT"],[2213,556,1,789,1,"LATENT"],[2214,789,0,591,1,"LATENT"],[2215,791,0,531,2,"INT"],[2216,792,0,531,3,"INT"],[2218,797,0,767,6,"FLOAT"],[2219,797,0,772,3,"FLOAT"],[2220,796,0,798,0,"INT"],[2221,798,1,538,0,"INT"]],"groups":[{"id":27,"title":"Main","bounding":[-1494.9909480082615,3161.937903314709,3804.2302523868716,2304.222673775987],"color":"#3f789e","flags":{}},{"id":23,"title":"Force First-pass Audio","bounding":[1228.7248700088794,2801.335742850906,537.4514362675345,323.87650696724893],"color":"#3f789e","flags":{}},{"id":29,"title":"Final Upscale Pass","bounding":[2374.9914506547025,3167.3647232660733,2118.8972740162553,1955.3713894900138],"color":"#3f789e","flags":{}},{"id":30,"title":"Distilation Break-up (Thanks Lightricks For This Wonderful Distilled Lora)","bounding":[-167.0518161898845,2841.0578682916353,1330.1373185243046,278.6386435094696],"color":"#3f789e","flags":{}},{"id":32,"title":"Post Options, Add interpolation here before RTX","bounding":[4693.505674922303,3438.368751942385,687.0438904788525,342.37525726057083],"color":"#3f789e","flags":{}},{"id":33,"title":"Parameters","bounding":[-1444.103685997252,3815.005816195167,642.0089842779156,461.4995675177156],"color":"#3f789e","flags":{}}],"config":{},"extra":{"ds":{"scale":0.9090909090909107,"offset":[1758.6267071402829,-3662.4662806693705]},"frontendVersion":"1.43.17","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":false,"VHS_latentpreviewrate":0,"VHS_MetadataImage":true,"VHS_KeepIntermediate":true,"ue_links":[],"links_added_by_ue":[]},"version":0.4} \ No newline at end of file diff --git a/custom_nodes/10S-Comfy-nodes/workflows/10Eros_10SNodes_LikenessGuideHelper_I2V_v3.2.json b/custom_nodes/10S-Comfy-nodes/workflows/10Eros_10SNodes_LikenessGuideHelper_I2V_v3.2.json new file mode 100644 index 0000000000000000000000000000000000000000..dfd63e87e1a185805bfeca79627b4ad83d0004eb --- /dev/null +++ b/custom_nodes/10S-Comfy-nodes/workflows/10Eros_10SNodes_LikenessGuideHelper_I2V_v3.2.json @@ -0,0 +1 @@ +{"id":"ee686d30-b3a6-4a60-991e-b8bba9c03fd5","revision":0,"last_node_id":834,"last_link_id":2372,"nodes":[{"id":560,"type":"GetNode","pos":[231.68427124392713,3365.8233044983185],"size":[210,58],"flags":{"collapsed":true},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"FLOAT","type":"FLOAT","links":[1320]}],"title":"Get_fps_float","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":546,"type":"GetNode","pos":[-124.41455081715429,4554.729113523804],"size":[210,50],"flags":{"collapsed":true},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","links":[1297]}],"title":"Get_length_0","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["length_0"],"color":"#1b4669","bgcolor":"#29699c"},{"id":630,"type":"SolidMask","pos":[1238.7248700088776,2874.9357428509056],"size":[270,106],"flags":{},"order":64,"mode":0,"inputs":[{"localized_name":"value","name":"value","type":"FLOAT","widget":{"name":"value"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":1566},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":1567}],"outputs":[{"localized_name":"MASK","name":"MASK","type":"MASK","links":[1457]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"SolidMask","ue_properties":{"widget_ue_connectable":{"value":true,"width":true,"height":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0,1024,1024]},{"id":279,"type":"GetNode","pos":[1254.5088394005425,3043.669001545407],"size":[210,58],"flags":{"collapsed":true},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[586]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":275,"type":"SetLatentNoiseMask","pos":[1491.3415734193702,3061.244993543124],"size":[236.13729359019888,46],"flags":{},"order":81,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":583},{"localized_name":"mask","name":"mask","type":"MASK","link":1457}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"SetLatentNoiseMask","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[]},{"id":525,"type":"RandomNoise","pos":[-113.16650352315176,3653.016847379711],"size":[270,82],"flags":{"collapsed":true},"order":46,"mode":0,"inputs":[{"localized_name":"noise_seed","name":"noise_seed","type":"INT","widget":{"name":"noise_seed"},"link":1273}],"outputs":[{"localized_name":"NOISE","name":"NOISE","type":"NOISE","links":[1272]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"RandomNoise","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[728885192875474,"randomize"]},{"id":526,"type":"SetNode","pos":[-119.6034969855223,3592.656729380967],"size":[210,58],"flags":{"collapsed":true},"order":54,"mode":0,"inputs":[{"name":"NOISE","type":"NOISE","link":1272}],"outputs":[{"name":"NOISE","type":"NOISE","links":[1274]}],"title":"Set_noise_seed","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"noise_seed","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["noise_seed"],"color":"#2e2e2e","bgcolor":"#242121"},{"id":214,"type":"SetNode","pos":[-1117.3074592682947,3477.466425119632],"size":[210,60],"flags":{"collapsed":true},"order":39,"mode":0,"inputs":[{"name":"VAE","type":"VAE","link":1548}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_vid_vae","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"vid_vae","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":207,"type":"SetNode","pos":[-1131.973635706669,3343.177927274481],"size":[236.85766571079353,58],"flags":{"collapsed":true},"order":40,"mode":0,"inputs":[{"name":"VAE","type":"VAE","link":1423}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_aud_vae","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"aud_vae","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":652,"type":"Sigmas Easing","pos":[739.795612293746,4201.997117923222],"size":[447.6995362297288,154],"flags":{},"order":38,"mode":0,"inputs":[{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":1568},{"localized_name":"easing_type","name":"easing_type","type":"COMBO","widget":{"name":"easing_type"},"link":null},{"localized_name":"easing_mode","name":"easing_mode","type":"COMBO","widget":{"name":"easing_mode"},"link":null},{"localized_name":"normalize_input","name":"normalize_input","type":"BOOLEAN","widget":{"name":"normalize_input"},"link":null},{"localized_name":"normalize_output","name":"normalize_output","type":"BOOLEAN","widget":{"name":"normalize_output"},"link":null},{"localized_name":"strength","name":"strength","type":"FLOAT","widget":{"name":"strength"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[1569]}],"title":"Sigmas Easing (cubic, sine, quad) don't mess with this tbh","properties":{"cnr_id":"RES4LYF","ver":"0dc91c00c4c3fb38e7874fcd7a2a327765e8882c","Node name for S&R":"Sigmas Easing","ue_properties":{"widget_ue_connectable":{"easing_type":true,"easing_mode":true,"normalize_input":true,"normalize_output":true,"strength":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["cubic","in_out",true,true,0.7],"color":"#233","bgcolor":"#355"},{"id":545,"type":"GetNode","pos":[-115.01505791637044,4796.53233460559],"size":[210,58],"flags":{"collapsed":true},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","links":[1296]}],"title":"Get_fps_int","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_int"],"color":"#1b4669","bgcolor":"#29699c"},{"id":543,"type":"GetNode","pos":[-112.3520085965529,4844.393106338195],"size":[210,50],"flags":{"collapsed":true},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"INT","type":"INT","links":[1295]}],"title":"Get_length_0","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["length_0"],"color":"#1b4669","bgcolor":"#29699c"},{"id":558,"type":"GetNode","pos":[-111.18858564259628,4901.712773792447],"size":[210,50],"flags":{"collapsed":true},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1318]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":635,"type":"GetNode","pos":[3675.6779403363385,3944.3050892612646],"size":[210,58],"flags":{"collapsed":true},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"FLOAT","type":"FLOAT","links":[1481]}],"title":"Get_fps_float","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":714,"type":"GetNode","pos":[2813.360751805555,3608.779911647415],"size":[210,60],"flags":{"collapsed":false},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[2107]}],"title":"Get_resized_image","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["resized_image"],"color":"#2a363b","bgcolor":"#3f5159"},{"id":589,"type":"GetNode","pos":[2811.237769683004,3354.9545348795436],"size":[210,58],"flags":{"collapsed":false},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1957,2109]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":553,"type":"GetNode","pos":[766.7614713430766,3904.3966824075724],"size":[210,50],"flags":{"collapsed":true},"order":9,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1307]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":554,"type":"GetNode","pos":[767.012010065247,4049.870641019706],"size":[210,34],"flags":{"collapsed":true},"order":10,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1305]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":769,"type":"PreviewImage","pos":[-596.8441361687978,4886.676255650047],"size":[294.6088617022692,468.8946724679554],"flags":{},"order":56,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":2103}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"PreviewImage","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#222","bgcolor":"#000"},{"id":527,"type":"ManualSigmas","pos":[194.9532242386173,3763.4626399016347],"size":[279.41750632544984,59.3366153498655],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"sigmas","name":"sigmas","type":"STRING","widget":{"name":"sigmas"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[1568,1888]}],"title":"Sigmas First Pass","properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"ManualSigmas","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.1","input_ue_unconnectable":{}}},"widgets_values":["1.0,0.99375,0.9875,0.98125,0.9550,0.8925,0.8120,0.7150,0.6030,0.4824,0.3618,0.2412,0.1206,0.0"],"color":"#323","bgcolor":"#535"},{"id":520,"type":"KSamplerSelect","pos":[205.96839723877576,3632.5948112199703],"size":[261.46963267468345,59.13737920167023],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null}],"outputs":[{"localized_name":"SAMPLER","name":"SAMPLER","type":"SAMPLER","links":[1276]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"KSamplerSelect","ue_properties":{"widget_ue_connectable":{"sampler_name":true},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":["euler_ancestral"],"color":"#233","bgcolor":"#355"},{"id":585,"type":"KSamplerSelect","pos":[2813.2370757937415,3478.790375760474],"size":[270,58],"flags":{"collapsed":false},"order":13,"mode":0,"inputs":[{"localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null}],"outputs":[{"localized_name":"SAMPLER","name":"SAMPLER","type":"SAMPLER","links":[2267]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.75","Node name for S&R":"KSamplerSelect","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{"sampler_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["euler_ancestral_cfg_pp"],"color":"#322","bgcolor":"#533"},{"id":584,"type":"GetNode","pos":[3658.1883945443683,3604.7452337560367],"size":[210,50],"flags":{"collapsed":true},"order":14,"mode":0,"inputs":[],"outputs":[{"name":"NOISE","type":"NOISE","links":[2265]}],"title":"Get_noise_seed","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["noise_seed"],"color":"#2e2e2e","bgcolor":"#242121"},{"id":783,"type":"easy loraNames","pos":[-923.1107370119908,3255.5052938715753],"size":[573.4679999999979,64.3888243713377],"flags":{},"order":15,"mode":0,"inputs":[{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":null}],"outputs":[{"localized_name":"lora_name","name":"lora_name","type":"*","links":[2184,2185,2186,2187]}],"title":"Distilled Lora","properties":{"cnr_id":"comfyui-easy-use","ver":"1.3.5","Node name for S&R":"easy loraNames","ue_properties":{"widget_ue_connectable":{"lora_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors"],"color":"#222","bgcolor":"#000"},{"id":780,"type":"JWStringToFloat","pos":[188.4736608306006,2933.8859567773025],"size":[270,58],"flags":{},"order":16,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2181]}],"title":"First Pass Audio Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["0.9"],"color":"#222","bgcolor":"#000"},{"id":722,"type":"LTX2LoraLoaderAdvanced","pos":[-117.5080221723439,3066.4864860588896],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":52,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1859},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2184},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2180},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[1860]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,1,0,0,0,1]},{"id":723,"type":"LTX2LoraLoaderAdvanced","pos":[201.42706107783408,3065.5746013380663],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":60,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1860},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2185},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2181},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[2088]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,0,1,1,1,0]},{"id":719,"type":"LTX2LoraLoaderAdvanced","pos":[533.6158560203178,3067.449408361733],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":51,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1852},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2186},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2183},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[1853]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,1,0,0,0,1]},{"id":718,"type":"LTX2LoraLoaderAdvanced","pos":[864.9976953461493,3064.0054448447786],"size":[293.6646484375,242],"flags":{"collapsed":true},"order":59,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":1853},{"localized_name":"opt_lora_path","name":"opt_lora_path","shape":7,"type":"STRING","link":null},{"localized_name":"blocks","name":"blocks","shape":7,"type":"SELECTEDDITBLOCKS","link":null},{"localized_name":"lora_name","name":"lora_name","type":"COMBO","widget":{"name":"lora_name"},"link":2187},{"localized_name":"strength_model","name":"strength_model","type":"FLOAT","widget":{"name":"strength_model"},"link":2182},{"localized_name":"video","name":"video","type":"FLOAT","widget":{"name":"video"},"link":null},{"localized_name":"video_to_audio","name":"video_to_audio","type":"FLOAT","widget":{"name":"video_to_audio"},"link":null},{"localized_name":"audio","name":"audio","type":"FLOAT","widget":{"name":"audio"},"link":null},{"localized_name":"audio_to_video","name":"audio_to_video","type":"FLOAT","widget":{"name":"audio_to_video"},"link":null},{"localized_name":"other","name":"other","type":"FLOAT","widget":{"name":"other"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[2366]},{"localized_name":"rank","name":"rank","type":"STRING","links":null},{"localized_name":"loaded_keys_info","name":"loaded_keys_info","type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTX2LoraLoaderAdvanced","ue_properties":{"widget_ue_connectable":{"lora_name":true,"strength_model":true,"video":true,"video_to_audio":true,"audio":true,"audio_to_video":true,"other":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx23\\ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors",1,0,1,1,1,0]},{"id":743,"type":"LatentUpscaleModelLoader","pos":[2742.801612317037,3240.087891289923],"size":[353.96248779296855,58],"flags":{},"order":17,"mode":0,"inputs":[{"localized_name":"model_name","name":"model_name","type":"COMBO","widget":{"name":"model_name"},"link":null}],"outputs":[{"localized_name":"LATENT_UPSCALE_MODEL","name":"LATENT_UPSCALE_MODEL","type":"LATENT_UPSCALE_MODEL","links":[1956]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LatentUpscaleModelLoader","ue_properties":{"widget_ue_connectable":{"model_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["ltx-2.3-spatial-upscaler-x2-1.1.safetensors"],"color":"#323","bgcolor":"#535"},{"id":274,"type":"LTXVAudioVAEEncode","pos":[1547.6286442388596,2949.1333592283954],"size":[201.4126953125,46],"flags":{},"order":79,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":1602},{"label":"Audio VAE","localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":586}],"outputs":[{"localized_name":"Audio Latent","name":"Audio Latent","type":"LATENT","links":[583,2212]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"LTXVAudioVAEEncode","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[]},{"id":789,"type":"TwoWaySwitch","pos":[2100.2422289843566,3308.3526900407155],"size":[474.2561723693839,115.75201477050769],"flags":{"collapsed":false},"order":82,"mode":0,"inputs":[{"localized_name":"input_1","name":"input_1","shape":7,"type":"*","link":2212},{"localized_name":"input_2","name":"input_2","shape":7,"type":"*","link":2213},{"localized_name":"selection_setting","name":"selection_setting","type":"INT","widget":{"name":"selection_setting"},"link":null}],"outputs":[{"localized_name":"output","name":"output","type":"*","links":[2214]}],"title":"Resample Audio (2) or Force First Pass Audio (1)","properties":{"cnr_id":"controlaltai-nodes","ver":"721492b66c9cede8ae23ae10615462ad80cfd061","Node name for S&R":"TwoWaySwitch","ue_properties":{"widget_ue_connectable":{"selection_setting":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[2],"color":"#222","bgcolor":"#000"},{"id":646,"type":"CheckpointLoaderSimple","pos":[-1458.3873184326671,3368.418161889962],"size":[270,98],"flags":{},"order":18,"mode":0,"inputs":[{"localized_name":"ckpt_name","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null}],"outputs":[{"localized_name":"MODEL","name":"MODEL","type":"MODEL","links":[1547]},{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":null},{"localized_name":"VAE","name":"VAE","type":"VAE","links":[1548]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"CheckpointLoaderSimple","ue_properties":{"widget_ue_connectable":{"ckpt_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["10Eros_v1-fp8mixed_learned.safetensors"],"color":"#432","bgcolor":"#653"},{"id":617,"type":"LTXVAudioVAELoader","pos":[-1456.4873518000977,3244.6519920330684],"size":[270,58],"flags":{},"order":19,"mode":0,"inputs":[{"localized_name":"ckpt_name","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null}],"outputs":[{"localized_name":"Audio VAE","name":"Audio VAE","type":"VAE","links":[1423]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LTXVAudioVAELoader","ue_properties":{"widget_ue_connectable":{"ckpt_name":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["10Eros_v1-fp8mixed_learned.safetensors"],"color":"#332922","bgcolor":"#593930"},{"id":533,"type":"GetImageSize","pos":[-601.4523042529495,4667.276345609934],"size":[210,136],"flags":{"collapsed":true},"order":61,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":1282}],"outputs":[{"localized_name":"width","name":"width","type":"INT","links":[1289,1566]},{"localized_name":"height","name":"height","type":"INT","links":[1290,1567]},{"localized_name":"batch_size","name":"batch_size","type":"INT","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"GetImageSize","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#232","bgcolor":"#353"},{"id":621,"type":"SetNode","pos":[-958.3432191260194,4126.741475466073],"size":[210,58],"flags":{"collapsed":true},"order":47,"mode":0,"inputs":[{"name":"INT","type":"INT","link":1426}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_seed","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"seed","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["seed"],"color":"#1b4669","bgcolor":"#29699c"},{"id":539,"type":"SetNode","pos":[-1326.181128285111,3985.2457797596926],"size":[210,50],"flags":{"collapsed":true},"order":50,"mode":0,"inputs":[{"name":"INT","type":"INT","link":1291}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_fps_int","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"fps_int","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["fps_int"],"color":"#1b4669","bgcolor":"#29699c"},{"id":540,"type":"CM_FloatToInt","pos":[-1329.0308168146676,3988.8535893477842],"size":[269.31637824001507,60.0508913584531],"flags":{"collapsed":true},"order":41,"mode":0,"inputs":[{"localized_name":"a","name":"a","type":"FLOAT","widget":{"name":"a"},"link":1292}],"outputs":[{"localized_name":"INT","name":"INT","type":"INT","links":[1291]}],"properties":{"cnr_id":"ComfyMath","ver":"c01177221c31b8e5fbc062778fc8254aeb541638","Node name for S&R":"CM_FloatToInt","ue_properties":{"widget_ue_connectable":{"a":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0]},{"id":541,"type":"SetNode","pos":[-1324.7350287090942,3981.856291903243],"size":[210,58],"flags":{"collapsed":true},"order":42,"mode":0,"inputs":[{"name":"FLOAT","type":"FLOAT","link":1293}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_fps_float","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"fps_float","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":740,"type":"VAEDecode","pos":[4193.035737094888,3532.5951130775347],"size":[254.98011089086322,46],"flags":{},"order":86,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":2307},{"localized_name":"vae","name":"vae","type":"VAE","link":1942}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[2025]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"VAEDecode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.5.2"}},"widgets_values":[]},{"id":555,"type":"GetNode","pos":[1156.3202034737023,4098.632564761237],"size":[210,58],"flags":{"collapsed":true},"order":20,"mode":0,"inputs":[],"outputs":[{"name":"FLOAT","type":"FLOAT","links":[1308]}],"title":"Get_fps_float","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["fps_float"],"color":"#232","bgcolor":"#353"},{"id":538,"type":"SetNode","pos":[-1016.2055636732887,3959.9189761521307],"size":[210,50],"flags":{"collapsed":true},"order":58,"mode":0,"inputs":[{"name":"INT","type":"INT","link":2221}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_length","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["length_0"],"color":"#1b4669","bgcolor":"#29699c"},{"id":798,"type":"ComfyMathExpression","pos":[-1029.3164490888669,3984.937568842976],"size":[400,200],"flags":{"collapsed":true},"order":49,"mode":0,"inputs":[{"label":"a","localized_name":"values.a","name":"values.a","type":"FLOAT,INT,BOOLEAN","link":2220},{"label":"b","localized_name":"values.b","name":"values.b","shape":7,"type":"FLOAT,INT,BOOLEAN","link":null},{"localized_name":"expression","name":"expression","type":"STRING","widget":{"name":"expression"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":null},{"localized_name":"INT","name":"INT","type":"INT","links":[2221]},{"localized_name":"BOOL","name":"BOOL","type":"BOOLEAN","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"ComfyMathExpression"},"widgets_values":["a + 1"]},{"id":542,"type":"PrimitiveFloat","pos":[-1401.5390881967398,3937.0873152900876],"size":[263.24,58],"flags":{},"order":21,"mode":0,"inputs":[{"localized_name":"value","name":"value","type":"FLOAT","widget":{"name":"value"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[1292,1293]}],"title":"Base Frame Rate (24 Default)","properties":{"cnr_id":"comfy-core","ver":"0.8.0","Node name for S&R":"PrimitiveFloat","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[25],"color":"#222","bgcolor":"#000"},{"id":559,"type":"GetNode","pos":[-434.9434296471235,4269.844380713661],"size":[210,58],"flags":{"collapsed":true},"order":22,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1886,2118,2288]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":331,"type":"SetNode","pos":[-358.9001328669263,4800.997387078811],"size":[210,60],"flags":{"collapsed":true},"order":57,"mode":0,"inputs":[{"name":"IMAGE","type":"IMAGE","link":2106}],"outputs":[{"name":"*","type":"*","links":null}],"title":"Set_resized_image","properties":{"Node name for S&R":"SetNode","aux_id":"kijai/ComfyUI-KJNodes","previousName":"resized_image","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["resized_image"],"color":"#2a363b","bgcolor":"#3f5159"},{"id":531,"type":"ImageResizeKJv2","pos":[-598.7096254495578,4413.219332711192],"size":[288.23754794013735,336.00000000000045],"flags":{"collapsed":true},"order":48,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":2371},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":2215},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":2216},{"localized_name":"upscale_method","name":"upscale_method","type":"COMBO","widget":{"name":"upscale_method"},"link":null},{"localized_name":"keep_proportion","name":"keep_proportion","type":"COMBO","widget":{"name":"keep_proportion"},"link":null},{"localized_name":"pad_color","name":"pad_color","type":"STRING","widget":{"name":"pad_color"},"link":null},{"localized_name":"crop_position","name":"crop_position","type":"COMBO","widget":{"name":"crop_position"},"link":null},{"localized_name":"divisible_by","name":"divisible_by","type":"INT","widget":{"name":"divisible_by"},"link":null},{"localized_name":"device","name":"device","shape":7,"type":"COMBO","widget":{"name":"device"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[1281,2103,2106,2111,2140,2290]},{"localized_name":"width","name":"width","type":"INT","links":null},{"localized_name":"height","name":"height","type":"INT","links":null},{"localized_name":"mask","name":"mask","type":"MASK","links":null}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"4dfb85dcc52e4315c33170d97bb987baa46d128b","Node name for S&R":"ImageResizeKJv2","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[0,0,"area","crop","0, 0, 0","center",32,"cpu"],"color":"#232","bgcolor":"#353"},{"id":535,"type":"LTXVEmptyLatentAudio","pos":[-123.40904679592069,4605.252969807392],"size":[270,120],"flags":{"collapsed":false},"order":37,"mode":0,"inputs":[{"localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":1318},{"localized_name":"frames_number","name":"frames_number","type":"INT","widget":{"name":"frames_number"},"link":1295},{"localized_name":"frame_rate","name":"frame_rate","type":"INT","widget":{"name":"frame_rate"},"link":1296},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"localized_name":"Latent","name":"Latent","type":"LATENT","links":[2236]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"LTXVEmptyLatentAudio","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{"frames_number":true,"frame_rate":true},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[1,24,1]},{"id":534,"type":"EmptyLTXVLatentVideo","pos":[-129.22085515127537,4328.931784050993],"size":[270,146.66666666666669],"flags":{},"order":63,"mode":0,"inputs":[{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":1289},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":1290},{"localized_name":"length","name":"length","type":"INT","widget":{"name":"length"},"link":1297},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[2289]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.60","Node name for S&R":"EmptyLTXVLatentVideo","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{"width":true,"height":true,"length":true},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[704,64,97,1]},{"id":596,"type":"LTXVSeparateAVLatent","pos":[4187.725994744947,3295.4824999445127],"size":[200.43247236282696,46],"flags":{},"order":85,"mode":0,"inputs":[{"localized_name":"av_latent","name":"av_latent","type":"LATENT","link":2270}],"outputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","links":[2307]},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","links":[1611]}],"properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"LTXVSeparateAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":744,"type":"LTXVLatentUpsamplerTiled","pos":[3144.4283806312183,3251.3863845050637],"size":[315.929296875,194],"flags":{},"order":74,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":2259},{"localized_name":"upscale_model","name":"upscale_model","type":"LATENT_UPSCALE_MODEL","link":1956},{"localized_name":"vae","name":"vae","type":"VAE","link":1957},{"localized_name":"tile_size","name":"tile_size","shape":7,"type":"INT","widget":{"name":"tile_size"},"link":null},{"localized_name":"overlap","name":"overlap","shape":7,"type":"INT","widget":{"name":"overlap"},"link":null},{"localized_name":"max_size_for_no_tile","name":"max_size_for_no_tile","shape":7,"type":"INT","widget":{"name":"max_size_for_no_tile"},"link":null},{"localized_name":"rotate_for_landscape","name":"rotate_for_landscape","shape":7,"type":"BOOLEAN","widget":{"name":"rotate_for_landscape"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[2271]}],"properties":{"cnr_id":"10s-comfy-nodes","ver":"47852934037fbf96ed6b6572dded3c1da37bf36d","Node name for S&R":"LTXVLatentUpsamplerTiled","aux_id":"TenStrip/10S-Comfy-nodes"},"widgets_values":[11,6,22,false,false],"color":"#223","bgcolor":"#335"},{"id":770,"type":"LTXVImgToVideoInplaceKJ","pos":[3148.455220746667,3526.853708945063],"size":[294.98046875,146],"flags":{},"order":77,"mode":0,"inputs":[{"localized_name":"vae","name":"vae","type":"VAE","link":2109},{"localized_name":"latent","name":"latent","type":"LATENT","link":2271},{"localized_name":"num_images","name":"num_images","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"num_images"},"link":null},{"localized_name":"num_images.strength_1","name":"num_images.strength_1","type":"FLOAT","widget":{"name":"num_images.strength_1"},"link":null},{"label":"image_1","localized_name":"num_images.image_1","name":"num_images.image_1","shape":7,"type":"IMAGE","link":2107},{"localized_name":"num_images.index_1","name":"num_images.index_1","shape":7,"type":"INT","widget":{"name":"num_images.index_1"},"link":null}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[2272]}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTXVImgToVideoInplaceKJ","ue_properties":{"widget_ue_connectable":{"num_images":true,"num_images.strength_1":true,"num_images.index_1":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["1",1,0],"color":"#233","bgcolor":"#355"},{"id":593,"type":"LTXVAudioVAEDecode","pos":[4190.703613972216,3658.441870311331],"size":[240,46],"flags":{},"order":87,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":1611},{"label":"Audio VAE","localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":1382}],"outputs":[{"localized_name":"Audio","name":"Audio","type":"AUDIO","links":[1387,2179]}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"LTXVAudioVAEDecode","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":[]},{"id":594,"type":"GetNode","pos":[4017.8094594305244,3699.089876138927],"size":[210,50],"flags":{"collapsed":true},"order":23,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1942]}],"title":"Get_vid_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["vid_vae"],"color":"#322","bgcolor":"#533"},{"id":592,"type":"GetNode","pos":[4014.0155853784486,3817.035660623502],"size":[210,50],"flags":{"collapsed":true},"order":24,"mode":0,"inputs":[],"outputs":[{"name":"VAE","type":"VAE","links":[1382]}],"title":"Get_aud_vae","properties":{"Node name for S&R":"GetNode","aux_id":"kijai/ComfyUI-KJNodes","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["aud_vae"],"color":"#322","bgcolor":"#533"},{"id":548,"type":"LTXVConcatAVLatent","pos":[740.377257340901,4093.7142408333916],"size":[270,46],"flags":{},"order":68,"mode":0,"inputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","link":2304},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","link":2236}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[1531]}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"LTXVConcatAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":523,"type":"LTXVConditioning","pos":[204.255143869948,3497.5645449918684],"size":[259.4071611759441,79.50641394128222],"flags":{"collapsed":false},"order":53,"mode":0,"inputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2142},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2143},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":1320}],"outputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","links":[2286]},{"localized_name":"negative","name":"negative","type":"CONDITIONING","links":[2287]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"LTXVConditioning","ue_properties":{"widget_ue_connectable":{"frame_rate":true},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":[24],"color":"#233","bgcolor":"#355"},{"id":781,"type":"JWStringToFloat","pos":[838.2050706633592,2933.050555207784],"size":[270,58],"flags":{"collapsed":false},"order":25,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2182]}],"title":"Upscale Audio Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["1"],"color":"#222","bgcolor":"#000"},{"id":532,"type":"ResizeImageMaskNode","pos":[-624.754876970794,4501.054122430277],"size":[270,106],"flags":{},"order":55,"mode":0,"inputs":[{"localized_name":"input","name":"input","type":"IMAGE,MASK","link":1281},{"localized_name":"resize_type","name":"resize_type","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"resize_type"},"link":null},{"localized_name":"resize_type.multiplier","name":"resize_type.multiplier","type":"FLOAT","widget":{"name":"resize_type.multiplier"},"link":null},{"localized_name":"scale_method","name":"scale_method","type":"COMBO","widget":{"name":"scale_method"},"link":null}],"outputs":[{"localized_name":"resized","name":"resized","type":"IMAGE","links":[1282,2310]}],"properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"ResizeImageMaskNode","ue_properties":{"widget_ue_connectable":{"resize_type":true,"resize_type.multiplier":true,"scale_method":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["scale by multiplier",0.5,"area"]},{"id":616,"type":"LTXAVTextEncoderLoader","pos":[-1460.686865946572,3546.187363427289],"size":[318.444921875,106],"flags":{},"order":26,"mode":0,"inputs":[{"localized_name":"text_encoder","name":"text_encoder","type":"COMBO","widget":{"name":"text_encoder"},"link":null},{"localized_name":"ckpt_name","name":"ckpt_name","type":"COMBO","widget":{"name":"ckpt_name"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null}],"outputs":[{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":[1418,1419,1420]}],"properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"LTXAVTextEncoderLoader","ue_properties":{"widget_ue_connectable":{"text_encoder":true,"ckpt_name":true,"device":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["gemma_3_12B_it_fp8_scaled.safetensors","10Eros_v1-fp8mixed_learned.safetensors","default"],"color":"#222","bgcolor":"#000"},{"id":772,"type":"LTXVImgToVideoInplaceKJ","pos":[211.96099232357093,4364.559159371701],"size":[294.98046875,146],"flags":{},"order":67,"mode":0,"inputs":[{"localized_name":"vae","name":"vae","type":"VAE","link":2118},{"localized_name":"latent","name":"latent","type":"LATENT","link":2303},{"localized_name":"num_images","name":"num_images","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"num_images"},"link":null},{"localized_name":"num_images.strength_1","name":"num_images.strength_1","type":"FLOAT","widget":{"name":"num_images.strength_1"},"link":2219},{"label":"image_1","localized_name":"num_images.image_1","name":"num_images.image_1","shape":7,"type":"IMAGE","link":2140},{"localized_name":"num_images.index_1","name":"num_images.index_1","shape":7,"type":"INT","widget":{"name":"num_images.index_1"},"link":null}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[2304,2305]}],"properties":{"cnr_id":"comfyui-kjnodes","ver":"a9ab19d2a7933e03f72bc00103863f06b9483300","Node name for S&R":"LTXVImgToVideoInplaceKJ","ue_properties":{"widget_ue_connectable":{"num_images":true,"num_images.strength_1":true,"num_images.index_1":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["1",0.7499999999999999,0],"color":"#233","bgcolor":"#355"},{"id":653,"type":"STGGuiderAdvanced","pos":[608.8690776753175,4490.35707107023],"size":[684.9663455547113,358],"flags":{"collapsed":false},"order":71,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2372},{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2292},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2294},{"localized_name":"preset","name":"preset","shape":7,"type":"STG_ADVANCED_PRESET","link":null},{"localized_name":"skip_steps_sigma_threshold","name":"skip_steps_sigma_threshold","type":"FLOAT","widget":{"name":"skip_steps_sigma_threshold"},"link":null},{"localized_name":"cfg_star_rescale","name":"cfg_star_rescale","type":"BOOLEAN","widget":{"name":"cfg_star_rescale"},"link":null},{"localized_name":"sigmas","name":"sigmas","type":"STRING","widget":{"name":"sigmas"},"link":null},{"localized_name":"cfg_values","name":"cfg_values","type":"STRING","widget":{"name":"cfg_values"},"link":null},{"localized_name":"stg_scale_values","name":"stg_scale_values","type":"STRING","widget":{"name":"stg_scale_values"},"link":null},{"localized_name":"stg_rescale_values","name":"stg_rescale_values","type":"STRING","widget":{"name":"stg_rescale_values"},"link":null},{"localized_name":"stg_layers_indices","name":"stg_layers_indices","type":"STRING","widget":{"name":"stg_layers_indices"},"link":null},{"localized_name":"apply_apg","name":"apply_apg","shape":7,"type":"BOOLEAN","widget":{"name":"apply_apg"},"link":null},{"localized_name":"apg_cfg_scale","name":"apg_cfg_scale","shape":7,"type":"FLOAT","widget":{"name":"apg_cfg_scale"},"link":null},{"localized_name":"eta","name":"eta","shape":7,"type":"FLOAT","widget":{"name":"eta"},"link":null},{"localized_name":"norm_threshold","name":"norm_threshold","shape":7,"type":"FLOAT","widget":{"name":"norm_threshold"},"link":null}],"outputs":[{"localized_name":"GUIDER","name":"GUIDER","type":"GUIDER","links":[2043]}],"properties":{"cnr_id":"ComfyUI-LTXVideo","ver":"82bd963cdeb66d023bed8c99324a307020907ef8","Node name for S&R":"STGGuiderAdvanced","ue_properties":{"widget_ue_connectable":{"skip_steps_sigma_threshold":true,"cfg_star_rescale":true,"sigmas":true,"cfg_values":true,"stg_scale_values":true,"stg_rescale_values":true,"stg_layers_indices":true,"apply_apg":true,"apg_cfg_scale":true,"eta":true,"norm_threshold":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[1,false,"1.0,0.99375,0.9875,0.98125,0.9550,0.8925,0.8120,0.7150,0.6030,0.4824,0.3618,0.2412,0.1206,0.0","2,1.5,1,1,1,1,1,1,1,1,1,1,1","2,1.5,1,1,1,1,1,1,1,1,1,1,1","1,1,1,1,1,1,1,1,1,1,1,1,1","[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999],[9999]",false,1,1,1],"color":"#233","bgcolor":"#355"},{"id":808,"type":"LTXFaceDetector","pos":[-174.7346487441511,4962.321269802819],"size":[270,106],"flags":{},"order":62,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":2310},{"localized_name":"padding","name":"padding","shape":7,"type":"FLOAT","widget":{"name":"padding"},"link":null},{"localized_name":"fallback_bbox","name":"fallback_bbox","shape":7,"type":"STRING","widget":{"name":"fallback_bbox"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"face_bbox","name":"face_bbox","type":"STRING","links":[2311,2353,2361]}],"properties":{"cnr_id":"10s-comfy-nodes","ver":"2fcc4ac6018ae6fc6c2ff64c2e2a58f9e50ba30a","Node name for S&R":"LTXFaceDetector"},"widgets_values":[0.15000000000000002,"",false],"color":"#233","bgcolor":"#355"},{"id":731,"type":"LTXLatentAnchorAware","pos":[-123.45457058252265,3811.609849146388],"size":[290.362890625,450],"flags":{},"order":69,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2088},{"localized_name":"reference_image","name":"reference_image","shape":7,"type":"IMAGE","link":2111},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":1886},{"localized_name":"energy_latent","name":"energy_latent","shape":7,"type":"LATENT","link":2305},{"localized_name":"sigmas","name":"sigmas","shape":7,"type":"SIGMAS","link":1888},{"localized_name":"strength","name":"strength","shape":7,"type":"FLOAT","widget":{"name":"strength"},"link":null},{"localized_name":"cache_at_step","name":"cache_at_step","shape":7,"type":"INT","widget":{"name":"cache_at_step"},"link":null},{"localized_name":"similarity_threshold","name":"similarity_threshold","shape":7,"type":"FLOAT","widget":{"name":"similarity_threshold"},"link":null},{"localized_name":"decay_with_distance","name":"decay_with_distance","shape":7,"type":"FLOAT","widget":{"name":"decay_with_distance"},"link":null},{"localized_name":"energy_threshold","name":"energy_threshold","shape":7,"type":"FLOAT","widget":{"name":"energy_threshold"},"link":null},{"localized_name":"bypass","name":"bypass","shape":7,"type":"BOOLEAN","widget":{"name":"bypass"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null},{"localized_name":"advanced_mode","name":"advanced_mode","shape":7,"type":"BOOLEAN","widget":{"name":"advanced_mode"},"link":null},{"localized_name":"cache_mode","name":"cache_mode","shape":7,"type":"COMBO","widget":{"name":"cache_mode"},"link":null},{"localized_name":"forwards_per_step","name":"forwards_per_step","shape":7,"type":"INT","widget":{"name":"forwards_per_step"},"link":null},{"localized_name":"cache_warmup","name":"cache_warmup","shape":7,"type":"INT","widget":{"name":"cache_warmup"},"link":null},{"localized_name":"anchor_frame","name":"anchor_frame","shape":7,"type":"INT","widget":{"name":"anchor_frame"},"link":null},{"localized_name":"depth_curve","name":"depth_curve","shape":7,"type":"COMBO","widget":{"name":"depth_curve"},"link":null},{"localized_name":"block_index_filter","name":"block_index_filter","shape":7,"type":"STRING","widget":{"name":"block_index_filter"},"link":null}],"outputs":[{"localized_name":"model","name":"model","type":"MODEL","links":[2299,2359]}],"properties":{"cnr_id":"10s-comfy-nodes","ver":"47852934037fbf96ed6b6572dded3c1da37bf36d","Node name for S&R":"LTXLatentAnchorAware","aux_id":"TenStrip/10S-Comfy-nodes"},"widgets_values":[0.11,5,0.5,0.15000000000000002,0.3,false,false,true,"schedule",2,50,0,"flat",""],"color":"#223","bgcolor":"#335"},{"id":583,"type":"CFGGuider","pos":[3542.6365852942336,3663.4475380786234],"size":[270,98],"flags":{},"order":66,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2366},{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2291},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2293},{"localized_name":"cfg","name":"cfg","type":"FLOAT","widget":{"name":"cfg"},"link":null}],"outputs":[{"localized_name":"GUIDER","name":"GUIDER","type":"GUIDER","links":[2266]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.71","Node name for S&R":"CFGGuider","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[1]},{"id":597,"type":"VHS_VideoCombine","pos":[3870.928705603576,3877.0993595930177],"size":[1186.2974423514481,1784.4567369074332],"flags":{"collapsed":false},"order":90,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":2026},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":2179},{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":1481},{"localized_name":"loop_count","name":"loop_count","type":"INT","widget":{"name":"loop_count"},"link":null},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"pingpong","name":"pingpong","type":"BOOLEAN","widget":{"name":"pingpong"},"link":null},{"localized_name":"save_output","name":"save_output","type":"BOOLEAN","widget":{"name":"save_output"},"link":null},{"name":"pix_fmt","type":["yuv420p","yuv420p10le"],"widget":{"name":"pix_fmt"},"link":null},{"name":"crf","type":"INT","widget":{"name":"crf"},"link":null},{"name":"save_metadata","type":"BOOLEAN","widget":{"name":"save_metadata"},"link":null},{"name":"trim_to_audio","type":"BOOLEAN","widget":{"name":"trim_to_audio"},"link":null}],"outputs":[{"localized_name":"Filenames","name":"Filenames","type":"VHS_FILENAMES","links":null}],"properties":{"cnr_id":"comfyui-videohelpersuite","ver":"1.7.9","Node name for S&R":"VHS_VideoCombine","ue_properties":{"widget_ue_connectable":{"frame_rate":true,"loop_count":true,"filename_prefix":true,"format":true,"pingpong":true,"save_output":true,"pix_fmt":true,"crf":true,"save_metadata":true,"trim_to_audio":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":{"frame_rate":24,"loop_count":0,"filename_prefix":"Eros/I2V","format":"video/h264-mp4","pix_fmt":"yuv420p","crf":6,"save_metadata":true,"trim_to_audio":false,"pingpong":false,"save_output":true,"videopreview":{"hidden":false,"paused":false,"params":{"filename":"I2V_00088-audio.mp4","subfolder":"Eros","type":"output","format":"video/h264-mp4","frame_rate":25,"workflow":"I2V_00088.png","fullpath":"H:\\New_AI_hub\\matrix2\\Data\\Packages\\ComfyUI2\\output\\Eros\\I2V_00088-audio.mp4"}}},"color":"#222","bgcolor":"#000"},{"id":549,"type":"VHS_VideoCombine","pos":[1464.4416673997355,4119.990205878587],"size":[812.2700816993929,1326.6865715110148],"flags":{"collapsed":false},"order":80,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":1303},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":1714},{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":1308},{"localized_name":"loop_count","name":"loop_count","type":"INT","widget":{"name":"loop_count"},"link":null},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"pingpong","name":"pingpong","type":"BOOLEAN","widget":{"name":"pingpong"},"link":null},{"localized_name":"save_output","name":"save_output","type":"BOOLEAN","widget":{"name":"save_output"},"link":null},{"name":"pix_fmt","type":["yuv420p","yuv420p10le"],"widget":{"name":"pix_fmt"},"link":null},{"name":"crf","type":"INT","widget":{"name":"crf"},"link":null},{"name":"save_metadata","type":"BOOLEAN","widget":{"name":"save_metadata"},"link":null},{"name":"trim_to_audio","type":"BOOLEAN","widget":{"name":"trim_to_audio"},"link":null}],"outputs":[{"localized_name":"Filenames","name":"Filenames","type":"VHS_FILENAMES","links":null}],"properties":{"cnr_id":"comfyui-videohelpersuite","ver":"1.7.9","Node name for S&R":"VHS_VideoCombine","ue_properties":{"widget_ue_connectable":{"frame_rate":true,"loop_count":true,"filename_prefix":true,"format":true,"pingpong":true,"save_output":true,"pix_fmt":true,"crf":true,"save_metadata":true,"trim_to_audio":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":{"frame_rate":1,"loop_count":0,"filename_prefix":"10E_firstpass","format":"video/h264-mp4","pix_fmt":"yuv420p","crf":8,"save_metadata":true,"trim_to_audio":false,"pingpong":false,"save_output":false,"videopreview":{"hidden":false,"paused":false,"params":{"filename":"10E_firstpass_00002-audio.mp4","subfolder":"","type":"temp","format":"video/h264-mp4","frame_rate":25,"workflow":"10E_firstpass_00002.png","fullpath":"H:\\New_AI_hub\\matrix2\\Data\\Packages\\ComfyUI2\\temp\\10E_firstpass_00002-audio.mp4"}}},"color":"#222","bgcolor":"#000"},{"id":551,"type":"AudioAdjustVolume","pos":[1127.1639915606715,3863.7019091264565],"size":[270,58],"flags":{},"order":78,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":1306},{"localized_name":"volume","name":"volume","type":"INT","widget":{"name":"volume"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[1714]}],"properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"AudioAdjustVolume","ue_properties":{"widget_ue_connectable":{"volume":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0],"color":"#222","bgcolor":"#000"},{"id":598,"type":"AudioAdjustVolume","pos":[3543.226714786005,3823.3257733842042],"size":[270,58],"flags":{},"order":89,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":1387},{"localized_name":"volume","name":"volume","type":"INT","widget":{"name":"volume"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[]}],"properties":{"cnr_id":"comfy-core","ver":"0.14.1","Node name for S&R":"AudioAdjustVolume","ue_properties":{"widget_ue_connectable":{"volume":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[0],"color":"#222","bgcolor":"#000"},{"id":755,"type":"RTXVideoSuperResolution","pos":[5042.307315358794,3574.972988633893],"size":[285.81015625,106],"flags":{},"order":88,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":2025},{"localized_name":"resize_type","name":"resize_type","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"resize_type"},"link":null},{"localized_name":"resize_type.scale","name":"resize_type.scale","type":"FLOAT","widget":{"name":"resize_type.scale"},"link":null},{"localized_name":"quality","name":"quality","type":"COMBO","widget":{"name":"quality"},"link":null}],"outputs":[{"localized_name":"upscaled_images","name":"upscaled_images","type":"IMAGE","links":[2026]}],"properties":{"cnr_id":"comfyui_nvidia_rtx_nodes","ver":"892515e3eb9a4920a131a502a047e47adca9eb0d","Node name for S&R":"RTXVideoSuperResolution"},"widgets_values":["scale by multiplier",1.3,"ULTRA"]},{"id":782,"type":"JWStringToFloat","pos":[516.2577673613606,2935.444411519735],"size":[270,58],"flags":{},"order":27,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2183]}],"title":"Upscale Video Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["0.75"],"color":"#222","bgcolor":"#000"},{"id":827,"type":"LTXLikenessAnchor","pos":[205.94489789464737,4566.408542279273],"size":[302.99609375,414],"flags":{},"order":70,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":2359},{"localized_name":"reference_info","name":"reference_info","shape":7,"type":"REFERENCE_INFO","link":2360},{"localized_name":"strength","name":"strength","type":"FLOAT","widget":{"name":"strength"},"link":null},{"localized_name":"reference_source","name":"reference_source","shape":7,"type":"COMBO","widget":{"name":"reference_source"},"link":null},{"localized_name":"frame_0_bbox","name":"frame_0_bbox","shape":7,"type":"STRING","widget":{"name":"frame_0_bbox"},"link":2361},{"localized_name":"similarity_threshold","name":"similarity_threshold","shape":7,"type":"FLOAT","widget":{"name":"similarity_threshold"},"link":null},{"localized_name":"decay_with_distance","name":"decay_with_distance","shape":7,"type":"FLOAT","widget":{"name":"decay_with_distance"},"link":null},{"localized_name":"bypass","name":"bypass","shape":7,"type":"BOOLEAN","widget":{"name":"bypass"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null},{"localized_name":"advanced_mode","name":"advanced_mode","shape":7,"type":"BOOLEAN","widget":{"name":"advanced_mode"},"link":null},{"localized_name":"depth_curve","name":"depth_curve","shape":7,"type":"COMBO","widget":{"name":"depth_curve"},"link":null},{"localized_name":"block_index_filter","name":"block_index_filter","shape":7,"type":"STRING","widget":{"name":"block_index_filter"},"link":null},{"localized_name":"similarity_sharpness","name":"similarity_sharpness","shape":7,"type":"FLOAT","widget":{"name":"similarity_sharpness"},"link":null},{"localized_name":"override_face_bbox","name":"override_face_bbox","shape":7,"type":"STRING","widget":{"name":"override_face_bbox"},"link":null},{"localized_name":"skip_when_sigma_above","name":"skip_when_sigma_above","shape":7,"type":"FLOAT","widget":{"name":"skip_when_sigma_above"},"link":null},{"localized_name":"pull_mode","name":"pull_mode","shape":7,"type":"COMBO","widget":{"name":"pull_mode"},"link":null},{"localized_name":"late_block_falloff","name":"late_block_falloff","shape":7,"type":"FLOAT","widget":{"name":"late_block_falloff"},"link":null}],"outputs":[{"localized_name":"MODEL","name":"MODEL","type":"MODEL","links":[2372]}],"properties":{"cnr_id":"10s-comfy-nodes","ver":"2fcc4ac6018ae6fc6c2ff64c2e2a58f9e50ba30a","Node name for S&R":"LTXLikenessAnchor"},"widgets_values":[0.5,"auto","",0.5,0,false,false,false,"middle","",8,"",0,"directional",0.4],"color":"#233","bgcolor":"#355"},{"id":806,"type":"LTXLikenessGuide","pos":[213.56489728537017,3875.3816888634133],"size":[322.642578125,426],"flags":{},"order":65,"mode":0,"inputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":2286},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":2287},{"localized_name":"vae","name":"vae","type":"VAE","link":2288},{"localized_name":"latent","name":"latent","type":"LATENT","link":2289},{"localized_name":"image","name":"image","type":"IMAGE","link":2290},{"localized_name":"strength","name":"strength","shape":7,"type":"FLOAT","widget":{"name":"strength"},"link":null},{"localized_name":"placement_mode","name":"placement_mode","shape":7,"type":"COMBO","widget":{"name":"placement_mode"},"link":null},{"localized_name":"face_detect","name":"face_detect","shape":7,"type":"COMBO","widget":{"name":"face_detect"},"link":null},{"localized_name":"reference_mask_mode","name":"reference_mask_mode","shape":7,"type":"COMBO","widget":{"name":"reference_mask_mode"},"link":null},{"localized_name":"face_padding","name":"face_padding","shape":7,"type":"FLOAT","widget":{"name":"face_padding"},"link":null},{"localized_name":"crf","name":"crf","shape":7,"type":"INT","widget":{"name":"crf"},"link":null},{"localized_name":"blur_radius","name":"blur_radius","shape":7,"type":"INT","widget":{"name":"blur_radius"},"link":null},{"localized_name":"interpolation","name":"interpolation","shape":7,"type":"COMBO","widget":{"name":"interpolation"},"link":null},{"localized_name":"crop","name":"crop","shape":7,"type":"COMBO","widget":{"name":"crop"},"link":null},{"localized_name":"attention_strength","name":"attention_strength","shape":7,"type":"FLOAT","widget":{"name":"attention_strength"},"link":null},{"localized_name":"face_bbox_within_reference","name":"face_bbox_within_reference","shape":7,"type":"STRING","widget":{"name":"face_bbox_within_reference"},"link":2353},{"localized_name":"emit_latent","name":"emit_latent","shape":7,"type":"COMBO","widget":{"name":"emit_latent"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"positive","name":"positive","type":"CONDITIONING","links":[2291,2292]},{"localized_name":"negative","name":"negative","type":"CONDITIONING","links":[2293,2294]},{"localized_name":"latent","name":"latent","type":"LATENT","links":[2303]},{"localized_name":"reference_info","name":"reference_info","type":"REFERENCE_INFO","links":[2300,2360]}],"properties":{"cnr_id":"10s-comfy-nodes","ver":"2fcc4ac6018ae6fc6c2ff64c2e2a58f9e50ba30a","Node name for S&R":"LTXLikenessGuide"},"widgets_values":[0.8999999999999999,"silent_reference","manual","bbox_only",0.15,24,0,"area","center",1,"","passthrough",false],"color":"#233","bgcolor":"#355"},{"id":582,"type":"ManualSigmas","pos":[3512.1741421202964,3254.6564277665666],"size":[322.9496561279293,69.61597045898407],"flags":{},"order":28,"mode":0,"inputs":[{"localized_name":"sigmas","name":"sigmas","type":"STRING","widget":{"name":"sigmas"},"link":null}],"outputs":[{"localized_name":"SIGMAS","name":"SIGMAS","type":"SIGMAS","links":[2268]}],"title":"Sigmas Upscale","properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"ManualSigmas","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["0.715, 0.4824, 0.2412, 0.0"],"color":"#323","bgcolor":"#535"},{"id":797,"type":"mxSlider","pos":[-1248.5673821105443,4239.186948807451],"size":[239.23083288068028,30],"flags":{"collapsed":false},"order":29,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"FLOAT","links":[2219]}],"title":"First Frame Strength","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":0.77,"min":0.4,"max":1,"step":0.01,"decimals":2,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[0,0.77,1],"color":"#222","bgcolor":"#000"},{"id":557,"type":"Power Lora Loader (rgthree)","pos":[-772.3542382588616,3380.3379658151052],"size":[428.9642467678492,322.75598817305604],"flags":{},"order":43,"mode":0,"inputs":[{"dir":3,"name":"model","type":"MODEL","link":1547},{"dir":3,"name":"clip","type":"CLIP","link":1418}],"outputs":[{"dir":4,"name":"MODEL","shape":3,"type":"MODEL","links":[1852,1859]},{"dir":4,"name":"CLIP","shape":3,"type":"CLIP","links":null}],"properties":{"cnr_id":"rgthree-comfy","ver":"1.0.2512112053","Show Strengths":"Single Strength","Match":"","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[{},{"type":"PowerLoraLoaderHeaderWidget"},{"on":true,"lora":"ltx23\\OmniNFT_converted_lora.safetensors","strength":0.8,"strengthTwo":null},{},""],"color":"#222","bgcolor":"#000"},{"id":524,"type":"Seed (rgthree)","pos":[-1098.5329029499685,4040.8382724031653],"size":[288.55371900826447,130],"flags":{},"order":30,"mode":0,"inputs":[],"outputs":[{"dir":4,"name":"SEED","shape":3,"type":"INT","links":[1273,1426]}],"properties":{"cnr_id":"rgthree-comfy","ver":"1.0.2511270846","randomMax":1125899906842624,"randomMin":0,"ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[-1,"","",""],"color":"#222","bgcolor":"#000"},{"id":779,"type":"JWStringToFloat","pos":[-131.4793794648068,2934.4688117206065],"size":[270,58],"flags":{},"order":31,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"FLOAT","name":"FLOAT","type":"FLOAT","links":[2180]}],"title":"First Pass Video Strength","properties":{"cnr_id":"comfyui-various","ver":"5bd85aaf7616878471469c4ec7e11bbd0cef3bf2","Node name for S&R":"JWStringToFloat","ue_properties":{"widget_ue_connectable":{"text":true},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["0.84"],"color":"#222","bgcolor":"#000"},{"id":791,"type":"mxSlider","pos":[-1408.6283418078767,4062.6513388998037],"size":[280.438609128239,30],"flags":{"collapsed":false},"order":32,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"INT","links":[2215]}],"title":"Video Width","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":1120,"min":480,"max":2400,"step":32,"decimals":0,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[1120,1120,0],"color":"#222","bgcolor":"#000"},{"id":792,"type":"mxSlider","pos":[-1409.273145451266,4160.645577420733],"size":[280.438609128239,30],"flags":{"collapsed":false},"order":33,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"INT","links":[2216]}],"title":"Video Height","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":1344,"min":480,"max":2400,"step":32,"decimals":0,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[1344,1344,0],"color":"#222","bgcolor":"#000"},{"id":834,"type":"LoadImage","pos":[-1313.9358666037338,4471.535870455174],"size":[557.1987915039062,864.3999328613281],"flags":{},"order":34,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[2371]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.21.1","Node name for S&R":"LoadImage"},"widgets_values":["sex-with-ellie-mr-strauz-v0-2lp0irchqrxg1.webp","image"],"color":"#222","bgcolor":"#000"},{"id":795,"type":"MarkdownNote","pos":[-2021.1608937808342,3197.3792734479475],"size":[521.3413669628835,1161.0100683552218],"flags":{},"order":35,"mode":0,"inputs":[],"outputs":[],"properties":{"ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":["# 10S Workflows\n\nThese are advanced quality-focused conditioned workflows for 10Eros/Sulphur/Base LTX 2.3. They make some compromises on speed for quality. \nSteal these or study 10S Node intergation if you are a workflow maker or want to customize and release them as your own. \n**They are NOT FAST.** I find with LTX, fast ≠ good.\n\n---\n\n## 📦 Downloads\n\n| Resource | Link |\n|----------|------|\n| 🎬 Workflows | [LTX2.3-10Eros Workflows](https://huggingface.co/TenStrip/LTX2.3-10Eros_Workflows) |\n| 🧠 Models (10Eros/Sulphur) | [huggingface.co/TenStrip/LTX2.3-10Eros](https://huggingface.co/TenStrip/LTX2.3-10Eros) |\n| 🔧 10S Comfy Nodes | [TenStrip/10S-Comfy-nodes](https://github.com/TenStrip/10S-Comfy-nodes) |\n| ⚗️ Distilled LoRAs | [LTX2.3 Distilled LoRA 1.1 Experiments](https://huggingface.co/TenStrip/LTX2.3_Distilled_Lora_1.1_Experiments/tree/main) |\n| 🔼 OMNINFT LoRA | [Omni NFT 2.3 Conversion - Quality/Reward LoRA](https://huggingface.co/VasiliyWeb/OmniNFT_ComfyUI/tree/main) |\n---\n\n## ⚙️ How It Works\n\nThese workflows rely on **[10S-Comfy-nodes]**\n\nMouse over parameters for info.\n\n\nFace Guider/Likeness Anchor are there to **reinforce likeness** they do not teach the model what the initial figure looks like but they curb it's attempts to overwrite finer facial details. Complex face altering movements like turning away, camera zoom outs, and transitions will destroy the initial reference regardless. The only way to inject a true character likeness is to train a small and quick character lora (5-30 minutes ~ 10-30 images) for Sulphur or base LTX2.3 and use that to alter the actual underlying weights with gradient data.\n\nGuided prompt conditioning from the guide is leveraged against the i2v conditioned latent with Latent anchor. Increasing strength on Latent anchor as well as **depth_curve** gives you a more manageable control over less/more movement and gives you 3 points of strength to control the i2v freedom. \n\nLatent Anchor node needs paper-level study in what it does, but for most calculations turn debug on and follow this formula to set it up correctly: \n*target_call_count = cache_at_step × forwards_per_step*\n\nIt is currently configured for 13 steps at step 6 cache with 1 CFG, although forwards per step must be increased to 3 to account for CFG,STG,audioCFG-but isn't impactful on just first 2 CFG steps.\n\nUpscale distortion is removed and detail is increased with tiled sampler, this is a default-level improvement to quality. At it's current setting it will only trigger when distortion inducing resolution is reached for vertical ratios, if working in landscape increase max_tile_size accordingly to bypass it. Audio-sync issues can arise if the **tile_carrying** mode does not target the tile with talking heads or sound-causing actors. Generally you will select first for vertical resolutions and middle for huge landscape passes.\nBecause of how tiled sampler works you will want to avoid hand motion and movement near the middle of the composition.\n\nFor Landscape formats **bypass tiling**, it is not necessary to correct that ratio.\n\nDistillation is handled at full break down. Lower pass video distillation will increase motion. Upscale pass video distillation will add detail, but also add artifacts and unaesthetic blur artifacts. Audio weight must be kept high for audio quality.\n\nIn some cases first pass audio may be more desirable, that can be toggled to be hard encoded instead of passed and resampled fully.\n\nThe first pass output is a full **motion** preview, if it contains mutation, distortion, bad motion, bad audio, etc... The upscale pass will not fix it. Only allow it to run to upscale when first pass motion looks good. **Visual quality and style will be restored** at upscale.\n\n\n---\n\n## ✍️ Prompt Enhancement\n\nThese use STG Guider and have CFG for initial steps which is where most broad motion and audio form. **Your negative prompt is impactful, so use it.**\n\nFor best positive prompt results, use this foreword in **Grok** or an **uncensored LLM**:\n\n---\n\n*Generate a video scene script with a description based on the attached image for an LLM that has a tokenizer that uses interleaved attention to support long-context understanding that is fed into a multimodal video model. Strict specification, follow up to the word: No timestamps. No unnecessary embellishment. Output only plain English text and make it a copy box.*\n\n*First, describe the image initial scene in concise natural language; subject(s), subject(s) appearance, subject(s) composition and pose, background, and context.*\n\n*Next, formulate a naturally evolving scenario that would take place describing every moving body part, composition change, and manipulation from the uploaded initial frame that would be reflected in the video models post-latent evolution output. If the image is explicit or sexual in nature, use full anatomical terminology and spice it up slightly with visually representable erotic themes.*\n\n*Center the prompt around this basic idea: **[ concept ]***\n\n*Interweave this dialogue or sound concept into the scene with descriptions of voice tone followed by the lines delivered in quotations, in a temporal sequence between or during motions. Dialogue should be concise and non-rambling as it will take away from video quality: **[ dialogue ]***\n\n*Inside that prompt describe only notable audio and audio cues, both normal and explicit; background noise as well as foley and natural sounds. In a temporal sequence paired with coinciding motions. In the case of absent dialogue or soundscapes and only if background music is fitting; describe a fitting genre and melodic tone with matching mood.*\n\n*Output only text following above instruction. Follow-up suggestions should be on the topic of expanding or changing motion or dialogue from the output text.*\n"],"color":"#432","bgcolor":"#653"},{"id":796,"type":"mxSlider","pos":[-1115.3480311403682,3957.2240104480748],"size":[329.41939288068033,30],"flags":{"collapsed":false},"order":36,"mode":0,"inputs":[{"localized_name":"Xi","name":"Xi","type":"INT","widget":{"name":"Xi"},"link":null},{"localized_name":"Xf","name":"Xf","type":"FLOAT","widget":{"name":"Xf"},"link":null},{"localized_name":"isfloatX","name":"isfloatX","type":"INT","widget":{"name":"isfloatX"},"link":null}],"outputs":[{"localized_name":"","name":"","type":"INT","links":[2220]}],"title":"Length (Frame Count X/24 = Seconds)","properties":{"cnr_id":"comfyui-mxtoolkit","ver":"0.9.92","Node name for S&R":"mxSlider","value":240,"min":48,"max":960,"step":24,"decimals":0,"snap":true,"ue_properties":{"version":"7.1","widget_ue_connectable":{},"input_ue_unconnectable":{}},"aux_id":"Smirnov75/ComfyUI-mxToolkit"},"widgets_values":[240,240,0],"color":"#222","bgcolor":"#000"},{"id":537,"type":"CLIPTextEncode","pos":[-622.2981226504402,4084.6379532384985],"size":[397.87115230398985,118.93963289899375],"flags":{"collapsed":false},"order":45,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":1420},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","links":[2143]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"CLIPTextEncode","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":["captions, music, transition, VR, bad quality, subtitles, text, watermark, overlay effects, pc game, yelling, console game, video game, cartoon, childish, ugly, link, text, blur, logo, wordmark, static, low quality, noise, white noise, bleep, censoring, censor, bleeping, beep, beeping, newscast, interview, podcast, mutant, horror, 70's, film grain, cinematic, comedy, "],"color":"#322","bgcolor":"#533"},{"id":550,"type":"LTXVAudioVAEDecode","pos":[734.1180720394716,3961.1649068339216],"size":[240,46],"flags":{},"order":76,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":1608},{"label":"Audio VAE","localized_name":"audio_vae","name":"audio_vae","type":"VAE","link":1305}],"outputs":[{"localized_name":"Audio","name":"Audio","type":"AUDIO","links":[1306,1602]}],"properties":{"cnr_id":"comfy-core","ver":"0.7.0","Node name for S&R":"LTXVAudioVAEDecode","ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}},"enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65},"widgets_values":[]},{"id":552,"type":"VAEDecode","pos":[732.7767779131676,3815.2725002657667],"size":[254.98011089086322,46],"flags":{},"order":75,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":2306},{"localized_name":"vae","name":"vae","type":"VAE","link":1307}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[1303]}],"properties":{"cnr_id":"comfy-core","ver":"0.9.2","Node name for S&R":"VAEDecode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.5.2"}},"widgets_values":[]},{"id":556,"type":"LTXVSeparateAVLatent","pos":[740.0229196133681,3693.150701447153],"size":[244.39530248670872,49.77267672823109],"flags":{},"order":73,"mode":0,"inputs":[{"localized_name":"av_latent","name":"av_latent","type":"LATENT","link":1546}],"outputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","links":[2259,2306]},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","links":[1608,2213]}],"properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"LTXVSeparateAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":510,"type":"SamplerCustomAdvanced","pos":[737.7202132201912,3461.078040820131],"size":[212.3638671875,106],"flags":{},"order":72,"mode":0,"inputs":[{"localized_name":"noise","name":"noise","type":"NOISE","link":1274},{"localized_name":"guider","name":"guider","type":"GUIDER","link":2043},{"localized_name":"sampler","name":"sampler","type":"SAMPLER","link":1276},{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":1569},{"localized_name":"latent_image","name":"latent_image","type":"LATENT","link":1531}],"outputs":[{"localized_name":"output","name":"output","type":"LATENT","links":[]},{"localized_name":"denoised_output","name":"denoised_output","type":"LATENT","links":[1546]}],"title":"First Pass","properties":{"cnr_id":"comfy-core","ver":"0.20.1","Node name for S&R":"SamplerCustomAdvanced","ue_properties":{"widget_ue_connectable":{},"version":"7.8","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#323","bgcolor":"#535"},{"id":802,"type":"LTXTiledSampler","pos":[3868.938281719194,3309.5082249176958],"size":[279.60546875,306],"flags":{},"order":84,"mode":0,"inputs":[{"localized_name":"noise","name":"noise","type":"NOISE","link":2265},{"localized_name":"guider","name":"guider","type":"GUIDER","link":2266},{"localized_name":"sampler","name":"sampler","type":"SAMPLER","link":2267},{"localized_name":"sigmas","name":"sigmas","type":"SIGMAS","link":2268},{"localized_name":"latent_image","name":"latent_image","type":"LATENT","link":2269},{"localized_name":"bypass_tiling","name":"bypass_tiling","shape":7,"type":"BOOLEAN","widget":{"name":"bypass_tiling"},"link":null},{"localized_name":"tile_axis","name":"tile_axis","shape":7,"type":"COMBO","widget":{"name":"tile_axis"},"link":null},{"localized_name":"n_tiles","name":"n_tiles","shape":7,"type":"INT","widget":{"name":"n_tiles"},"link":null},{"localized_name":"tile_overlap","name":"tile_overlap","shape":7,"type":"INT","widget":{"name":"tile_overlap"},"link":null},{"localized_name":"max_size_for_no_tile","name":"max_size_for_no_tile","shape":7,"type":"INT","widget":{"name":"max_size_for_no_tile"},"link":null},{"localized_name":"audio_pass","name":"audio_pass","shape":7,"type":"COMBO","widget":{"name":"audio_pass"},"link":null},{"localized_name":"audio_carrier_tile","name":"audio_carrier_tile","shape":7,"type":"COMBO","widget":{"name":"audio_carrier_tile"},"link":null},{"localized_name":"debug","name":"debug","shape":7,"type":"BOOLEAN","widget":{"name":"debug"},"link":null}],"outputs":[{"localized_name":"output","name":"output","type":"LATENT","links":null},{"localized_name":"denoised_output","name":"denoised_output","type":"LATENT","links":[2270]}],"properties":{"cnr_id":"10s-comfy-nodes","ver":"1f58fea9fba0714538f9e83ddddd89bce3068361","Node name for S&R":"LTXTiledSampler"},"widgets_values":[false,"auto",2,8,24,"tile_carrying","first",false],"color":"#323","bgcolor":"#535"},{"id":591,"type":"LTXVConcatAVLatent","pos":[3146.6765223307043,3745.3069037036903],"size":[270,46],"flags":{},"order":83,"mode":0,"inputs":[{"localized_name":"video_latent","name":"video_latent","type":"LATENT","link":2272},{"localized_name":"audio_latent","name":"audio_latent","type":"LATENT","link":2214}],"outputs":[{"localized_name":"latent","name":"latent","type":"LATENT","links":[2269]}],"properties":{"cnr_id":"comfy-core","ver":"0.5.1","Node name for S&R":"LTXVConcatAVLatent","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":536,"type":"CLIPTextEncode","pos":[-622.7470436804905,3775.9097381631527],"size":[394.7031737883649,255.53167788014707],"flags":{},"order":44,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":1419},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":null}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","links":[2142]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.56","Node name for S&R":"CLIPTextEncode","enableTabs":false,"tabWidth":65,"tabXOffset":10,"hasSecondTab":false,"secondTabText":"Send Back","secondTabOffset":80,"secondTabWidth":65,"ue_properties":{"widget_ue_connectable":{},"version":"7.5.2","input_ue_unconnectable":{}}},"widgets_values":[""],"color":"#232","bgcolor":"#353"}],"links":[[583,274,0,275,0,"LATENT"],[586,279,0,274,1,"VAE"],[1272,525,0,526,0,"NOISE"],[1273,524,0,525,0,"INT"],[1274,526,0,510,0,"NOISE"],[1276,520,0,510,2,"SAMPLER"],[1281,531,0,532,0,"IMAGE"],[1282,532,0,533,0,"IMAGE"],[1289,533,0,534,0,"INT"],[1290,533,1,534,1,"INT"],[1291,540,0,539,0,"INT"],[1292,542,0,540,0,"FLOAT"],[1293,542,0,541,0,"FLOAT"],[1295,543,0,535,1,"INT"],[1296,545,0,535,2,"INT"],[1297,546,0,534,2,"INT"],[1303,552,0,549,0,"IMAGE"],[1305,554,0,550,1,"VAE"],[1306,550,0,551,0,"AUDIO"],[1307,553,0,552,1,"VAE"],[1308,555,0,549,4,"FLOAT"],[1318,558,0,535,0,"VAE"],[1320,560,0,523,2,"FLOAT"],[1382,592,0,593,1,"VAE"],[1387,593,0,598,0,"AUDIO"],[1418,616,0,557,1,"CLIP"],[1419,616,0,536,0,"CLIP"],[1420,616,0,537,0,"CLIP"],[1423,617,0,207,0,"VAE"],[1426,524,0,621,0,"INT"],[1457,630,0,275,1,"MASK"],[1481,635,0,597,4,"FLOAT"],[1531,548,0,510,4,"LATENT"],[1546,510,1,556,0,"LATENT"],[1547,646,0,557,0,"MODEL"],[1548,646,2,214,0,"VAE"],[1566,533,0,630,1,"INT"],[1567,533,1,630,2,"INT"],[1568,527,0,652,0,"SIGMAS"],[1569,652,0,510,3,"SIGMAS"],[1602,550,0,274,0,"AUDIO"],[1608,556,1,550,0,"LATENT"],[1611,596,1,593,0,"LATENT"],[1714,551,0,549,1,"AUDIO"],[1852,557,0,719,0,"MODEL"],[1853,719,0,718,0,"MODEL"],[1859,557,0,722,0,"MODEL"],[1860,722,0,723,0,"MODEL"],[1886,559,0,731,2,"VAE"],[1888,527,0,731,4,"SIGMAS"],[1942,594,0,740,1,"VAE"],[1956,743,0,744,1,"LATENT_UPSCALE_MODEL"],[1957,589,0,744,2,"VAE"],[2025,740,0,755,0,"IMAGE"],[2026,755,0,597,0,"IMAGE"],[2043,653,0,510,1,"GUIDER"],[2088,723,0,731,0,"MODEL"],[2103,531,0,769,0,"IMAGE"],[2106,531,0,331,0,"IMAGE"],[2107,714,0,770,4,"IMAGE"],[2109,589,0,770,0,"VAE"],[2111,531,0,731,1,"IMAGE"],[2118,559,0,772,0,"VAE"],[2140,531,0,772,4,"IMAGE"],[2142,536,0,523,0,"CONDITIONING"],[2143,537,0,523,1,"CONDITIONING"],[2179,593,0,597,1,"AUDIO"],[2180,779,0,722,4,"FLOAT"],[2181,780,0,723,4,"FLOAT"],[2182,781,0,718,4,"FLOAT"],[2183,782,0,719,4,"FLOAT"],[2184,783,0,722,3,"COMBO"],[2185,783,0,723,3,"COMBO"],[2186,783,0,719,3,"COMBO"],[2187,783,0,718,3,"COMBO"],[2212,274,0,789,0,"LATENT"],[2213,556,1,789,1,"LATENT"],[2214,789,0,591,1,"LATENT"],[2215,791,0,531,2,"INT"],[2216,792,0,531,3,"INT"],[2219,797,0,772,3,"FLOAT"],[2220,796,0,798,0,"INT"],[2221,798,1,538,0,"INT"],[2236,535,0,548,1,"LATENT"],[2259,556,0,744,0,"LATENT"],[2265,584,0,802,0,"NOISE"],[2266,583,0,802,1,"GUIDER"],[2267,585,0,802,2,"SAMPLER"],[2268,582,0,802,3,"SIGMAS"],[2269,591,0,802,4,"LATENT"],[2270,802,1,596,0,"LATENT"],[2271,744,0,770,1,"LATENT"],[2272,770,0,591,0,"LATENT"],[2286,523,0,806,0,"CONDITIONING"],[2287,523,1,806,1,"CONDITIONING"],[2288,559,0,806,2,"VAE"],[2289,534,0,806,3,"LATENT"],[2290,531,0,806,4,"IMAGE"],[2291,806,0,583,1,"CONDITIONING"],[2292,806,0,653,1,"CONDITIONING"],[2293,806,1,583,2,"CONDITIONING"],[2294,806,1,653,2,"CONDITIONING"],[2303,806,2,772,1,"LATENT"],[2304,772,0,548,0,"LATENT"],[2305,772,0,731,3,"LATENT"],[2306,556,0,552,0,"LATENT"],[2307,596,0,740,0,"LATENT"],[2310,532,0,808,0,"IMAGE"],[2353,808,0,806,15,"STRING"],[2359,731,0,827,0,"MODEL"],[2360,806,3,827,1,"REFERENCE_INFO"],[2361,808,0,827,4,"STRING"],[2366,718,0,583,0,"MODEL"],[2371,834,0,531,0,"IMAGE"],[2372,827,0,653,0,"MODEL"]],"groups":[{"id":27,"title":"Main","bounding":[-1470.686865946572,3171.0519920330685,3804.2302523868716,2304.222673775987],"color":"#3f789e","flags":{}},{"id":23,"title":"Force First-pass Audio","bounding":[1228.7248700088794,2801.335742850906,537.4514362675345,323.87650696724893],"color":"#3f789e","flags":{}},{"id":29,"title":"Final Upscale Pass","bounding":[2374.9914506547025,3167.3647232660733,2118.8972740162553,1955.3713894900138],"color":"#3f789e","flags":{}},{"id":30,"title":"Distilation Break-up (Thanks Lightricks For This Wonderful Distilled Lora)","bounding":[-167.0518161898845,2841.0578682916353,1330.1373185243046,278.6386435094696],"color":"#3f789e","flags":{}},{"id":32,"title":"Post Options, Add interpolation here before RTX","bounding":[4693.505674922303,3438.368751942385,687.0438904788525,342.37525726057083],"color":"#3f789e","flags":{}},{"id":33,"title":"Parameters","bounding":[-1419.7996039355623,3824.1199049135266,642.0089842779156,461.4995675177156],"color":"#3f789e","flags":{}}],"config":{},"extra":{"ds":{"scale":1,"offset":[1347.2068398473516,-3519.212712047732]},"frontendVersion":"1.43.18","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":false,"VHS_latentpreviewrate":0,"VHS_MetadataImage":true,"VHS_KeepIntermediate":true,"ue_links":[],"links_added_by_ue":[]},"version":0.4} \ No newline at end of file diff --git a/custom_nodes/ComfyUI-GGUF/tools/README.md b/custom_nodes/ComfyUI-GGUF/tools/README.md new file mode 100644 index 0000000000000000000000000000000000000000..228bc220593d7878ca854032638bdc2682c76a71 --- /dev/null +++ b/custom_nodes/ComfyUI-GGUF/tools/README.md @@ -0,0 +1,93 @@ +## Converting initial model + +To convert your initial safetensors/ckpt model to FP16/BF16 GGUF, run the following command: + +``` +python convert.py --src E:\models\unet\flux1-dev.safetensors +``` +Make sure `gguf>=0.13.0` is installed for this step. Optionally, specify the output gguf file with the `--dst` arg. + +> [!NOTE] +> Do not use the diffusers UNET format for flux, it won't work, use the default/reference checkpoint key format. This is due to q/k/v being merged into one qkv key. +> You can convert it by loading it in ComfyUI and saving it using the built-in "ModelSave" node. + +> [!WARNING] +> For hunyuan video/wan 2.1, you will see a warning about 5D tensors. This means the script will save a **non functional** model to disk first, that you can quantize. I recommend saving these in a separate `raw` folder to avoid confusion. +> +> After quantization, you will have to run `fix_5d_tensor.py` manually to add back the missing key that was saved by the conversion code. + +## Quantizing using custom llama.cpp + +Depending on your git settings, you may need to run the following script first in order to make sure the patch file is valid. It will convert Windows (CRLF) line endings to Unix (LF) ones. + +``` +python fix_lines_ending.py +``` + +Git clone llama.cpp into the current folder: + +``` +git clone https://github.com/ggerganov/llama.cpp +``` + +Check out the correct branch, then apply the custom patch needed to add image model support to the repo you just cloned. + +``` +cd llama.cpp +git checkout tags/b3962 +git apply ..\lcpp.patch +``` + +Compile the llama-quantize binary. This example uses cmake, on linux you can just use make. + +### Visual Studio 2019, Linux, etc... + +``` +mkdir build +cmake -B build +cmake --build build --config Debug -j10 --target llama-quantize +cd .. +``` + +### Visual Studio 2022 + +``` +mkdir build +cmake -B build -DCMAKE_CXX_STANDARD=17 -DCMAKE_CXX_STANDARD_REQUIRED=ON -DCMAKE_CXX_FLAGS="-std=c++17" +``` + +Edit the `llama.cpp\common\log.cpp` file, inserts two lines after the existing first line: + +``` +#include "log.h" + +#define _SILENCE_CXX23_CHRONO_DEPRECATION_WARNING +#include +``` + +Then you can build the project: +``` +cmake --build build --config Debug -j10 --target llama-quantize +cd .. +``` + +### Quantize your model + + +Now you can use the newly build binary to quantize your model to the desired format: +``` +llama.cpp\build\bin\Debug\llama-quantize.exe E:\models\unet\flux1-dev-BF16.gguf E:\models\unet\flux1-dev-Q4_K_S.gguf Q4_K_S +``` + +You can extract the patch again with `git diff src\llama.cpp > lcpp.patch` if you wish to change something and contribute back. + +> [!WARNING] +> For hunyuan video/wan 2.1, you will have to run `fix_5d_tensor.py` after the quantization step is done. +> +> Example usage: `fix_5d_tensors.py --src E:\models\video\raw\wan2.1-t2v-1.3b-Q8_0.gguf --dst E:\models\video\wan2.1-t2v-1.3b-Q8_0.gguf` +> +> By default, this also saves a `fix_5d_tensors_[arch].safetensors` file in the `ComfyUI-GGUF/tools` folder, it's recommended to delete this after all models have been converted. + +> [!NOTE] +> Do not quantize SDXL / SD1 / other Conv2D heavy models. If you do, make sure to **extract the UNET model first**. +>This should be obvious, but also don't use the resulting llama-quantize binary with LLMs. diff --git a/custom_nodes/ComfyUI-GGUF/tools/convert.py b/custom_nodes/ComfyUI-GGUF/tools/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..5029c874277358559f8855d3f1032963437a3e91 --- /dev/null +++ b/custom_nodes/ComfyUI-GGUF/tools/convert.py @@ -0,0 +1,365 @@ +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import os +import gguf +import torch +import logging +import argparse +from tqdm import tqdm +from safetensors.torch import load_file, save_file + +QUANTIZATION_THRESHOLD = 1024 +REARRANGE_THRESHOLD = 512 +MAX_TENSOR_NAME_LENGTH = 127 +MAX_TENSOR_DIMS = 4 + +class ModelTemplate: + arch = "invalid" # string describing architecture + shape_fix = False # whether to reshape tensors + keys_detect = [] # list of lists to match in state dict + keys_banned = [] # list of keys that should mark model as invalid for conversion + keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason + keys_ignore = [] # list of strings to ignore keys by when found + + def handle_nd_tensor(self, key, data): + raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})") + +class ModelFlux(ModelTemplate): + arch = "flux" + keys_detect = [ + ("transformer_blocks.0.attn.norm_added_k.weight",), + ("double_blocks.0.img_attn.proj.weight",), + ] + keys_banned = ["transformer_blocks.0.attn.norm_added_k.weight",] + +class ModelSD3(ModelTemplate): + arch = "sd3" + keys_detect = [ + ("transformer_blocks.0.attn.add_q_proj.weight",), + ("joint_blocks.0.x_block.attn.qkv.weight",), + ] + keys_banned = ["transformer_blocks.0.attn.add_q_proj.weight",] + +class ModelAura(ModelTemplate): + arch = "aura" + keys_detect = [ + ("double_layers.3.modX.1.weight",), + ("joint_transformer_blocks.3.ff_context.out_projection.weight",), + ] + keys_banned = ["joint_transformer_blocks.3.ff_context.out_projection.weight",] + +class ModelHiDream(ModelTemplate): + arch = "hidream" + keys_detect = [ + ( + "caption_projection.0.linear.weight", + "double_stream_blocks.0.block.ff_i.shared_experts.w3.weight" + ) + ] + keys_hiprec = [ + # nn.parameter, can't load from BF16 ver + ".ff_i.gate.weight", + "img_emb.emb_pos" + ] + +class CosmosPredict2(ModelTemplate): + arch = "cosmos" + keys_detect = [ + ( + "blocks.0.mlp.layer1.weight", + "blocks.0.adaln_modulation_cross_attn.1.weight", + ) + ] + keys_hiprec = ["pos_embedder"] + keys_ignore = ["_extra_state", "accum_"] + +class ModelHyVid(ModelTemplate): + arch = "hyvid" + keys_detect = [ + ( + "double_blocks.0.img_attn_proj.weight", + "txt_in.individual_token_refiner.blocks.1.self_attn_qkv.weight", + ) + ] + + def handle_nd_tensor(self, key, data): + # hacky but don't have any better ideas + path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? + if os.path.isfile(path): + raise RuntimeError(f"5D tensor fix file already exists! {path}") + fsd = {key: torch.from_numpy(data)} + tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") + save_file(fsd, path) + +class ModelWan(ModelHyVid): + arch = "wan" + keys_detect = [ + ( + "blocks.0.self_attn.norm_q.weight", + "text_embedding.2.weight", + "head.modulation", + ) + ] + keys_hiprec = [ + ".modulation" # nn.parameter, can't load from BF16 ver + ] + +class ModelLTXV(ModelTemplate): + arch = "ltxv" + keys_detect = [ + ( + "adaln_single.emb.timestep_embedder.linear_2.weight", + "transformer_blocks.27.scale_shift_table", + "caption_projection.linear_2.weight", + ) + ] + keys_hiprec = [ + "scale_shift_table" # nn.parameter, can't load from BF16 base quant + ] + +class ModelSDXL(ModelTemplate): + arch = "sdxl" + shape_fix = True + keys_detect = [ + ("down_blocks.0.downsamplers.0.conv.weight", "add_embedding.linear_1.weight",), + ( + "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", + "output_blocks.2.2.conv.weight", "output_blocks.5.2.conv.weight", + ), # Non-diffusers + ("label_emb.0.0.weight",), + ] + +class ModelSD1(ModelTemplate): + arch = "sd1" + shape_fix = True + keys_detect = [ + ("down_blocks.0.downsamplers.0.conv.weight",), + ( + "input_blocks.3.0.op.weight", "input_blocks.6.0.op.weight", "input_blocks.9.0.op.weight", + "output_blocks.2.1.conv.weight", "output_blocks.5.2.conv.weight", "output_blocks.8.2.conv.weight" + ), # Non-diffusers + ] + +class ModelLumina2(ModelTemplate): + arch = "lumina2" + keys_detect = [ + ("cap_embedder.1.weight", "context_refiner.0.attention.qkv.weight") + ] + +arch_list = [ModelFlux, ModelSD3, ModelAura, ModelHiDream, CosmosPredict2, + ModelLTXV, ModelHyVid, ModelWan, ModelSDXL, ModelSD1, ModelLumina2] + +def is_model_arch(model, state_dict): + # check if model is correct + matched = False + invalid = False + for match_list in model.keys_detect: + if all(key in state_dict for key in match_list): + matched = True + invalid = any(key in state_dict for key in model.keys_banned) + break + assert not invalid, "Model architecture not allowed for conversion! (i.e. reference VS diffusers format)" + return matched + +def detect_arch(state_dict): + model_arch = None + for arch in arch_list: + if is_model_arch(arch, state_dict): + model_arch = arch() + break + assert model_arch is not None, "Unknown model architecture!" + return model_arch + +def parse_args(): + parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET") + parser.add_argument("--src", required=True, help="Source model ckpt file.") + parser.add_argument("--dst", help="Output unet gguf file.") + args = parser.parse_args() + + if not os.path.isfile(args.src): + parser.error("No input provided!") + + return args + +def strip_prefix(state_dict): + # prefix for mixed state dict + prefix = None + for pfx in ["model.diffusion_model.", "model."]: + if any([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # prefix for uniform state dict + if prefix is None: + for pfx in ["net."]: + if all([x.startswith(pfx) for x in state_dict.keys()]): + prefix = pfx + break + + # strip prefix if found + if prefix is not None: + logging.info(f"State dict prefix found: '{prefix}'") + sd = {} + for k, v in state_dict.items(): + if prefix not in k: + continue + k = k.replace(prefix, "") + sd[k] = v + else: + logging.debug("State dict has no prefix") + sd = state_dict + + return sd + +def load_state_dict(path): + if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): + state_dict = torch.load(path, map_location="cpu", weights_only=True) + for subkey in ["model", "module"]: + if subkey in state_dict: + state_dict = state_dict[subkey] + break + if len(state_dict) < 20: + raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}") + else: + state_dict = load_file(path) + + return strip_prefix(state_dict) + +def handle_tensors(writer, state_dict, model_arch): + name_lengths = tuple(sorted( + ((key, len(key)) for key in state_dict.keys()), + key=lambda item: item[1], + reverse=True, + )) + if not name_lengths: + return + max_name_len = name_lengths[0][1] + if max_name_len > MAX_TENSOR_NAME_LENGTH: + bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH) + raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}") + for key, data in tqdm(state_dict.items()): + old_dtype = data.dtype + + if any(x in key for x in model_arch.keys_ignore): + tqdm.write(f"Filtering ignored key: '{key}'") + continue + + if data.dtype == torch.bfloat16: + data = data.to(torch.float32).numpy() + # this is so we don't break torch 2.0.X + elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]: + data = data.to(torch.float16).numpy() + else: + data = data.numpy() + + n_dims = len(data.shape) + data_shape = data.shape + if old_dtype == torch.bfloat16: + data_qtype = gguf.GGMLQuantizationType.BF16 + # elif old_dtype == torch.float32: + # data_qtype = gguf.GGMLQuantizationType.F32 + else: + data_qtype = gguf.GGMLQuantizationType.F16 + + # The max no. of dimensions that can be handled by the quantization code is 4 + if len(data.shape) > MAX_TENSOR_DIMS: + model_arch.handle_nd_tensor(key, data) + continue # needs to be added back later + + # get number of parameters (AKA elements) in this tensor + n_params = 1 + for dim_size in data_shape: + n_params *= dim_size + + if old_dtype in (torch.float32, torch.bfloat16): + if n_dims == 1: + # one-dimensional tensors should be kept in F32 + # also speeds up inference due to not dequantizing + data_qtype = gguf.GGMLQuantizationType.F32 + + elif n_params <= QUANTIZATION_THRESHOLD: + # very small tensors + data_qtype = gguf.GGMLQuantizationType.F32 + + elif any(x in key for x in model_arch.keys_hiprec): + # tensors that require max precision + data_qtype = gguf.GGMLQuantizationType.F32 + + if (model_arch.shape_fix # NEVER reshape for models such as flux + and n_dims > 1 # Skip one-dimensional tensors + and n_params >= REARRANGE_THRESHOLD # Only rearrange tensors meeting the size requirement + and (n_params / 256).is_integer() # Rearranging only makes sense if total elements is divisible by 256 + and not (data.shape[-1] / 256).is_integer() # Only need to rearrange if the last dimension is not divisible by 256 + ): + orig_shape = data.shape + data = data.reshape(n_params // 256, 256) + writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) + + try: + data = gguf.quants.quantize(data, data_qtype) + except (AttributeError, gguf.QuantError) as e: + tqdm.write(f"falling back to F16: {e}") + data_qtype = gguf.GGMLQuantizationType.F16 + data = gguf.quants.quantize(data, data_qtype) + + new_name = key # do we need to rename? + + shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" + tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") + + writer.add_tensor(new_name, data, raw_dtype=data_qtype) + +def convert_file(path, dst_path=None, interact=True, overwrite=False): + # load & run model detection logic + state_dict = load_state_dict(path) + model_arch = detect_arch(state_dict) + logging.info(f"* Architecture detected from input: {model_arch.arch}") + + # detect & set dtype for output file + dtypes = [x.dtype for x in state_dict.values()] + dtypes = {x:dtypes.count(x) for x in set(dtypes)} + main_dtype = max(dtypes, key=dtypes.get) + + if main_dtype == torch.bfloat16: + ftype_name = "BF16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 + # elif main_dtype == torch.float32: + # ftype_name = "F32" + # ftype_gguf = None + else: + ftype_name = "F16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_F16 + + if dst_path is None: + dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" + elif "{ftype}" in dst_path: # lcpp logic + dst_path = dst_path.replace("{ftype}", ftype_name) + + if os.path.isfile(dst_path) and not overwrite: + if interact: + input("Output exists enter to continue or ctrl+c to abort!") + else: + raise OSError("Output exists and overwriting is disabled!") + + # handle actual file + writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + if ftype_gguf is not None: + writer.add_file_type(ftype_gguf) + + handle_tensors(writer, state_dict, model_arch) + writer.write_header_to_file(path=dst_path) + writer.write_kv_data_to_file() + writer.write_tensors_to_file(progress=True) + writer.close() + + fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" + if os.path.isfile(fix): + logging.warning(f"\n### Warning! Fix file found at '{fix}'") + logging.warning(" you most likely need to run 'fix_5d_tensors.py' after quantization.") + + return dst_path, model_arch + +if __name__ == "__main__": + args = parse_args() + convert_file(args.src, args.dst) + diff --git a/custom_nodes/ComfyUI-GGUF/tools/fix_5d_tensors.py b/custom_nodes/ComfyUI-GGUF/tools/fix_5d_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..0e61d1c2e5f2572c3d9fa12eba38c84f13689a53 --- /dev/null +++ b/custom_nodes/ComfyUI-GGUF/tools/fix_5d_tensors.py @@ -0,0 +1,82 @@ +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import os +import gguf +import torch +import argparse +from tqdm import tqdm +from safetensors.torch import load_file + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--src", required=True) + parser.add_argument("--dst", required=True) + parser.add_argument("--fix", required=False, help="Defaults to ./fix_5d_tensors_[arch].pt") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + if not os.path.isfile(args.src): + parser.error(f"Invalid source file '{args.src}'") + if not args.overwrite and os.path.exists(args.dst): + parser.error(f"Output exists, use '--overwrite' ({args.dst})") + + return args + +def get_arch_str(reader): + field = reader.get_field("general.architecture") + return str(field.parts[field.data[-1]], encoding="utf-8") + +def get_file_type(reader): + field = reader.get_field("general.file_type") + ft = int(field.parts[field.data[-1]]) + return gguf.LlamaFileType(ft) + +if __name__ == "__main__": + args = get_args() + + # read existing + reader = gguf.GGUFReader(args.src) + arch = get_arch_str(reader) + file_type = get_file_type(reader) + print(f"Detected arch: '{arch}' (ftype: {str(file_type)})") + + # prep fix + if args.fix is None: + args.fix = f"./fix_5d_tensors_{arch}.safetensors" + + if not os.path.isfile(args.fix): + raise OSError(f"No 5D tensor fix file: {args.fix}") + + sd5d = load_file(args.fix) + sd5d = {k:v.numpy() for k,v in sd5d.items()} + print("5D tensors:", sd5d.keys()) + + # prep output + writer = gguf.GGUFWriter(path=None, arch=arch) + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + writer.add_file_type(file_type) + + added = [] + def add_extra_key(writer, key, data): + global added + data_qtype = gguf.GGMLQuantizationType.F32 + data = gguf.quants.quantize(data, data_qtype) + tqdm.write(f"Adding key {key} ({data.shape})") + writer.add_tensor(key, data, raw_dtype=data_qtype) + added.append(key) + + # main loop to add missing 5D tensor(s) + for tensor in tqdm(reader.tensors): + writer.add_tensor(tensor.name, tensor.data, raw_dtype=tensor.tensor_type) + key5d = tensor.name.replace(".bias", ".weight") + if key5d in sd5d.keys(): + add_extra_key(writer, key5d, sd5d[key5d]) + + # brute force for any missed + for key, data in sd5d.items(): + if key not in added: + add_extra_key(writer, key, data) + + writer.write_header_to_file(path=args.dst) + writer.write_kv_data_to_file() + writer.write_tensors_to_file(progress=True) + writer.close() diff --git a/custom_nodes/ComfyUI-GGUF/tools/fix_lines_ending.py b/custom_nodes/ComfyUI-GGUF/tools/fix_lines_ending.py new file mode 100644 index 0000000000000000000000000000000000000000..346e3501fb7682fa3754f175965aa750241fe4ac --- /dev/null +++ b/custom_nodes/ComfyUI-GGUF/tools/fix_lines_ending.py @@ -0,0 +1,31 @@ +import os + +files = ["lcpp.patch", "lcpp_sd3.patch"] + +def has_unix_line_endings(file_path): + try: + with open(file_path, 'rb') as file: + content = file.read() + return b'\r\n' not in content + except Exception as e: + print(f"Error checking '{file_path}': {e}") + return False + +def convert_to_linux_format(file_path): + try: + with open(file_path, 'rb') as file: + content = file.read().replace(b'\r\n', b'\n') + with open(file_path, 'wb') as file: + file.write(content) + print(f"'{file_path}' converted to Linux line endings (LF).") + except Exception as e: + print(f"Error processing '{file_path}': {e}") + +for file in files: + if os.path.exists(file): + if has_unix_line_endings(file): + print(f"'{file}' already has Unix line endings (LF). No conversion needed.") + else: + convert_to_linux_format(file) + else: + print(f"File '{file}' does not exist.") diff --git a/custom_nodes/ComfyUI-GGUF/tools/lcpp.patch b/custom_nodes/ComfyUI-GGUF/tools/lcpp.patch new file mode 100644 index 0000000000000000000000000000000000000000..92396e17b4ecf281265f07603232d23111ee9baa --- /dev/null +++ b/custom_nodes/ComfyUI-GGUF/tools/lcpp.patch @@ -0,0 +1,451 @@ +diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h +index de3c706f..0267c1fa 100644 +--- a/ggml/include/ggml.h ++++ b/ggml/include/ggml.h +@@ -223,7 +223,7 @@ + #define GGML_MAX_OP_PARAMS 64 + + #ifndef GGML_MAX_NAME +-# define GGML_MAX_NAME 64 ++# define GGML_MAX_NAME 128 + #endif + + #define GGML_DEFAULT_N_THREADS 4 +@@ -2449,6 +2449,7 @@ extern "C" { + + // manage tensor info + GGML_API void gguf_add_tensor(struct gguf_context * ctx, const struct ggml_tensor * tensor); ++ GGML_API void gguf_set_tensor_ndim(struct gguf_context * ctx, const char * name, int n_dim); + GGML_API void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type); + GGML_API void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data, size_t size); + +diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c +index b16c462f..6d1568f1 100644 +--- a/ggml/src/ggml.c ++++ b/ggml/src/ggml.c +@@ -22960,6 +22960,14 @@ void gguf_add_tensor( + ctx->header.n_tensors++; + } + ++void gguf_set_tensor_ndim(struct gguf_context * ctx, const char * name, const int n_dim) { ++ const int idx = gguf_find_tensor(ctx, name); ++ if (idx < 0) { ++ GGML_ABORT("tensor not found"); ++ } ++ ctx->infos[idx].n_dims = n_dim; ++} ++ + void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type) { + const int idx = gguf_find_tensor(ctx, name); + if (idx < 0) { +diff --git a/src/llama.cpp b/src/llama.cpp +index 24e1f1f0..25db4c69 100644 +--- a/src/llama.cpp ++++ b/src/llama.cpp +@@ -205,6 +205,17 @@ enum llm_arch { + LLM_ARCH_GRANITE, + LLM_ARCH_GRANITE_MOE, + LLM_ARCH_CHAMELEON, ++ LLM_ARCH_FLUX, ++ LLM_ARCH_SD1, ++ LLM_ARCH_SDXL, ++ LLM_ARCH_SD3, ++ LLM_ARCH_AURA, ++ LLM_ARCH_LTXV, ++ LLM_ARCH_HYVID, ++ LLM_ARCH_WAN, ++ LLM_ARCH_HIDREAM, ++ LLM_ARCH_COSMOS, ++ LLM_ARCH_LUMINA2, + LLM_ARCH_UNKNOWN, + }; + +@@ -258,6 +269,17 @@ static const std::map LLM_ARCH_NAMES = { + { LLM_ARCH_GRANITE, "granite" }, + { LLM_ARCH_GRANITE_MOE, "granitemoe" }, + { LLM_ARCH_CHAMELEON, "chameleon" }, ++ { LLM_ARCH_FLUX, "flux" }, ++ { LLM_ARCH_SD1, "sd1" }, ++ { LLM_ARCH_SDXL, "sdxl" }, ++ { LLM_ARCH_SD3, "sd3" }, ++ { LLM_ARCH_AURA, "aura" }, ++ { LLM_ARCH_LTXV, "ltxv" }, ++ { LLM_ARCH_HYVID, "hyvid" }, ++ { LLM_ARCH_WAN, "wan" }, ++ { LLM_ARCH_HIDREAM, "hidream" }, ++ { LLM_ARCH_COSMOS, "cosmos" }, ++ { LLM_ARCH_LUMINA2, "lumina2" }, + { LLM_ARCH_UNKNOWN, "(unknown)" }, + }; + +@@ -1531,6 +1553,17 @@ static const std::map> LLM_TENSOR_N + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, + }, + }, ++ { LLM_ARCH_FLUX, {}}, ++ { LLM_ARCH_SD1, {}}, ++ { LLM_ARCH_SDXL, {}}, ++ { LLM_ARCH_SD3, {}}, ++ { LLM_ARCH_AURA, {}}, ++ { LLM_ARCH_LTXV, {}}, ++ { LLM_ARCH_HYVID, {}}, ++ { LLM_ARCH_WAN, {}}, ++ { LLM_ARCH_HIDREAM, {}}, ++ { LLM_ARCH_COSMOS, {}}, ++ { LLM_ARCH_LUMINA2, {}}, + { + LLM_ARCH_UNKNOWN, + { +@@ -5403,6 +5436,25 @@ static void llm_load_hparams( + // get general kv + ml.get_key(LLM_KV_GENERAL_NAME, model.name, false); + ++ // Disable LLM metadata for image models ++ switch (model.arch) { ++ case LLM_ARCH_FLUX: ++ case LLM_ARCH_SD1: ++ case LLM_ARCH_SDXL: ++ case LLM_ARCH_SD3: ++ case LLM_ARCH_AURA: ++ case LLM_ARCH_LTXV: ++ case LLM_ARCH_HYVID: ++ case LLM_ARCH_WAN: ++ case LLM_ARCH_HIDREAM: ++ case LLM_ARCH_COSMOS: ++ case LLM_ARCH_LUMINA2: ++ model.ftype = ml.ftype; ++ return; ++ default: ++ break; ++ } ++ + // get hparams kv + ml.get_key(LLM_KV_VOCAB_SIZE, hparams.n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, hparams.n_vocab); + +@@ -18016,6 +18068,134 @@ static void llama_tensor_dequantize_internal( + workers.clear(); + } + ++static ggml_type img_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { ++ // Special function for quantizing image model tensors ++ const std::string name = ggml_get_name(tensor); ++ const llm_arch arch = qs.model.arch; ++ ++ // Sanity check ++ if ( ++ (name.find("model.diffusion_model.") != std::string::npos) || ++ (name.find("first_stage_model.") != std::string::npos) || ++ (name.find("single_transformer_blocks.") != std::string::npos) || ++ (name.find("joint_transformer_blocks.") != std::string::npos) ++ ) { ++ throw std::runtime_error("Invalid input GGUF file. This is not a supported UNET model"); ++ } ++ ++ // Unsupported quant types - exclude all IQ quants for now ++ if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ2_XS || ++ ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M || ++ ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ1_S || ++ ftype == LLAMA_FTYPE_MOSTLY_IQ1_M || ftype == LLAMA_FTYPE_MOSTLY_IQ4_NL || ++ ftype == LLAMA_FTYPE_MOSTLY_IQ4_XS || ftype == LLAMA_FTYPE_MOSTLY_IQ3_S || ++ ftype == LLAMA_FTYPE_MOSTLY_IQ3_M || ftype == LLAMA_FTYPE_MOSTLY_Q4_0_4_4 || ++ ftype == LLAMA_FTYPE_MOSTLY_Q4_0_4_8 || ftype == LLAMA_FTYPE_MOSTLY_Q4_0_8_8) { ++ throw std::runtime_error("Invalid quantization type for image model (Not supported)"); ++ } ++ ++ if ( // Rules for to_v attention ++ (name.find("attn_v.weight") != std::string::npos) || ++ (name.find(".to_v.weight") != std::string::npos) || ++ (name.find(".v.weight") != std::string::npos) || ++ (name.find(".attn.w1v.weight") != std::string::npos) || ++ (name.find(".attn.w2v.weight") != std::string::npos) || ++ (name.find("_attn.v_proj.weight") != std::string::npos) ++ ){ ++ if (ftype == LLAMA_FTYPE_MOSTLY_Q2_K) { ++ new_type = GGML_TYPE_Q3_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M) { ++ new_type = qs.i_attention_wv < 2 ? GGML_TYPE_Q5_K : GGML_TYPE_Q4_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) { ++ new_type = GGML_TYPE_Q5_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M) { ++ new_type = GGML_TYPE_Q6_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_S && qs.i_attention_wv < 4) { ++ new_type = GGML_TYPE_Q5_K; ++ } ++ ++qs.i_attention_wv; ++ } else if ( // Rules for fused qkv attention ++ (name.find("attn_qkv.weight") != std::string::npos) || ++ (name.find("attn.qkv.weight") != std::string::npos) || ++ (name.find("attention.qkv.weight") != std::string::npos) ++ ) { ++ if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M || ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) { ++ new_type = GGML_TYPE_Q4_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M) { ++ new_type = GGML_TYPE_Q5_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M) { ++ new_type = GGML_TYPE_Q6_K; ++ } ++ } else if ( // Rules for ffn ++ (name.find("ffn_down") != std::string::npos) || ++ ((name.find("experts.") != std::string::npos) && (name.find(".w2.weight") != std::string::npos)) || ++ (name.find(".ffn.2.weight") != std::string::npos) || // is this even the right way around? ++ (name.find(".ff.net.2.weight") != std::string::npos) || ++ (name.find(".mlp.layer2.weight") != std::string::npos) || ++ (name.find(".adaln_modulation_mlp.2.weight") != std::string::npos) || ++ (name.find(".feed_forward.w2.weight") != std::string::npos) ++ ) { ++ // TODO: add back `layer_info` with some model specific logic + logic further down ++ if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_M) { ++ new_type = GGML_TYPE_Q4_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q3_K_L) { ++ new_type = GGML_TYPE_Q5_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_S) { ++ new_type = GGML_TYPE_Q5_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_K_M) { ++ new_type = GGML_TYPE_Q6_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q5_K_M) { ++ new_type = GGML_TYPE_Q6_K; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q4_0) { ++ new_type = GGML_TYPE_Q4_1; ++ } ++ else if (ftype == LLAMA_FTYPE_MOSTLY_Q5_0) { ++ new_type = GGML_TYPE_Q5_1; ++ } ++ ++qs.i_ffn_down; ++ } ++ ++ // Sanity check for row shape ++ bool convert_incompatible_tensor = false; ++ if (new_type == GGML_TYPE_Q2_K || new_type == GGML_TYPE_Q3_K || new_type == GGML_TYPE_Q4_K || ++ new_type == GGML_TYPE_Q5_K || new_type == GGML_TYPE_Q6_K) { ++ int nx = tensor->ne[0]; ++ int ny = tensor->ne[1]; ++ if (nx % QK_K != 0) { ++ LLAMA_LOG_WARN("\n\n%s : tensor cols %d x %d are not divisible by %d, required for %s", __func__, nx, ny, QK_K, ggml_type_name(new_type)); ++ convert_incompatible_tensor = true; ++ } else { ++ ++qs.n_k_quantized; ++ } ++ } ++ if (convert_incompatible_tensor) { ++ // TODO: Possibly reenable this in the future ++ // switch (new_type) { ++ // case GGML_TYPE_Q2_K: ++ // case GGML_TYPE_Q3_K: ++ // case GGML_TYPE_Q4_K: new_type = GGML_TYPE_Q5_0; break; ++ // case GGML_TYPE_Q5_K: new_type = GGML_TYPE_Q5_1; break; ++ // case GGML_TYPE_Q6_K: new_type = GGML_TYPE_Q8_0; break; ++ // default: throw std::runtime_error("\nUnsupported tensor size encountered\n"); ++ // } ++ new_type = GGML_TYPE_F16; ++ LLAMA_LOG_WARN(" - using fallback quantization %s\n", ggml_type_name(new_type)); ++ ++qs.n_fallback; ++ } ++ return new_type; ++} ++ + static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type new_type, const ggml_tensor * tensor, llama_ftype ftype) { + const std::string name = ggml_get_name(tensor); + +@@ -18513,7 +18693,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s + if (llama_model_has_encoder(&model)) { + n_attn_layer *= 3; + } +- GGML_ASSERT((qs.n_attention_wv == n_attn_layer) && "n_attention_wv is unexpected"); ++ if (model.arch != LLM_ARCH_HYVID) { // TODO: Check why this fails ++ GGML_ASSERT((qs.n_attention_wv == n_attn_layer) && "n_attention_wv is unexpected"); ++ } + } + + size_t total_size_org = 0; +@@ -18547,6 +18729,51 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s + ctx_outs[i_split] = gguf_init_empty(); + } + gguf_add_tensor(ctx_outs[i_split], tensor); ++ // SD3 pos_embed needs special fix as first dim is 1, which gets truncated here ++ if (model.arch == LLM_ARCH_SD3) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "pos_embed" && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "pos_embed", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting pos_embed shape for SD3: [key:%s]\n", __func__, tensor->name); ++ } ++ } ++ // same goes for auraflow ++ if (model.arch == LLM_ARCH_AURA) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "positional_encoding" && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "positional_encoding", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting positional_encoding shape for AuraFlow: [key:%s]\n", __func__, tensor->name); ++ } ++ if (name == "register_tokens" && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], "register_tokens", n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting register_tokens shape for AuraFlow: [key:%s]\n", __func__, tensor->name); ++ } ++ } ++ // conv3d fails due to max dims - unsure what to do here as we never even reach this check ++ if (model.arch == LLM_ARCH_HYVID) { ++ const std::string name = ggml_get_name(tensor); ++ if (name == "img_in.proj.weight" && tensor->ne[5] != 1 ) { ++ throw std::runtime_error("img_in.proj.weight size failed for HyVid"); ++ } ++ } ++ // All the modulation layers also have dim1, and I think conv3d fails here too but we segfaul way before that... ++ if (model.arch == LLM_ARCH_WAN) { ++ const std::string name = ggml_get_name(tensor); ++ if (name.find(".modulation") != std::string::npos && tensor->ne[2] == 1) { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting shape for Wan: [key:%s]\n", __func__, tensor->name); ++ } ++ // FLF2V model only ++ if (name == "img_emb.emb_pos") { ++ const int n_dim = 3; ++ gguf_set_tensor_ndim(ctx_outs[i_split], tensor->name, n_dim); ++ LLAMA_LOG_INFO("\n%s: Correcting shape for Wan FLF2V: [key:%s]\n", __func__, tensor->name); ++ } ++ } + } + + // Set split info if needed +@@ -18647,6 +18874,110 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s + // do not quantize relative position bias (T5) + quantize &= name.find("attn_rel_b.weight") == std::string::npos; + ++ // rules for image models ++ bool image_model = false; ++ if (model.arch == LLM_ARCH_FLUX) { ++ image_model = true; ++ quantize &= name.find("txt_in.") == std::string::npos; ++ quantize &= name.find("img_in.") == std::string::npos; ++ quantize &= name.find("time_in.") == std::string::npos; ++ quantize &= name.find("vector_in.") == std::string::npos; ++ quantize &= name.find("guidance_in.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ } ++ if (model.arch == LLM_ARCH_SD1 || model.arch == LLM_ARCH_SDXL) { ++ image_model = true; ++ quantize &= name.find("class_embedding.") == std::string::npos; ++ quantize &= name.find("time_embedding.") == std::string::npos; ++ quantize &= name.find("add_embedding.") == std::string::npos; ++ quantize &= name.find("time_embed.") == std::string::npos; ++ quantize &= name.find("label_emb.") == std::string::npos; ++ quantize &= name.find("conv_in.") == std::string::npos; ++ quantize &= name.find("conv_out.") == std::string::npos; ++ quantize &= name != "input_blocks.0.0.weight"; ++ quantize &= name != "out.2.weight"; ++ } ++ if (model.arch == LLM_ARCH_SD3) { ++ image_model = true; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ quantize &= name.find("time_text_embed.") == std::string::npos; ++ quantize &= name.find("context_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("y_embedder.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name != "proj_out.weight"; ++ quantize &= name != "pos_embed"; ++ } ++ if (model.arch == LLM_ARCH_AURA) { ++ image_model = true; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("init_x_linear.") == std::string::npos; ++ quantize &= name != "modF.1.weight"; ++ quantize &= name != "cond_seq_linear.weight"; ++ quantize &= name != "final_linear.weight"; ++ quantize &= name != "final_linear.weight"; ++ quantize &= name != "positional_encoding"; ++ quantize &= name != "register_tokens"; ++ } ++ if (model.arch == LLM_ARCH_LTXV) { ++ image_model = true; ++ quantize &= name.find("adaln_single.") == std::string::npos; ++ quantize &= name.find("caption_projection.") == std::string::npos; ++ quantize &= name.find("patchify_proj.") == std::string::npos; ++ quantize &= name.find("proj_out.") == std::string::npos; ++ quantize &= name.find("scale_shift_table") == std::string::npos; // last block too ++ } ++ if (model.arch == LLM_ARCH_HYVID) { ++ image_model = true; ++ quantize &= name.find("txt_in.") == std::string::npos; ++ quantize &= name.find("img_in.") == std::string::npos; ++ quantize &= name.find("time_in.") == std::string::npos; ++ quantize &= name.find("vector_in.") == std::string::npos; ++ quantize &= name.find("guidance_in.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ } ++ if (model.arch == LLM_ARCH_WAN) { ++ image_model = true; ++ quantize &= name.find("modulation.") == std::string::npos; ++ quantize &= name.find("patch_embedding.") == std::string::npos; ++ quantize &= name.find("text_embedding.") == std::string::npos; ++ quantize &= name.find("time_projection.") == std::string::npos; ++ quantize &= name.find("time_embedding.") == std::string::npos; ++ quantize &= name.find("img_emb.") == std::string::npos; ++ quantize &= name.find("head.") == std::string::npos; ++ } ++ if (model.arch == LLM_ARCH_HIDREAM) { ++ image_model = true; ++ quantize &= name.find("p_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ quantize &= name.find(".ff_i.gate.weight") == std::string::npos; ++ quantize &= name.find("caption_projection.") == std::string::npos; ++ } ++ if (model.arch == LLM_ARCH_COSMOS) { ++ image_model = true; ++ quantize &= name.find("p_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("t_embedding_norm.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name.find("pos_embedder.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ } ++ if (model.arch == LLM_ARCH_LUMINA2) { ++ image_model = true; ++ quantize &= name.find("t_embedder.") == std::string::npos; ++ quantize &= name.find("x_embedder.") == std::string::npos; ++ quantize &= name.find("final_layer.") == std::string::npos; ++ quantize &= name.find("cap_embedder.") == std::string::npos; ++ quantize &= name.find("context_refiner.") == std::string::npos; ++ quantize &= name.find("noise_refiner.") == std::string::npos; ++ } ++ // ignore 3D/4D tensors for image models as the code was never meant to handle these ++ if (image_model) { ++ quantize &= ggml_n_dims(tensor) == 2; ++ } ++ + enum ggml_type new_type; + void * new_data; + size_t new_size; +@@ -18655,6 +18986,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s + new_type = default_type; + + // get more optimal quantization type based on the tensor shape, layer, etc. ++ if (image_model) { ++ new_type = img_tensor_get_type(qs, new_type, tensor, ftype); ++ } else { + if (!params->pure && ggml_is_quantized(default_type)) { + new_type = llama_tensor_get_type(qs, new_type, tensor, ftype); + } +@@ -18664,6 +18998,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s + if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) { + new_type = params->output_tensor_type; + } ++ } + + // If we've decided to quantize to the same type the tensor is already + // in then there's nothing to do. diff --git a/custom_nodes/ComfyUI-GGUF/tools/read_tensors.py b/custom_nodes/ComfyUI-GGUF/tools/read_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..1bdff028a787c09b38e5616ef75a2f070c672445 --- /dev/null +++ b/custom_nodes/ComfyUI-GGUF/tools/read_tensors.py @@ -0,0 +1,21 @@ +#!/usr/bin/python3 +import os +import sys +import gguf + +def read_tensors(path): + reader = gguf.GGUFReader(path) + for tensor in reader.tensors: + if tensor.tensor_type == gguf.GGMLQuantizationType.F32: + continue + print(f"{str(tensor.tensor_type):32}: {tensor.name}") + +try: + path = sys.argv[1] + assert os.path.isfile(path), "Invalid path" + print(f"input: {path}") +except Exception as e: + input(f"failed: {e}") +else: + read_tensors(path) + input() diff --git a/custom_nodes/ComfyUI-HuggingFace/.github/workflows/publish.yml b/custom_nodes/ComfyUI-HuggingFace/.github/workflows/publish.yml new file mode 100644 index 0000000000000000000000000000000000000000..8e92e1f9caf921f4144100c4707b8879a50a6b6c --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/.github/workflows/publish.yml @@ -0,0 +1,28 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + - master + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'MoonGoblinDev' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + submodules: true + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/custom_nodes/ComfyUI-HuggingFace/api/__init__.py b/custom_nodes/ComfyUI-HuggingFace/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/custom_nodes/ComfyUI-HuggingFace/api/huggingface.py b/custom_nodes/ComfyUI-HuggingFace/api/huggingface.py new file mode 100644 index 0000000000000000000000000000000000000000..fe63f54ebd3b5a5f8a552e3f0b02d75ce7f99117 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/api/huggingface.py @@ -0,0 +1,219 @@ +# ================================================ +# File: api/huggingface.py +# ================================================ +import requests +import json +from typing import List, Optional, Dict, Any, Union + +# Try to import huggingface_hub for better downloads +try: + from huggingface_hub import hf_hub_download, snapshot_download + HF_HUB_AVAILABLE = True +except ImportError: + HF_HUB_AVAILABLE = False + print("[HuggingFace API] huggingface_hub not available, falling back to manual downloads") + +# Try to use huggingface_hub CLI as fallback +import subprocess +import sys + +class HuggingFaceAPI: + """Simple wrapper for interacting with the HuggingFace API.""" + BASE_URL = "https://huggingface.co/api" + + def __init__(self, api_key: Optional[str] = None): + self.api_key = api_key + self.base_headers = {'Content-Type': 'application/json'} + if api_key: + self.base_headers["Authorization"] = f"Bearer {api_key}" + print("[HuggingFace API] Using HF token for private repositories.") + else: + print("[HuggingFace API] No HF token provided. Only public repositories accessible.") + + def _get_request_headers(self, method: str, has_json_data: bool) -> Dict[str, str]: + """Returns headers for a specific request.""" + headers = self.base_headers.copy() + # Don't send content-type for GET/HEAD if no json_data + if method.upper() in ["GET", "HEAD"] and not has_json_data: + headers.pop('Content-Type', None) + return headers + + def _request(self, method: str, endpoint: str, params: Optional[Dict] = None, + json_data: Optional[Dict] = None, stream: bool = False, + allow_redirects: bool = True, timeout: int = 30) -> Union[Dict[str, Any], requests.Response, None]: + """Makes a request to the HuggingFace API and handles basic errors.""" + url = f"{self.BASE_URL}/{endpoint.lstrip('/')}" + request_headers = self._get_request_headers(method, json_data is not None) + + try: + response = requests.request( + method, + url, + headers=request_headers, + params=params, + json=json_data, + stream=stream, + allow_redirects=allow_redirects, + timeout=timeout + ) + response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx) + + if stream: + return response # Return the response object for streaming + + # Handle No Content response (e.g., 204) + if response.status_code == 204 or not response.content: + return None + + return response.json() + + except requests.exceptions.HTTPError as http_err: + error_detail = None + status_code = http_err.response.status_code + try: + error_detail = http_err.response.json() + except json.JSONDecodeError: + error_detail = http_err.response.text[:200] # First 200 chars + print(f"HuggingFace API HTTP Error ({method} {url}): Status {status_code}, Response: {error_detail}") + # Return a structured error dictionary + return {"error": f"HTTP Error: {status_code}", "details": error_detail, "status_code": status_code} + + except requests.exceptions.RequestException as req_err: + print(f"HuggingFace API Request Error ({method} {url}): {req_err}") + return {"error": str(req_err), "details": None, "status_code": None} + + except json.JSONDecodeError as json_err: + print(f"HuggingFace API Error: Failed to decode JSON response from {url}: {json_err}") + # Include response text if possible and not streaming + response_text = response.text[:200] if not stream and hasattr(response, 'text') else "N/A" + return {"error": "Invalid JSON response", "details": response_text, "status_code": response.status_code if hasattr(response, 'status_code') else None} + + def search_models(self, query: str, limit: int = 20) -> Optional[Dict[str, Any]]: + """Searches for models on HuggingFace. (GET /models)""" + endpoint = "/models" + params = { + "search": query, + "limit": limit + } + result = self._request("GET", endpoint, params=params) + if isinstance(result, dict) and "error" in result: + return result + return result + + def search_models_meili(self, query: str = None, types: Optional[List[str]] = None, + base_models: Optional[List[str]] = None, + sort: str = 'Most Downloaded', limit: int = 20, page: int = 1, + nsfw: Optional[bool] = None) -> Optional[Dict[str, Any]]: + """Searches models using HuggingFace's Meilisearch endpoint.""" + meili_url = "https://huggingface.co/multi-search" + headers = {'Content-Type': 'application/json'} + if self.api_key: + headers['Authorization'] = f'Bearer {self.api_key}' + + # Build search query + search_query = { + "q": query or "", + "limit": limit, + "offset": (page - 1) * limit + } + + # Add filters + filters = [] + + # Type filters + if types: + type_filter = {"type": types} + filters.append(type_filter) + + # Base model filters + if base_models: + base_filter = {"base_model": base_models} + filters.append(base_filter) + + # NSFW filter + if nsfw is not None: + nsfw_filter = {"nsfw": nsfw} + filters.append(nsfw_filter) + + if filters: + search_query["filters"] = filters + + # Add sorting + sort_mapping = { + "Relevancy": "id:desc", + "Most Downloaded": "metrics.downloadCount:desc", + "Highest Rated": "metrics.thumbsUpCount:desc", + "Most Liked": "metrics.favoriteCount:desc", + "Most Discussed": "metrics.commentCount:desc", + "Most Collected": "metrics.collectedCount:desc", + "Most Buzz": "metrics.tippedAmountCount:desc", + "Newest": "createdAt:desc", + } + + if sort in sort_mapping: + search_query["sort"] = [sort_mapping[sort]] + + try: + response = requests.post(meili_url, json=search_query, headers=headers, timeout=30) + response.raise_for_status() + return response.json() + except requests.exceptions.RequestException as e: + print(f"HuggingFace Meili API Error: {e}") + return {"error": str(e), "status_code": getattr(e.response, 'status_code', None) if hasattr(e, 'response') else None} + + def get_model_files(self, model_id: str) -> Optional[Dict[str, Any]]: + """Gets files for a specific HuggingFace model. Uses huggingface_hub instead of API.""" + # Skip API calls, let huggingface_hub handle everything + print(f"[HuggingFace API] Skipping file listing, using huggingface_hub auto-detect for {model_id}") + return {"auto_detect": True} + + def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]: + """Gets information about a model by its ID. Uses huggingface_hub instead of API.""" + # Skip API calls, let huggingface_hub handle everything + print(f"[HuggingFace API] Skipping model info, using huggingface_hub auto-detect for {model_id}") + return {"id": model_id, "name": model_id.split('/')[-1]} + + def get_model_version_info(self, version_id: str) -> Optional[Dict[str, Any]]: + """Gets version information - not applicable for HuggingFace, returns empty dict""" + # HuggingFace doesn't have version IDs like Civitai + # This method is kept for compatibility but returns empty + return {} + + def download_file(self, model_id: str, filename: str, local_dir: str = None) -> Optional[Union[requests.Response, str]]: + """Downloads a specific file from HuggingFace. Uses only huggingface_hub.""" + if not HF_HUB_AVAILABLE: + print("[HuggingFace API] huggingface_hub not available") + return None + + if not local_dir: + print("[HuggingFace API] local_dir not specified") + return None + + try: + print(f"[HuggingFace API] Using huggingface_hub for download: {model_id}/{filename}") + + if filename is None: + # Download entire repo using snapshot_download + print(f"[HuggingFace API] Downloading entire repo {model_id}") + result = snapshot_download( + repo_id=model_id, + local_dir=local_dir, + token=self.api_key + ) + print(f"[HuggingFace API] snapshot_download success: {result}") + return result + else: + # Download specific file using hf_hub_download + result = hf_hub_download( + repo_id=model_id, + filename=filename, + local_dir=local_dir, + resume_download=True, + token=self.api_key + ) + print(f"[HuggingFace API] hf_hub_download success: {result}") + return result + + except Exception as e: + print(f"[HuggingFace API] huggingface_hub download failed: {e}") + return None diff --git a/custom_nodes/ComfyUI-HuggingFace/downloader/__init__.py b/custom_nodes/ComfyUI-HuggingFace/downloader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/custom_nodes/ComfyUI-HuggingFace/downloader/chunk_downloader.py b/custom_nodes/ComfyUI-HuggingFace/downloader/chunk_downloader.py new file mode 100644 index 0000000000000000000000000000000000000000..44764b1df80b53345e648319bddc8d5b37468476 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/downloader/chunk_downloader.py @@ -0,0 +1,574 @@ +# ================================================ +# File: downloader/chunk_downloader.py +# Multi-connection somehow still not working +# ================================================ + +import requests +import threading +import time +import shutil +from pathlib import Path +import os +from typing import Optional, Dict, Tuple, Union, TYPE_CHECKING + +# Import manager type hint without circular dependency during type checking +if TYPE_CHECKING: + from .manager import DownloadManager + +# Import config values +from ..config import DEFAULT_CHUNK_SIZE, DOWNLOAD_TIMEOUT, HEAD_REQUEST_TIMEOUT + +class ChunkDownloader: + """Handles downloading files in chunks using multiple connections or fallback.""" + # Constants + STATUS_UPDATE_INTERVAL = 0.5 + HEAD_REQUEST_TIMEOUT = HEAD_REQUEST_TIMEOUT + DOWNLOAD_TIMEOUT = DOWNLOAD_TIMEOUT + MIN_SIZE_FOR_MULTI_MB = 100 # Minimum file size for multi-connection download + + def __init__(self, url: str, output_path: str, num_connections: int = 4, + chunk_size: int = DEFAULT_CHUNK_SIZE, manager: 'DownloadManager' = None, + download_id: str = None, api_key: Optional[str] = None, + known_size: Optional[int] = None): + # URLs + self.initial_url = url + self.url = url + + # Paths + self.output_path = Path(output_path) + self.temp_dir = self.output_path.parent / f".{self.output_path.name}.parts_{download_id or int(time.time())}" + + # Download configuration + self.num_connections = num_connections # Use the provided number of connections + self.chunk_size = chunk_size + self.manager = manager + self.download_id = download_id + self.api_key = api_key + self.known_size = known_size if known_size and known_size > 0 else None + + # Download state + self.total_size = self.known_size or 0 + self.downloaded = 0 + self.connection_type = "N/A" + self.error = None + + # Thread management + self.threads = [] + self.lock = threading.Lock() + self.cancel_event = threading.Event() + self.part_files = [] + + # Performance tracking + self._start_time = 0 + self._last_update_time = 0 + self._last_downloaded_bytes = 0 + self._speed = 0 + + def _get_request_headers(self, add_range: Optional[str] = None) -> Dict[str, str]: + """Constructs request headers with optional auth and range.""" + headers = {} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + if add_range: + headers['Range'] = add_range + return headers + + @property + def is_cancelled(self) -> bool: + """Check if download has been cancelled.""" + return self.cancel_event.is_set() + + def cancel(self): + """Signal the download to cancel.""" + if not self.is_cancelled: + print(f"[Downloader {self.download_id or 'N/A'}] Cancellation requested by user.") + self.cancel_event.set() + self.error = "Download cancelled by user" + if self.manager and self.download_id: + self.manager._update_download_status(self.download_id, status="cancelled", error=self.error) + + def _cleanup_temp(self, success: bool): + """Remove temporary directory and potentially the output file.""" + # Clean up temp directory + if self.temp_dir.exists(): + try: + shutil.rmtree(self.temp_dir) + except Exception as e: + print(f"[Downloader {self.download_id}] Warning: Could not remove temp directory {self.temp_dir}: {e}") + + # Remove output file if download failed + if not success and self.output_path.exists(): + try: + self.output_path.unlink() + print(f"[Downloader {self.download_id}] Removed incomplete/failed output file: {self.output_path}") + except Exception as e: + print(f"[Downloader {self.download_id}] Warning: Could not remove incomplete output file {self.output_path}: {e}") + + def _get_range_support_and_url(self) -> Tuple[str, bool]: + """Check for range support and get final URL after redirects.""" + final_url = self.initial_url + supports_ranges = False + + try: + request_headers = self._get_request_headers() + + print(f"[Downloader {self.download_id}] Checking range support/redirects for: {self.initial_url} (Timeout: {self.HEAD_REQUEST_TIMEOUT}s)") + response = requests.head( + self.initial_url, + allow_redirects=True, + timeout=self.HEAD_REQUEST_TIMEOUT, + headers=request_headers + ) + response.raise_for_status() + + # Update URL after redirects + final_url = response.url + self.url = final_url + + # Check range support + accept_ranges = response.headers.get('accept-ranges', 'none').lower() + supports_ranges = accept_ranges == 'bytes' + + # Get file size if not already known + if self.total_size <= 0: + head_size = int(response.headers.get('Content-Length', 0)) + if head_size > 0: + self.total_size = head_size + print(f"[Downloader {self.download_id}] Got file size from HEAD: {self.total_size} bytes") + + print(f"[Downloader {self.download_id}] HEAD Check OK - Final URL: {final_url}, Range Support: {supports_ranges}") + return final_url, supports_ranges + + except requests.exceptions.Timeout: + print(f"[Downloader {self.download_id}] Warning: HEAD request timed out. Proceeding with Single connection.") + return self.initial_url, False + + except requests.exceptions.RequestException as e: + http_error_details = "" + if hasattr(e, 'response') and e.response is not None: + status_code = e.response.status_code + http_error_details = f" (Status Code: {status_code})" + print(f"[Downloader {self.download_id}] Warning: HEAD request failed{http_error_details}. Proceeding with Single connection.") + return self.initial_url, False + + except Exception as e: + print(f"[Downloader {self.download_id}] Warning: Unexpected error during HEAD request: {e}. Proceeding with Single connection.") + return self.initial_url, False + + def _update_progress(self, chunk_len: int): + """Thread-safe update of download progress and speed calculation.""" + with self.lock: + self.downloaded += chunk_len + current_time = time.monotonic() + time_diff = current_time - self._last_update_time + + # Update speed and notify manager periodically + if time_diff >= self.STATUS_UPDATE_INTERVAL or self.downloaded == self.total_size: + progress = min((self.downloaded / self.total_size) * 100, 100.0) if self.total_size > 0 else 0 + + # Calculate speed + if time_diff > 0: + bytes_diff = self.downloaded - self._last_downloaded_bytes + self._speed = bytes_diff / time_diff + + self._last_update_time = current_time + self._last_downloaded_bytes = self.downloaded + + if self.manager and self.download_id: + self.manager._update_download_status( + self.download_id, + progress=progress, + speed=self._speed, + status="downloading" + ) + + def download_segment(self, segment_index: int, start_byte: int, end_byte: int): + """Downloads a specific segment of the file.""" + part_file_path = self.temp_dir / f"part_{segment_index}" + request_headers = self._get_request_headers(add_range=f'bytes={start_byte}-{end_byte}') + retries = 3 + + for current_try in range(retries): + if self.is_cancelled: + print(f"[Downloader {self.download_id}] Segment {segment_index} cancelled before request (Try {current_try+1}).") + # Ensure error is set if not already + if not self.error: self.error = "Cancelled during segment download" + return + + response = None + try: + response = requests.get(self.url, headers=request_headers, stream=True, timeout=self.DOWNLOAD_TIMEOUT) + response.raise_for_status() + + bytes_written_this_segment = 0 + with open(part_file_path, 'wb') as f: + for chunk in response.iter_content(self.chunk_size): + if self.is_cancelled: + print(f"[Downloader {self.download_id}] Segment {segment_index} cancelled mid-stream.") + # Ensure error is set if not already + if not self.error: self.error = "Cancelled during segment download" + + + if chunk: + bytes_written = f.write(chunk) + bytes_written_this_segment += bytes_written + self._update_progress(bytes_written) + + # Verify segment size + expected_size = (end_byte - start_byte) + 1 + if bytes_written_this_segment != expected_size: + if response: + response.close() + raise ValueError(f"Size mismatch. Expected {expected_size}, got {bytes_written_this_segment}") + + return # Success + + + except (requests.exceptions.RequestException, ValueError) as e: + # Handle HTTP status codes + status_code = None + error_msg_detail = f"{e}" + + if isinstance(e, requests.exceptions.RequestException) and hasattr(e, 'response') and e.response is not None: + status_code = e.response.status_code + if status_code == 401: + error_msg_detail += " (Unauthorized)" + elif status_code == 403: + error_msg_detail += " (Forbidden)" + elif status_code == 416: + error_msg_detail += " (Range Not Satisfiable)" + self.error = f"Segment {segment_index} failed: {error_msg_detail}" + self.cancel() + return + + print(f"[Downloader {self.download_id}] Warning: Segment {segment_index} failed (Try {current_try+1}/{retries}): {error_msg_detail}") + + if current_try >= retries - 1: # Last attempt failed + self.error = f"Segment {segment_index} failed after {retries} attempts: {error_msg_detail}" + self.cancel() + return + + # Exponential backoff before retry + time.sleep(min(2 ** current_try, 10)) + + except Exception as e: + self.error = f"Segment {segment_index} critical error: {e}" + print(f"[Downloader {self.download_id}] Error: {self.error}") + self.cancel() + return + + finally: + if response: + response.close() + + def merge_parts(self) -> bool: + """Merges all downloaded part files into the final output file.""" + print(f"[Downloader {self.download_id}] Merging {len(self.part_files)} parts for {self.output_path.name}...") + + # Check if we have parts to merge + if not self.part_files: + if self.is_cancelled: + self.error = self.error or "Cancelled before any parts downloaded." + elif not self.error: + self.error = "No part files were created to merge." + print(f"[Downloader {self.download_id}] Error during merge: {self.error}") + return False + + try: + # Sort part files numerically + sorted_part_files = sorted(self.part_files, key=lambda p: int(p.name.split('_')[-1])) + + with open(self.output_path, 'wb') as outfile: + for part_file in sorted_part_files: + # Check if part exists + if not part_file.exists(): + if self.is_cancelled: + self.error = self.error or "Cancelled during download, a part file is missing." + elif not self.error: + self.error = f"Merge failed, required part file is missing: {part_file}." + print(f"[Downloader {self.download_id}] Warning: Aborting merge. Missing part: {part_file.name}.") + return False + + # Copy part data to output file + try: + with open(part_file, 'rb') as infile: + shutil.copyfileobj(infile, outfile, length=1024*1024*2) # 2MB buffer + except Exception as copy_e: + self.error = f"Error copying data from part {part_file.name} during merge: {copy_e}" + print(f"[Downloader {self.download_id}] Error: {self.error}") + return False + + print(f"[Downloader {self.download_id}] Merging complete.") + + # Verify final file size + final_size = self.output_path.stat().st_size + if self.total_size > 0 and abs(final_size - self.total_size) > 1: + self.error = f"Merged size ({final_size}) differs significantly from expected ({self.total_size}). File may be corrupt." + print(f"[Downloader {self.download_id}] Error: {self.error}") + return False + elif self.total_size > 0 and final_size != self.total_size: + print(f"[Downloader {self.download_id}] Warning: Final merged size ({final_size}) differs slightly from expected ({self.total_size}).") + + return True + + except Exception as e: + self.error = f"Failed to merge parts: {e}" + print(f"[Downloader {self.download_id}] Error: {self.error}") + return False + + def fallback_download(self) -> bool: + """Fallback to standard single-connection download.""" + self.connection_type = "Single" + if self.manager and self.download_id: + self.manager._update_download_status(self.download_id, connection_type=self.connection_type, status="downloading") + + print(f"[Downloader {self.download_id}] Using standard single-connection download for {self.output_path.name}...") + + self._start_time = self._start_time or time.monotonic() + self._last_update_time = self._start_time + self._last_downloaded_bytes = 0 + self.downloaded = 0 + + response = None + + try: + request_headers = self._get_request_headers() + response = requests.get(self.url, stream=True, timeout=self.DOWNLOAD_TIMEOUT, + allow_redirects=True, headers=request_headers) + response.raise_for_status() + + # Update URL after potential redirects + final_get_url = response.url + if final_get_url != self.url: + print(f"[Downloader {self.download_id}] URL redirected during GET to: {final_get_url}") + self.url = final_get_url + + # Get/confirm file size + if self.total_size <= 0: + get_size = int(response.headers.get('Content-Length', 0)) + if get_size > 0: + self.total_size = get_size + print(f"[Downloader {self.download_id}] Obtained file size via fallback GET: {self.total_size}") + else: + print(f"[Downloader {self.download_id}] Warning: File size unknown. Progress may be inaccurate.") + + # Ensure output directory exists + self.output_path.parent.mkdir(parents=True, exist_ok=True) + + with open(self.output_path, 'wb') as f: + for chunk in response.iter_content(self.chunk_size): + if self.is_cancelled: + print(f"[Downloader {self.download_id}] Fallback download cancelled.") + return False + if chunk: + bytes_written = f.write(chunk) + self._update_progress(bytes_written) + + + + # Verify download size if known + if self.total_size > 0 and self.downloaded != self.total_size and not self.error: + print(f"[Downloader {self.download_id}] Warning: Size mismatch. Expected {self.total_size}, got {self.downloaded}.") + + print(f"[Downloader {self.download_id}] Fallback download completed.") + return not self.error + + except requests.exceptions.RequestException as e: + error_msg_detail = f"{e}" + if hasattr(e, 'response') and e.response is not None: + status_code = e.response.status_code + if status_code == 401: + error_msg_detail += " (Unauthorized - Check HF token?)" + elif status_code == 403: + error_msg_detail += " (Forbidden - Permissions Issue?)" + + if not self.error: + self.error = f"Fallback download failed: {error_msg_detail}" + print(f"[Downloader {self.download_id}] Error during fallback download: {self.error}") + return False + + except Exception as e: + if not self.error: + self.error = f"Fallback download failed: {e}" + print(f"[Downloader {self.download_id}] Error during fallback download: {self.error}") + return False + + finally: + if response: + response.close() + + def download(self) -> bool: + """Main download method that chooses between multi-connection or fallback approach.""" + self._start_time = time.monotonic() + self.downloaded = 0 + self.error = None + self.threads = [] + self.part_files = [] + success = False + + # Clean up any existing temp directory + if self.temp_dir.exists(): + print(f"[Downloader {self.download_id}] Warning: Removing leftover temp directory: {self.temp_dir}") + self._cleanup_temp(success=False) + + # Check range support and get final URL + final_url, supports_ranges = self._get_range_support_and_url() + + # Decide on download strategy + use_multi_connection = False + if supports_ranges and self.num_connections > 1 and self.total_size > 0: + if self.total_size > self.MIN_SIZE_FOR_MULTI_MB * 1024 * 1024: + use_multi_connection = True + else: + print(f"[Downloader {self.download_id}] File size ({self.total_size / (1024*1024):.2f} MB) below threshold for multi-connection.") + + expected_final_size = self.total_size + + try: + if use_multi_connection: + # Multi-connection download approach + success = self._do_multi_connection_download() + else: + # Single connection fallback + reason = "Range requests not supported" if not supports_ranges else \ + "Single connection requested" if self.num_connections <= 1 else \ + "File size unknown or too small" + print(f"[Downloader {self.download_id}] ({reason}). Using fallback single-connection download.") + success = self.fallback_download() + expected_final_size = self.total_size + + if not success and not self.error: + self.error = "Single connection download failed." + + except KeyboardInterrupt: + print(f"[Downloader {self.download_id}] Interrupted! Signalling cancellation.") + self.cancel() + if not self.error: + self.error = "Download interrupted by user." + success = False + + except Exception as e: + import traceback + print(f"--- Critical Error in Download {self.download_id} ('{self.output_path.name}') ---") + traceback.print_exc() + print("--- End Error ---") + + if not self.error: + self.error = f"Unexpected download error: {str(e)}" + + success = False + if not self.is_cancelled: + self.cancel() + + finally: + # Cleanup and final status update + self._cleanup_temp(success=success and not self.is_cancelled and not self.error) + + if self.manager and self.download_id: + final_status = "completed" if success else ("cancelled" if self.is_cancelled else "failed") + final_progress = 100.0 if success else ((self.downloaded / self.total_size * 100) if self.total_size > 0 else 0) + + self.manager._update_download_status( + self.download_id, + status=final_status, + progress=min(100.0, final_progress), + speed=0, + error=self.error, + connection_type=self.connection_type + ) + + return success and not self.error and not self.is_cancelled + + def _do_multi_connection_download(self) -> bool: + """Handle multi-connection download process.""" + self.connection_type = f"Multi ({self.num_connections})" + if self.manager and self.download_id: + self.manager._update_download_status(self.download_id, connection_type=self.connection_type, status="downloading") + + print(f"[Downloader {self.download_id}] Starting multi-connection download for {self.output_path.name} " + f"({self.total_size / (1024 * 1024):.2f} MB) using {self.num_connections} connections.") + + # Create temp directory + try: + if self.temp_dir.exists(): + shutil.rmtree(self.temp_dir) + self.temp_dir.mkdir(parents=True) + except Exception as e: + self.error = f"Failed to create temp directory: {e}" + print(f"[Downloader {self.download_id}] Error: {self.error}") + return False + + # Calculate segments + segment_size = self.total_size // self.num_connections + + # Handle small files with many connections + if segment_size == 0 and self.total_size > 0: + segment_size = self.total_size // min(self.num_connections, self.total_size) if self.total_size >= self.num_connections else self.total_size + if segment_size == 0: + segment_size = self.total_size + self.num_connections = 1 + print(f"[Downloader {self.download_id}] Warning: Forcing single connection for very small file.") + return self.fallback_download() + + # Create segments + segments = [] + current_byte = 0 + for i in range(self.num_connections): + if current_byte >= self.total_size: + break + + start_byte = current_byte + end_byte = min(current_byte + segment_size - 1, self.total_size - 1) + + # Ensure last segment goes to the end + if i == self.num_connections - 1: + end_byte = self.total_size - 1 + + # Ensure segment is valid + if start_byte <= end_byte < self.total_size: + segments.append((i, start_byte, end_byte)) + self.part_files.append(self.temp_dir / f"part_{i}") + else: + print(f"[Downloader {self.download_id}] Warning: Skipping invalid segment {i}, start={start_byte}, end={end_byte}") + + current_byte = end_byte + 1 + + if not segments: + self.error = f"No valid download segments calculated (Total Size: {self.total_size})." + print(f"[Downloader {self.download_id}] Error: {self.error}") + return False + + # Start download threads + for index, start, end in segments: + if self.is_cancelled: + break + thread = threading.Thread(target=self.download_segment, args=(index, start, end), daemon=True) + self.threads.append(thread) + thread.start() + + # Wait for threads to complete + active_threads = list(self.threads) + while active_threads and not self.is_cancelled: + joined_threads = [] + for t in active_threads: + t.join(timeout=0.2) + if not t.is_alive(): + joined_threads.append(t) + + active_threads = [t for t in active_threads if t not in joined_threads] + + # Handle download completion + if self.is_cancelled: + print(f"[Downloader {self.download_id}] Download stopped (cancelled).") + self.error = self.error or "Download cancelled." + return False + elif self.error: + print(f"[Downloader {self.download_id}] Download stopped (error): {self.error}") + return False + elif self.total_size > 0 and self.downloaded != self.total_size: + self.error = f"Multi-download size mismatch. Expected {self.total_size}, got {self.downloaded}." + print(f"[Downloader {self.download_id}] Error: {self.error}") + return False + + # Merge parts + return self.merge_parts() \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/downloader/manager.py b/custom_nodes/ComfyUI-HuggingFace/downloader/manager.py new file mode 100644 index 0000000000000000000000000000000000000000..35bd7a955e21963c8622726911fe33464a34210d --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/downloader/manager.py @@ -0,0 +1,996 @@ +# ================================================ +# File: downloader/manager.py +# ================================================ +import threading +import time +import datetime +import os +import json +import requests +import subprocess +import platform +import sys +from typing import List, Dict, Any, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from .chunk_downloader import ChunkDownloader + +from ..config import ( + MAX_CONCURRENT_DOWNLOADS, DOWNLOAD_HISTORY_LIMIT, DEFAULT_CONNECTIONS, + METADATA_SUFFIX, PREVIEW_SUFFIX, METADATA_DOWNLOAD_TIMEOUT, PLUGIN_ROOT +) +try: + from folder_paths import get_directory_by_type, get_valid_path, base_path + COMFY_PATHS_AVAILABLE = True +except ImportError: + print("[HuggingFace Manager] Warning: ComfyUI folder_paths not available. Path validation/opening might be limited.") + COMFY_PATHS_AVAILABLE = False + base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) + +# --- Define History File Path --- +# Place it in the root of the extension directory +HISTORY_FILE_PATH = os.path.join(PLUGIN_ROOT, "download_history.json") + +class DownloadManager: + """Manages a queue of downloads, running them concurrently and saving metadata.""" + + def __init__(self, max_concurrent: int = MAX_CONCURRENT_DOWNLOADS): + self.queue: List[Dict[str, Any]] = [] + self.active_downloads: Dict[str, Dict[str, Any]] = {} # {download_id: download_info} + # History now stores more complete dictionaries for retry functionality + self.history: List[Dict[str, Any]] = [] + self.lock: threading.Lock = threading.Lock() + self.max_concurrent: int = max(1, max_concurrent) + self.running: bool = True + self._load_history_from_file() + self._process_thread: threading.Thread = threading.Thread(target=self._process_queue, daemon=True) + print(f"HuggingFace Download Manager starting (Max Concurrent: {self.max_concurrent}).") + self._process_thread.start() + + # --- add_to_queue remains largely the same, ensuring all necessary fields are initialized --- + def add_to_queue(self, download_info: Dict[str, Any]) -> str: + """Adds a download task to the queue.""" + with self.lock: + # Generate a unique ID + timestamp = int(time.time() * 1000) + file_hint = os.path.basename(download_info.get('output_path', 'file'))[:10] + unique_num = sum(1 for item in self.queue if file_hint in item.get("id", "") or any(file_hint in h.get("id","") for h in self.history)) # Check history too + download_id = f"dl_{timestamp}_{unique_num}_{file_hint}" + + # Set initial status and info + download_info["id"] = download_id + download_info["status"] = "queued" + download_info["added_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + download_info["progress"] = 0 + download_info["speed"] = 0 + download_info["error"] = None + download_info["start_time"] = None + download_info["end_time"] = None + download_info["connection_type"] = "N/A" + + # --- Ensure all fields potentially needed for retry exist --- + # (Most were likely already filled by the calling route, but double-check) + required_for_retry = [ + 'url', 'output_path', 'num_connections', 'api_key', 'known_size', + 'huggingface_model_info', 'huggingface_version_info', 'huggingface_primary_file', + 'thumbnail', 'filename', 'model_url_or_id', 'model_version_id', 'model_type', + 'custom_filename', 'force_redownload', 'huggingface_model_name' # Add huggingface_model_name! + ] + for key in required_for_retry: + if key not in download_info: + # Add default or None if missing. More robust handling might be needed + # depending on how routes.py prepares the dict. + if key in ['huggingface_model_info', 'huggingface_version_info', 'huggingface_primary_file']: + download_info[key] = {} + elif key == 'num_connections': + download_info[key] = DEFAULT_CONNECTIONS + elif key == 'force_redownload': + download_info[key] = False + elif key == 'huggingface_model_name': + download_info[key] = None + else: + download_info[key] = None + print(f"[Manager Warning] Queued item '{download_id}' missing '{key}', added default.") + + self.queue.append(download_info) + print(f"[Manager] Queued: {download_info.get('filename', 'N/A')} (ID: {download_id}, Size: {download_info.get('known_size', 'Unknown')})") + return download_id + + # --- cancel_download remains the same --- + def cancel_download(self, download_id: str) -> bool: + """Requests cancellation of a queued or active download.""" + # ... (no changes needed here) ... + print(f"[Manager] Received cancellation request for: {download_id}") # Moved print earlier + downloader_to_cancel: Optional['ChunkDownloader'] = None + found_in_active = False + + with self.lock: + # 1. Check queue first (can be fully handled under lock) + for i, item in enumerate(self.queue): + if item["id"] == download_id: + cancelled_info = self.queue.pop(i) + cancelled_info["status"] = "cancelled" + cancelled_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + cancelled_info["error"] = "Cancelled from queue" + self._add_to_history(cancelled_info) # _add_to_history is safe under lock + print(f"[Manager] Cancelled queued download: {download_id}") + return True # Cancelled from queue, we are done + + # 2. Check active downloads - Find the instance *under lock* + if download_id in self.active_downloads: + found_in_active = True + active_info = self.active_downloads[download_id] + downloader_to_cancel = active_info.get("downloader_instance") + current_status = active_info.get("status") + + # If downloader instance doesn't exist yet (status 'starting') + # or if already terminal, handle it here under lock + if not downloader_to_cancel and current_status == "starting": + print(f"[Manager] Marking 'starting' download as cancelled: {download_id}") + # Mark as cancelled, it won't start or will be caught by wrapper + active_info["status"] = "cancelled" + if not active_info.get("end_time"): + active_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + active_info["error"] = "Cancelled before download thread fully started" + # Don't assign downloader_to_cancel, let process_queue clean it up + downloader_to_cancel = None # Explicitly clear + # Indicate we handled it (or attempted to) + return True # Exit, processed within lock + elif current_status in ["completed", "failed", "cancelled"]: + print(f"[Manager] Download {download_id} is already in terminal state '{current_status}'. Cannot cancel.") + # Indicate we don't need to proceed further outside the lock + downloader_to_cancel = None # Explicitly clear + return False # Already finished/cancelled + + # If we found an instance and it's potentially running, + # store it to call cancel *after* releasing the lock. + print(f"[Manager] Found active downloader instance for {download_id}. Will signal outside lock.") + + # Lock is released automatically here when exiting 'with' + + # 3. Signal the downloader instance *outside* the lock + if downloader_to_cancel: + try: + # Check if already cancelled (to avoid duplicate logs/actions) - This check is thread-safe + if not downloader_to_cancel.is_cancelled: + print(f"[Manager] Calling downloader.cancel() for {download_id}") + downloader_to_cancel.cancel() # This can now call _update_download_status safely + print(f"[Manager] Signalled downloader.cancel() for {download_id}") + # Let the download thread's final status update handle moving to history + return True # Signal sent + else: + print(f"[Manager] Active download {download_id} was already cancelling.") + return True # Already in cancelling state is considered a success here + except Exception as e: + print(f"[Manager] Error calling downloader.cancel() for {download_id}: {e}") + # Update status to failed maybe? Or just log. + # Use _update_download_status directly here as we are outside the lock + self._update_download_status(download_id, status="failed", error=f"Error during cancel signaling: {e}") + return False # Failed to signal + + # 4. Handle cases where it wasn't in queue and wasn't running/starting + if not found_in_active: + print(f"[Manager] Could not cancel - ID not found in queue or active: {download_id}") + return False # Not found + + # It was found in active but was already terminal or couldn't be signalled + # Return value determined above + return False # Should have returned True earlier if successful + + # --- get_status remains the same (still strips data for UI) --- + def get_status(self) -> Dict[str, List[Dict[str, Any]]]: + """Returns the current state of the queue, active downloads, and history. + Strips sensitive/large data for UI efficiency.""" + with self.lock: + # Fields to exclude when sending status to UI + exclude_fields_for_ui = [ + 'downloader_instance', 'huggingface_model_info', 'huggingface_version_info', + 'api_key', # Don't send API key to frontend status + # Large potentially redundant fields: + 'url', 'output_path', 'custom_filename', 'model_url_or_id', + # Keep 'thumbnail', 'filename', 'model_name', 'version_name' etc for display + ] + + # Prepare active downloads list + active_list = [ + {k: v for k, v in item_data.items() if k not in exclude_fields_for_ui} + for item_id, item_data in self.active_downloads.items() + ] + + # Prepare history list similarly + history_list = [ + {k: v for k, v in item_data.items() if k not in exclude_fields_for_ui} + for item_data in self.history[:DOWNLOAD_HISTORY_LIMIT] + ] + + # Return copies + return { + "queue": [ + {k:v for k,v in item.items() if k not in exclude_fields_for_ui} + for item in self.queue + ], + "active": active_list, + "history": history_list, + } + + def _load_history_from_file(self): + """Loads download history from the JSON file.""" + # No lock needed here as it's called during __init__ before the thread starts + if not os.path.exists(HISTORY_FILE_PATH): + print(f"[Manager] History file not found ({HISTORY_FILE_PATH}). Starting with empty history.") + self.history = [] + return + + try: + with open(HISTORY_FILE_PATH, 'r', encoding='utf-8') as f: + loaded_data = json.load(f) + + if isinstance(loaded_data, list): + # Basic validation: Ensure items have IDs (optional but good) + validated_history = [item for item in loaded_data if isinstance(item, dict) and 'id' in item] + invalid_count = len(loaded_data) - len(validated_history) + if invalid_count > 0: + print(f"[Manager Warning] {invalid_count} items removed from loaded history due to missing 'id'.") + + # Ensure history limit + self.history = validated_history[:DOWNLOAD_HISTORY_LIMIT] + print(f"[Manager] Successfully loaded {len(self.history)} items from {HISTORY_FILE_PATH}.") + else: + print(f"[Manager Warning] History file ({HISTORY_FILE_PATH}) contained invalid data (not a list). Starting fresh.") + self.history = [] + # Optionally try to delete the corrupted file? + # try: os.remove(HISTORY_FILE_PATH) except Exception: pass + + except json.JSONDecodeError as e: + print(f"[Manager Error] Failed to parse history file ({HISTORY_FILE_PATH}): {e}. Starting fresh.") + self.history = [] + # Optionally try to delete the corrupted file? + except Exception as e: + print(f"[Manager Error] Failed to read history file ({HISTORY_FILE_PATH}): {e}. Starting fresh.") + self.history = [] + + # --- Save History Method --- + def _save_history_to_file(self): + """Saves the current in-memory history list to the JSON file.""" + # Assumes self.lock is HELD when this is called + history_to_save = self.history[:DOWNLOAD_HISTORY_LIMIT] # Ensure limit before saving + + try: + # Ensure directory exists (should already, but belt-and-suspenders) + os.makedirs(os.path.dirname(HISTORY_FILE_PATH), exist_ok=True) + + # Write atomically (write to temp then rename) to reduce corruption risk + temp_file_path = HISTORY_FILE_PATH + ".tmp" + with open(temp_file_path, 'w', encoding='utf-8') as f: + json.dump(history_to_save, f, indent=2, ensure_ascii=False) # Pretty print + + os.replace(temp_file_path, HISTORY_FILE_PATH) # Atomic rename/replace + # print(f"[Manager] Saved {len(history_to_save)} items to history file.") # Can be noisy + + except Exception as e: + # Log error, but don't crash the manager + print(f"[Manager Error] Failed to save history file ({HISTORY_FILE_PATH}): {e}") + # Attempt to remove temp file if it exists + if os.path.exists(temp_file_path): + try: os.remove(temp_file_path) + except Exception: pass + + # --- Updated _add_to_history Method --- + def _add_to_history(self, download_info: Dict[str, Any]): + """Adds a completed/failed/cancelled item to history (internal). + Stores most original parameters needed for potential retry. + NOW ALSO TRIGGERS SAVING HISTORY TO FILE.""" + # --- (Keep existing logic to prepare info_copy) --- + info_copy = { + k: v for k, v in download_info.items() + if k not in ['downloader_instance'] + } + if "end_time" not in info_copy or info_copy["end_time"] is None: + info_copy["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + if "status" not in info_copy: + info_copy["status"] = "unknown" + if info_copy["status"] != "completed" and not info_copy.get("error"): + info_copy["error"] = f"Finished with status '{info_copy['status']}' but no error recorded." + + # --- Update in-memory history --- + self.history.insert(0, info_copy) # Prepend + # Trim in-memory list (save will also respect limit) + if len(self.history) > DOWNLOAD_HISTORY_LIMIT: + self.history = self.history[:DOWNLOAD_HISTORY_LIMIT] + + # --- Trigger Save to File (still under lock) --- + self._save_history_to_file() # <--- Added call + + # --- Updated clear_history Method --- + def clear_history(self) -> Dict[str, Any]: + """Clears the download history (in-memory and the JSON file).""" + cleared_count = 0 + file_deleted = False + error_msg = None + + try: + with self.lock: + cleared_count = len(self.history) + if cleared_count > 0: + print(f"[Manager] Clearing {cleared_count} items from in-memory history.") + self.history = [] # Clear memory list first + + # Attempt to delete the history file + if os.path.exists(HISTORY_FILE_PATH): + try: + os.remove(HISTORY_FILE_PATH) + file_deleted = True + print(f"[Manager] Deleted history file: {HISTORY_FILE_PATH}") + except OSError as e: + error_msg = f"Failed to delete history file {HISTORY_FILE_PATH}: {e}" + print(f"[Manager Error] {error_msg}") + else: + print(f"[Manager] History file ({HISTORY_FILE_PATH}) did not exist, nothing to delete.") + file_deleted = True # Consider success if file wasn't there anyway + + else: + print("[Manager] History clear request received, but history was already empty.") + # Should we still check/delete the file just in case? Yes. + if os.path.exists(HISTORY_FILE_PATH): + try: + os.remove(HISTORY_FILE_PATH) + file_deleted = True + print(f"[Manager] Deleted potentially orphaned history file: {HISTORY_FILE_PATH}") + except OSError as e: + error_msg = f"Failed to delete potentially orphaned history file {HISTORY_FILE_PATH}: {e}" + print(f"[Manager Error] {error_msg}") + else: + file_deleted = True # Success if clear requested and neither memory/file had anything + + if error_msg: + return {"success": False, "error": f"History cleared from memory, but could not delete file: {error_msg}"} + elif cleared_count > 0: + return {"success": True, "message": f"Cleared {cleared_count} history items (memory and file)."} + else: + # If count was 0 but file deletion was attempted/succeeded + return {"success": True, "message": "History was already empty."} + + except Exception as e: + print(f"[Manager] Critical error during clear_history: {e}") + import traceback + traceback.print_exc() + return {"success": False, "error": f"Failed to clear history due to unexpected error: {e}"} + + # --- _process_queue remains the same --- + def _process_queue(self): + """Internal thread function to manage downloads.""" + # ... (no changes needed here) ... + print("[Manager] Process queue thread started.") + while self.running: + processed_something = False + with self.lock: + # 1. Check for finished/failed/cancelled active downloads to move to history + finished_ids = [ + dl_id for dl_id, info in self.active_downloads.items() + if info.get("status") in ["completed", "failed", "cancelled"] # Use .get() for safety + ] + for dl_id in finished_ids: + # Check if still in active_downloads before popping (might have been removed by another thread edge case?) + if dl_id in self.active_downloads: + finished_info = self.active_downloads.pop(dl_id) + self._add_to_history(finished_info) # Will now store more data + print(f"[Manager] Moved '{finished_info.get('filename', dl_id)}' to history (Status: {finished_info['status']})") + processed_something = True + else: + print(f"[Manager] Warning: Item {dl_id} intended for history was already removed from active list.") + + # 2. Start new downloads if slots available and queue has items + while len(self.active_downloads) < self.max_concurrent and self.queue: + download_info = self.queue.pop(0) + download_id = download_info["id"] + + # Double check if cancelled just before starting + if download_info["status"] == "cancelled": + # Ensure it has an end time before adding to history + if not download_info.get("end_time"): + download_info["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + self._add_to_history(download_info) + print(f"[Manager] Skipping cancelled item from queue: {download_id}") + processed_something = True + continue + + # Update status to 'starting' + download_info["status"] = "starting" + download_info["start_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + download_info["downloader_instance"] = None # Placeholder + + # Add to active downloads BEFORE starting thread + self.active_downloads[download_id] = download_info + + # Start download in a separate thread + thread = threading.Thread( + target=self._download_file_wrapper, + args=(download_info,), + daemon=True # Ensure thread exits if main program exits + ) + thread.start() + print(f"[Manager] Starting download thread for: {download_info.get('filename', 'N/A')} ({download_id})") + processed_something = True + + # Sleep only if nothing was processed to avoid busy-waiting + if not processed_something: + time.sleep(0.5) # Small delay before checking again + + print("[Manager] Process queue thread stopped.") + + # --- _update_download_status remains the same --- + def _update_download_status(self, download_id: str, status: Optional[str] = None, + progress: Optional[float] = None, speed: Optional[float] = None, + error: Optional[str] = None, connection_type: Optional[str] = None): + """Safely updates the status fields of an active download (thread-safe).""" + # ... (no changes needed here) ... + with self.lock: + if download_id in self.active_downloads: + item = self.active_downloads[download_id] + updated = False # Track if any field was actually updated + # Only update if value is provided + if status is not None and item.get("status") != status: + item["status"] = status + updated = True + # If status becomes terminal, record end time + if status in ["completed", "failed", "cancelled"] and not item.get("end_time"): + item["end_time"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + if progress is not None: + # Clamp progress between 0 and 100 + clamped_progress = max(0.0, min(100.0, progress)) + # Only update if different enough to avoid excessive updates + if abs(item.get("progress", 0) - clamped_progress) > 0.01: + item["progress"] = clamped_progress + updated = True + if speed is not None: + clamped_speed = max(0.0, speed) # Speed cannot be negative + if item.get("speed") != clamped_speed: # Update if different + item["speed"] = clamped_speed + updated = True + if error is not None and item.get("error") != error: + # Update error only if it's new or different + item["error"] = str(error)[:500] # Limit length + updated = True + if connection_type is not None and connection_type != "N/A" and item.get("connection_type") != connection_type: # Only update if not N/A and different + item["connection_type"] = connection_type + updated = True + + # --- _save_huggingface_metadata remains the same --- + def _save_huggingface_metadata(self, download_info: Dict[str, Any]): + """Saves the .cminfo.json file.""" + # ... (no changes needed here) ... + output_path = download_info.get('output_path') + model_info = download_info.get('huggingface_model_info', {}) + version_info = download_info.get('huggingface_version_info', {}) + primary_file = download_info.get('huggingface_primary_file', {}) + download_id = download_info.get('id', 'unknown') + + try: + file_meta = primary_file.get('metadata', {}) or {} # Ensure dict + creator_info = model_info.get('creator', {}) or {} + model_stats = model_info.get('stats', {}) or {} + version_stats = version_info.get('stats', {}) or {} + + metadata = { + "ModelId": model_info.get('id', version_info.get('modelId')) , # Use .get() on version info too + "ModelName": model_info.get('name', version_info.get('model',{}).get('name')), # Nested .get() + "ModelDescription": model_info.get('description'), + "CreatorUsername": creator_info.get('username'), + "Nsfw": model_info.get('nsfw', version_info.get('model',{}).get('nsfw')), + "Poi": model_info.get('poi', version_info.get('model',{}).get('poi')), + "AllowNoCredit": model_info.get('allowNoCredit', True), + "AllowCommercialUse": str(model_info.get('allowCommercialUse', 'Unknown')), # Ensure string + "AllowDerivatives": model_info.get('allowDerivatives', True), + "AllowDifferentLicense": model_info.get('allowDifferentLicense', True), + "Tags": model_info.get('tags', []), + "ModelType": model_info.get('type'), + "VersionId": version_info.get('id'), + "VersionName": version_info.get('name'), + "VersionDescription": version_info.get('description'), + "BaseModel": version_info.get('baseModel'), + "BaseModelType": version_info.get('baseModelType'), + "EarlyAccessDeadline": version_info.get('earlyAccessDeadline'), + "VersionPublishedAt": version_info.get('publishedAt'), + "VersionUpdatedAt": version_info.get('updatedAt'), + "VersionStatus": version_info.get('status'), + "IsPrimaryFile": primary_file.get('primary', False), + "PrimaryFileId": primary_file.get('id'), + "PrimaryFileName": primary_file.get('name'), + "FileMetadata": { + "fp": file_meta.get('fp'), + "size": file_meta.get('size'), + "format": file_meta.get('format', 'Unknown') + }, + "ImportedAt": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "Hashes": primary_file.get('hashes', {}), + "TrainedWords": version_info.get('trainedWords', []), + "Stats": { + "downloadCount": version_stats.get('downloadCount', model_stats.get('downloadCount', 0)), + "rating": version_stats.get('rating', model_stats.get('rating', 0)), + "ratingCount": version_stats.get('ratingCount', model_stats.get('ratingCount', 0)), + "favoriteCount": version_stats.get('favoriteCount', model_stats.get('favoriteCount', 0)), # Correct source needed? Check API docs + "commentCount": version_stats.get('commentCount', model_stats.get('commentCount', 0)), # Correct source needed? Check API docs + "thumbsUpCount": version_stats.get('thumbsUpCount', 0), + }, + "DownloadUrlUsed": download_info.get('url'), + } + + base, _ = os.path.splitext(output_path) + meta_filename = base + METADATA_SUFFIX + meta_path = os.path.join(os.path.dirname(output_path), meta_filename) + + print(f"[Manager Meta {download_id}] Saving metadata to: {meta_path}") + with open(meta_path, 'w', encoding='utf-8') as f: + json.dump(metadata, f, indent=2, ensure_ascii=False) + print(f"[Manager Meta {download_id}] Metadata saved successfully.") + + except Exception as e: + import traceback + print(f"[Manager Meta {download_id}] Error saving metadata file {meta_path}: {e}") + # traceback.print_exc() # Uncomment for full trace + + # --- _download_and_save_preview remains the same --- + def _download_and_save_preview(self, download_info: Dict[str, Any]): + """Downloads and saves the .preview.jpeg file.""" + # ... (no changes needed here) ... + output_path = download_info.get('output_path') + thumbnail_url = download_info.get('thumbnail') + api_key = download_info.get('api_key') + download_id = download_info.get('id', 'unknown') + + if not output_path: + print(f"[Manager Preview {download_id}] Skipping preview download: Missing output path.") + return + if not thumbnail_url: + print(f"[Manager Preview {download_id}] Skipping preview download: No thumbnail URL provided.") + # Optionally try to find one in version_info images again? Might be redundant. + version_info = download_info.get('huggingface_version_info', {}) + if version_info and isinstance(version_info.get('images'), list) and version_info['images']: + sorted_images = sorted([img for img in version_info['images'] if img and img.get("url")], key=lambda x: x.get('index', 0)) + img_data = next((img for img in sorted_images if img.get("type") == "image" and "/width=" in img.get("url","")), None) # Prefer image type with width param + if not img_data: img_data = next((img for img in sorted_images if img.get("type") == "image"), None) # Fallback to any image type + if not img_data: img_data = next((img for img in sorted_images), None) # Fallback to any image at all + if img_data and img_data.get('url'): + thumbnail_url = img_data['url'] + # Attempt to get a reasonable size (e.g. ~450px width) + if "/width=" in thumbnail_url: + thumbnail_url = thumbnail_url.split("/width=")[0] + "/width=450" + elif "/blob/" not in thumbnail_url: # Avoid adding params to blob URLs + separator = "&" if "?" in thumbnail_url else "?" + thumbnail_url += f"{separator}width=450" + print(f"[Manager Preview {download_id}] Found alternative thumbnail URL: {thumbnail_url}") + else: + print(f"[Manager Preview {download_id}] Still no thumbnail URL found in version info.") + return # Exit if still no URL + else: + return # Exit if no URL and no version info to search + + base, _ = os.path.splitext(output_path) + preview_filename = base + PREVIEW_SUFFIX + preview_path = os.path.join(os.path.dirname(output_path), preview_filename) + + print(f"[Manager Preview {download_id}] Downloading thumbnail from {thumbnail_url} to {preview_path}") + response = None + try: + headers = {} + if api_key: headers["Authorization"] = f"Bearer {api_key}" + response = requests.get(thumbnail_url, stream=True, headers=headers, timeout=METADATA_DOWNLOAD_TIMEOUT, allow_redirects=True) + response.raise_for_status() + content_type = response.headers.get('Content-Type', '').lower() + if not content_type.startswith('image/'): + print(f"[Manager Preview {download_id}] Warning: Thumbnail URL returned non-image content type '{content_type}'. Skipping save.") + return + with open(preview_path, 'wb') as f: + for chunk in response.iter_content(chunk_size=8192): f.write(chunk) + print(f"[Manager Preview {download_id}] Thumbnail downloaded successfully.") + except requests.exceptions.RequestException as e: + error_msg = f"Error downloading thumbnail {thumbnail_url}: {e}" + if hasattr(e, 'response') and e.response is not None: error_msg += f" (Status: {e.response.status_code})" + print(f"[Manager Preview {download_id}] {error_msg}") + except Exception as e: print(f"[Manager Preview {download_id}] Error saving thumbnail {preview_path}: {e}") + finally: + if response: response.close() + + # --- _download_file_wrapper remains the same --- + def _download_file_wrapper(self, download_info: Dict[str, Any]): + """Wraps the download execution, handles status updates, exceptions, and metadata saving.""" + # ... (no changes needed here) ... + download_id = download_info["id"] + filename = download_info.get('filename', download_id) + from .chunk_downloader import ChunkDownloader + downloader = None + success = False + final_status = "failed" + error_msg = None + + try: + print(f"[Downloader Wrapper {download_id}] Preparing download for '{filename}'.") + + # Handle case where url is None (repo downloads) + url = download_info["url"] + if url is None: + print(f"[Downloader Wrapper {download_id}] URL is None, using huggingface_hub") + # For repo downloads or single files, use huggingface_hub + from huggingface_hub import snapshot_download, hf_hub_download + + # Extract just the repo_id from model_url_or_id (remove URL part) + model_url_or_id = download_info["model_url_or_id"] + if model_url_or_id.startswith("https://huggingface.co/"): + # Extract repo_id and filename from URL + parts = model_url_or_id.replace("https://huggingface.co/", "").split("/") + model_id = f"{parts[0]}/{parts[1]}" # Get "Kijai/WanVideo_comfy" + + # Check if it's a specific file or entire repo + if len(parts) >= 4 and parts[2] == "resolve": + # It's a specific file: /resolve/commit/path/to/file + filename = "/".join(parts[4:]) # Get the file path after /resolve/commit/ + print(f"[Downloader Wrapper {download_id}] Downloading single file {filename} from repo {model_id}") + + # Use hf_hub_download for single file + result = hf_hub_download( + repo_id=model_id, + filename=filename, + local_dir=download_info["output_path"], + token=download_info.get("api_key") + ) + else: + # It's an entire repo download + print(f"[Downloader Wrapper {download_id}] Downloading entire repo {model_id}") + result = snapshot_download( + repo_id=model_id, + local_dir=download_info["output_path"], + token=download_info.get("api_key") + ) + else: + # Plain repo_id - assume entire repo download + model_id = model_url_or_id + print(f"[Downloader Wrapper {download_id}] Downloading entire repo {model_id}") + result = snapshot_download( + repo_id=model_id, + local_dir=download_info["output_path"], + token=download_info.get("api_key") + ) + + if result: + success = True + final_status = "completed" + print(f"[Downloader Wrapper {download_id}] Repo download completed: {result}") + else: + raise Exception("snapshot_download failed") + else: + # Normal file download + downloader = ChunkDownloader( + url=url, + output_path=download_info["output_path"], + num_connections=download_info.get("num_connections", DEFAULT_CONNECTIONS), + manager=self, + download_id=download_id, + api_key=download_info.get("api_key"), + known_size=download_info.get("known_size") + ) + + with self.lock: + if download_id not in self.active_downloads or self.active_downloads[download_id]["status"] == "cancelled": + print(f"[Downloader Wrapper {download_id}] Download was cancelled before instance could be fully linked/started.") + self._update_download_status(download_id, status="cancelled", error="Cancelled before start") + return + + # Only set downloader_instance for file downloads + if url is not None: + self.active_downloads[download_id]["downloader_instance"] = downloader + + self._update_download_status(download_id, status="downloading") + print(f"[Downloader Wrapper {download_id}] Starting download process for '{filename}'.") + + # For repo downloads, success is already determined + if url is None: + # Repo download already completed above + print(f"[Downloader Wrapper {download_id}] Repo download already completed") + else: + # File download - use ChunkDownloader + success = downloader.download() # Blocking call + error_msg = downloader.error + + if success: + final_status = "completed" + print(f"[Downloader Wrapper {download_id}] Download completed successfully for '{filename}'.") + try: + self._save_huggingface_metadata(download_info) + self._download_and_save_preview(download_info) + except Exception as meta_err: + print(f"[Downloader Wrapper {download_id}] Error during post-download metadata/preview saving: {meta_err}") + + elif downloader.is_cancelled: + final_status = "cancelled" + error_msg = downloader.error or "Download cancelled" + print(f"[Downloader Wrapper {download_id}] Download cancelled for '{filename}'. Reason: {error_msg}") + else: + final_status = "failed" + error_msg = downloader.error or "Download failed with unknown error" + print(f"[Downloader Wrapper {download_id}] Download failed for '{filename}'. Error: {error_msg}") + + except Exception as e: + import traceback + print(f"--- Critical Error in Download Wrapper {download_id} ('{filename}') ---") + traceback.print_exc() + print("--- End Error ---") + final_status = "failed" + error_msg = f"Unexpected wrapper error: {str(e)}" + if downloader and not downloader.is_cancelled: + try: downloader.cancel() + except: pass + + finally: + final_progress_percent = 0 + conn_type = download_info.get("connection_type", "N/A") + + if downloader: + conn_type = downloader.connection_type + if downloader.total_size and downloader.total_size > 0: + final_progress_percent = (downloader.downloaded / downloader.total_size * 100) + if final_status == "completed": final_progress_percent = 100.0 + final_progress_percent = min(100.0, max(0.0, final_progress_percent)) + + print(f"[Downloader Wrapper {download_id}] Finalizing status: {final_status}, Error: {error_msg}") + self._update_download_status( + download_id, status=final_status, progress=final_progress_percent, + speed=0, error=error_msg, connection_type=conn_type + ) + if final_status == "completed": + print(f"[Manager] Download {download_id} completed ('{filename}'). Manual ComfyUI refresh may be needed for model list.") + + # --- NEW: Retry Download Method --- + def retry_download(self, original_download_id: str) -> Dict[str, Any]: + """Finds a failed/cancelled download in history and re-queues it.""" + with self.lock: + # Find the original download info in history + original_info = next((item for item in self.history if item.get("id") == original_download_id), None) + + if not original_info: + return {"success": False, "error": f"Original download ID '{original_download_id}' not found in history."} + + original_status = original_info.get("status") + if original_status not in ["failed", "cancelled"]: + return {"success": False, "error": f"Cannot retry download with status '{original_status}'. Only 'failed' or 'cancelled' are retryable."} + + # --- Prepare the new download info dictionary --- + # Make a deep copy to avoid modifying the history item directly + try: + retry_info = json.loads(json.dumps(original_info)) + print(retry_info) + except Exception as e: + return {"success": False, "error": f"Failed to copy original download data: {e}"} + + # Remove fields specific to the *previous* attempt + retry_info.pop("id", None) # Will get a new ID + retry_info.pop("status", None) + retry_info.pop("progress", None) + retry_info.pop("speed", None) + retry_info.pop("error", None) + retry_info.pop("start_time", None) + retry_info.pop("end_time", None) + retry_info.pop("added_time", None) + retry_info.pop("connection_type", None) + retry_info.pop("downloader_instance", None) + # --- Crucially: Set force_redownload to True for retry --- + # This ensures it overwrites the potentially corrupted/partial file from the previous attempt. + retry_info["force_redownload"] = True + + # --- Validate required fields for queuing (redundant check, but safe) --- + required_for_retry = [ + 'url', 'output_path', 'num_connections', 'api_key', 'known_size', + 'huggingface_model_info', 'huggingface_version_info', 'huggingface_primary_file', + 'thumbnail', 'filename', 'model_url_or_id', 'model_version_id', 'model_type', + 'custom_filename', 'force_redownload' + ] + missing_keys = [key for key in required_for_retry if key not in retry_info or retry_info[key] is None and key != 'api_key' and key != 'custom_filename'] # Allow api_key/custom_filename to be None + #if missing_keys: + # return {"success": False, "error": f"Cannot retry: Original download data is missing required fields: {', '.join(missing_keys)}"} + + # --- Add the prepared info to the queue (outside the lock for add_to_queue's own lock) --- + # Note: add_to_queue acquires its own lock internally + try: + new_download_id = self.add_to_queue(retry_info) + if new_download_id: # Check if add_to_queue returned a valid ID (indicating success) + with self.lock: + original_len = len(self.history) + # Filter out the item matching the original ID + self.history = [item for item in self.history if item.get("id") != original_download_id] + items_removed = original_len - len(self.history) + + if items_removed == 1: + print(f"[Manager] Successfully removed original download '{original_download_id}' from history.") + return { + "success": True, + "message": f"Retry initiated. New download queued. Original removed from history.", + "new_download_id": new_download_id + } + + else: + # Should have been caught by the except block, but as a failsafe + print(f"[Manager] Retry queueing failed for '{original_download_id}' for an unknown reason (no ID returned).") + return {"success": False, "error": "Failed to queue retry (unknown internal error)."} + + except Exception as e: + print(f"[Manager] Error requeuing download for retry (Original ID: {original_download_id}): {e}") + return {"success": False, "error": f"Failed to queue retry: {e}"} + + # --- NEW: Open Containing Folder Method --- + def open_containing_folder(self, download_id: str) -> Dict[str, Any]: + """Opens the directory containing the specified completed download file.""" + file_path = None + with self.lock: + # Check history first (most likely location for completed items) + item_info = next((item for item in self.history if item.get("id") == download_id), None) + # Fallback to active (less likely, but possible if called very quickly after completion) + if not item_info and download_id in self.active_downloads: + item_info = self.active_downloads.get(download_id) + + if not item_info: + return {"success": False, "error": "Download ID not found."} + + if item_info.get("status") != "completed": + return {"success": False, "error": f"Cannot open path for download with status '{item_info.get('status')}'. Must be 'completed'."} + + file_path = item_info.get("output_path") # Get the full path to the file + + # --- Perform file operations outside the lock --- + if not file_path: + return {"success": False, "error": "Output path not found for this download."} + + # --- Security Check: Ensure the path is within expected ComfyUI directories --- + # (This is a basic check, might need refinement based on your setup) + is_safe_path = False + try: + # Get the absolute path + abs_file_path = os.path.abspath(file_path) + folder_path = os.path.dirname(abs_file_path) + + # Option 1: Check against ComfyUI's known directories (preferred) + if COMFY_PATHS_AVAILABLE: + # Check if the folder is within any known model type directory + known_types = [ + "checkpoints", "loras", "vae", "embeddings", "hypernetworks", + "controlnet", "upscale_models", "clip_vision", "gligen", "configs", + "unet", "diffusers", "motion_models", "poses", "wildcards" + ] + known_dirs = [os.path.abspath(get_directory_by_type(t)) for t in known_types if get_directory_by_type(t)] + # Also allow output and input directories + if get_directory_by_type("output"): known_dirs.append(os.path.abspath(get_directory_by_type("output"))) + if get_directory_by_type("input"): known_dirs.append(os.path.abspath(get_directory_by_type("input"))) + # Add the plugin's own 'other_models' directory as safe + known_dirs.append(os.path.abspath(os.path.join(PLUGIN_ROOT, "other_models"))) + # Add plugin-managed custom roots as safe + try: + import json as _json + _roots_file = os.path.join(PLUGIN_ROOT, "custom_roots.json") + if os.path.exists(_roots_file): + with open(_roots_file, 'r', encoding='utf-8') as _f: + _data = _json.load(_f) + if isinstance(_data, dict): + for _lst in _data.values(): + if isinstance(_lst, list): + for _p in _lst: + if isinstance(_p, str): + known_dirs.append(os.path.abspath(_p)) + except Exception as _e: + print(f"[Manager OpenPath] Warning: Failed to load custom roots: {_e}") + # Include all first-level subdirectories under models_dir as safe + try: + models_dir = getattr(__import__('folder_paths'), 'folder_paths').models_dir + except Exception: + models_dir = None + try: + if models_dir and os.path.isdir(models_dir): + for _name in os.listdir(models_dir): + _p = os.path.join(models_dir, _name) + if os.path.isdir(_p): + known_dirs.append(os.path.abspath(_p)) + except Exception as _e2: + print(f"[Manager OpenPath] Warning: Failed enumerating models_dir subfolders: {_e2}") + + for known_dir in known_dirs: + if os.path.commonpath([known_dir, folder_path]) == known_dir: + is_safe_path = True + break + else: + # Option 2: Fallback - Check if path is within the ComfyUI base directory + comfy_base = os.path.abspath(base_path) + if os.path.commonpath([comfy_base, folder_path]) == comfy_base: + is_safe_path = True + #print(f"[Manager OpenPath Warning] ComfyUI paths unavailable. Using base path check for {folder_path}") + + if not is_safe_path: + print(f"[Manager OpenPath Denied] Path '{folder_path}' is outside known safe ComfyUI directories.") + return {"success": False, "error": "Cannot open path: Directory is outside allowed locations."} + + # --- Open the directory --- + if not os.path.exists(folder_path): + return {"success": False, "error": f"Directory does not exist: {folder_path}"} + if not os.path.isdir(folder_path): + return {"success": False, "error": f"Path is not a directory: {folder_path}"} + + try: + system = platform.system() + print(f"[Manager OpenPath] Attempting to open folder '{folder_path}' on {system}...") + if system == "Windows": + # Use startfile for better handling of spaces etc. + os.startfile(folder_path) + elif system == "Darwin": # macOS + subprocess.check_call(["open", folder_path]) + elif system == "Linux": + # Use xdg-open, handle potential errors if command not found + try: + subprocess.check_call(["xdg-open", folder_path]) + except FileNotFoundError: + # Fallback for headless systems or if xdg-open isn't installed + print(f"[Manager OpenPath] 'xdg-open' not found. Cannot automatically open folder on this Linux system.") + return {"success": False, "error": "'xdg-open' command not found. Cannot open folder."} + else: + print(f"[Manager OpenPath] Unsupported operating system: {system}. Cannot open folder.") + return {"success": False, "error": f"Unsupported OS ({system}) for opening folder."} + + print(f"[Manager OpenPath] Successfully requested folder opening for '{folder_path}'.") + return {"success": True, "message": f"Opened directory: {folder_path}"} + + except Exception as e: + print(f"[Manager OpenPath] Failed to open directory '{folder_path}': {e}") + return {"success": False, "error": f"Failed to open directory: {e}"} + + except Exception as path_err: + # Catch errors during path validation itself + print(f"[Manager OpenPath] Error during path validation/retrieval for {download_id}: {path_err}") + return {"success": False, "error": f"Error processing path: {path_err}"} + +# --- Global Instance --- +manager = DownloadManager(max_concurrent=MAX_CONCURRENT_DOWNLOADS) + +# --- Graceful Shutdown --- +# (shutdown_manager remains the same) +def shutdown_manager(): + # ... (no changes) ... + print("[Manager] Shutdown requested.") + if manager: + manager.running = False + acquired_lock = False + try: acquired_lock = manager.lock.acquire(timeout=1.0) + except RuntimeError: pass # Lock might not be initialised if init failed + + if acquired_lock: + try: + active_ids = list(manager.active_downloads.keys()) + queue_ids = [item['id'] for item in manager.queue] + print(f"[Manager] Requesting cancellation for {len(active_ids)} active and {len(queue_ids)} queued downloads on shutdown...") + all_ids_to_cancel = active_ids + queue_ids + manager.lock.release() # Release lock BEFORE calling cancel_download + + for dl_id in all_ids_to_cancel: + try: + # Reuse cancel_download which handles both active and queued safely + manager.cancel_download(dl_id) + except Exception as e: + print(f"Error cancelling {dl_id} during shutdown: {e}") + except Exception as e: + print(f"[Manager] Error accessing lists during shutdown (after lock acquired): {e}") + try: manager.lock.release() # Ensure release on error + except RuntimeError: pass + # Give threads/tasks a brief moment to react + time.sleep(0.5) + else: + print("[Manager] Warning: Could not acquire lock to cancel downloads during shutdown.") + + # Attempt to join the manager's process thread (best effort) + try: + if manager._process_thread and manager._process_thread.is_alive(): + manager._process_thread.join(timeout=2.0) + if manager._process_thread.is_alive(): + print("[Manager] Warning: Process thread did not exit cleanly within timeout.") + except Exception as e: + print(f"[Manager] Error joining manager thread during shutdown: {e}") + print("[Manager] Shutdown complete.") + +import atexit +atexit.register(shutdown_manager) + +# Create the global manager instance +manager = DownloadManager() diff --git a/custom_nodes/ComfyUI-HuggingFace/server/__init__.py b/custom_nodes/ComfyUI-HuggingFace/server/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..406a5495761804a3421442a7c1ad954bb002fec1 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/__init__.py @@ -0,0 +1,5 @@ +# ================================================ +# File: server/__init__.py +# ================================================ +# Import the routes package to ensure all route decorators are executed. +from . import routes \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/CancelDownload.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/CancelDownload.py new file mode 100644 index 0000000000000000000000000000000000000000..1c139e2781b498cd27e011608f2caa5741d40daa --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/CancelDownload.py @@ -0,0 +1,44 @@ +# ================================================ +# File: server/routes/CancelDownload.py +# ================================================ +import json +from aiohttp import web +import server # ComfyUI server instance +from ..utils import get_request_json +from ...downloader.manager import manager as download_manager + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/cancel") +async def route_cancel_download(request): + """API Endpoint to cancel a download.""" + try: + data = await get_request_json(request) + download_id = data.get("download_id") + if not download_id: + print("not download id " + download_id) + raise web.HTTPBadRequest(reason="Missing 'download_id'") + + success = download_manager.cancel_download(download_id) + if success: + return web.json_response({ + "status": "cancelled", # Or "cancellation_requested" ? + "message": f"Cancellation requested for download ID: {download_id}.", + "download_id": download_id + }) + else: + # Might be already completed/failed/cancelled and in history, or invalid ID + raise web.HTTPNotFound(reason=f"Download ID {download_id} not found in active queue or running downloads.") + + except web.HTTPError as http_err: + # Consistent error handling + body_detail = "" + try: + body_detail = await http_err.text() if hasattr(http_err, 'text') else http_err.body.decode('utf-8', errors='ignore') if http_err.body else "" + if body_detail.startswith('{') and body_detail.endswith('}'): body_detail = json.loads(body_detail) + except Exception: pass + return web.json_response({"error": http_err.reason, "details": body_detail or "No details", "status_code": http_err.status}, status=http_err.status) + + except Exception as e: + print(f"Error cancelling download: {e}") + return web.json_response({"error": "Internal Server Error", "details": f"Failed to cancel download: {str(e)}", "status_code": 500}, status=500) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/ClearHistory.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/ClearHistory.py new file mode 100644 index 0000000000000000000000000000000000000000..e02708f8a9832b71553a7fb00e5c68b423e8f472 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/ClearHistory.py @@ -0,0 +1,32 @@ +# ================================================ +# File: server/routes/ClearHistory.py +# ================================================ +import asyncio +from aiohttp import web + +import server # ComfyUI server instance +from ...downloader.manager import manager as download_manager + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/clear_history") +async def route_clear_history(request): + """API Endpoint to clear the download history.""" + if not download_manager: + return web.json_response({"error": "Download Manager not initialized"}, status=500) + + try: + # No request body needed for this action + print(f"[API Route /huggingface/clear_history] Received clear history request.") + + # Call manager method in thread + result = await asyncio.to_thread(download_manager.clear_history) + + status_code = 200 if result.get("success") else 500 # Use 500 for internal clear error + return web.json_response(result, status=status_code) + + except Exception as e: + import traceback + print(f"Error handling /huggingface/clear_history request: {e}") + # traceback.print_exc() # Uncomment for detailed logs + return web.json_response({"error": "Internal Server Error", "details": f"An unexpected error occurred: {str(e)}"}, status=500) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/DownloadModel.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/DownloadModel.py new file mode 100644 index 0000000000000000000000000000000000000000..90a24b8bb3667e36b02dbffa8ed71e37ddbf64e4 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/DownloadModel.py @@ -0,0 +1,186 @@ +# ================================================ +# File: server/routes/DownloadModel.py +# ================================================ +import os +import json +import traceback +import re +from aiohttp import web + +import server # ComfyUI server instance +from ..utils import get_request_json, resolve_huggingface_api_key +from ...downloader.manager import manager as download_manager +from ...api.huggingface import HuggingFaceAPI +from ...utils.helpers import get_model_dir, parse_huggingface_input, sanitize_filename +from ...config import METADATA_SUFFIX, PREVIEW_SUFFIX + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/download") +async def route_download_model(request): + """API Endpoint to initiate a download.""" + try: + data = await get_request_json(request) + + model_url_or_id = data.get("model_url_or_id") + model_type_value = data.get("model_type", "checkpoint") + explicit_save_root = (data.get("save_root") or "").strip() + custom_filename_input = data.get("custom_filename", "").strip() + selected_subdir = (data.get("subdir") or "").strip() + num_connections = int(data.get("num_connections", 4)) + force_redownload = bool(data.get("force_redownload", False)) + resolved_api_key = resolve_huggingface_api_key(data) + + if not model_url_or_id: + raise web.HTTPBadRequest(reason="Missing 'model_url_or_id'") + + print(f"[HF Download] Request: {model_url_or_id}, SaveType: {model_type_value}") + + # Parse HuggingFace URL/ID + parsed_model_id, parsed_filename = parse_huggingface_input(model_url_or_id) + + if not parsed_model_id: + raise web.HTTPBadRequest(reason=f"Could not parse HuggingFace model ID from: {model_url_or_id}") + + target_model_id = parsed_model_id + print(f"[HF Download] Parsed Model ID: {target_model_id}") + + # Initialize API + api = HuggingFaceAPI(resolved_api_key) + + if parsed_filename: + # Direct download from URL - skip API calls + target_filename = parsed_filename + # Try to get model name from the model_id itself + model_name = target_model_id.split('/')[-1] if target_model_id else "Unknown Model" + model_info = {"id": target_model_id, "name": model_name} + print(f"[HF Download] Direct download file: {target_filename}") + print(f"[HF Download] Using extracted model name: {model_name}") + else: + # Skip API calls for public repos, use only huggingface_hub + if resolved_api_key: + # For private repos, try API calls to get model info + api = HuggingFaceAPI(resolved_api_key) + model_info = api.get_model_info(target_model_id) + + if not model_info or "error" in model_info: + print(f"[HF Download] Model info failed, using huggingface_hub directly") + target_filename = parsed_filename if parsed_filename else None + model_info = {"id": target_model_id, "name": target_model_id.split('/')[-1]} + else: + # For private repos, still use huggingface_hub for download + target_filename = parsed_filename if parsed_filename else None + else: + # For public repos, skip API calls entirely + print(f"[HF Download] Public repo, using huggingface_hub directly") + target_filename = parsed_filename if parsed_filename else None + model_info = {"id": target_model_id, "name": target_model_id.split('/')[-1]} + + if not target_filename: + print(f"[HF Download] No specific file found, letting huggingface_hub auto-detect") + target_filename = None + + print(f"[HF Download] Target file: {target_filename}") + + # Determine save directory + target_dir = get_model_dir(model_type_value, explicit_save_root, selected_subdir) + if not target_dir: + raise web.HTTPBadRequest(reason=f"Invalid model type: {model_type_value}") + + # Determine filename + if custom_filename_input: + final_filename = sanitize_filename(custom_filename_input) + elif target_filename is None: + # For repo downloads, use model name as folder + final_filename = model_info.get("name", target_model_id.split('/')[-1]) + else: + final_filename = os.path.basename(target_filename) + + save_path = os.path.join(target_dir, final_filename) + + # Check if file exists + if os.path.exists(save_path) and not force_redownload: + raise web.HTTPBadRequest(reason=f"File already exists: {final_filename}") + + # Start download + if target_filename is None: + # For repo downloads, don't construct URL - let huggingface_hub handle it + download_url = None + else: + download_url = f"https://huggingface.co/{target_model_id}/resolve/main/{target_filename}" + + download_info = { + "model_url_or_id": model_url_or_id, + "save_path": save_path, + "output_path": save_path, # Add this for ChunkDownloader + "url": download_url, # Add this for ChunkDownloader + "filename": final_filename, + "model_type": model_type_value, + "download_url": download_url, + "huggingface_model_info": model_info, + "huggingface_filename": target_filename, + "num_connections": num_connections, + "force_redownload": force_redownload, + # Add missing fields to prevent warnings + "api_key": resolved_api_key, + "known_size": None, + "huggingface_version_info": {}, + "huggingface_primary_file": None, + "thumbnail": None, + "model_version_id": None, + "custom_filename": custom_filename_input, + "huggingface_model_name": model_info.get("name", target_model_id.split('/')[-1]) + } + + download_id = download_manager.add_to_queue(download_info) + + # Extract model name from model_info if available, otherwise from model_id + print(f"[DEBUG] model_info: {model_info}") + print(f"[DEBUG] target_model_id: {target_model_id}") + + # Try to parse model_info as JSON if it's a string + parsed_model_info = None + if isinstance(model_info, str): + try: + import json + parsed_model_info = json.loads(model_info) + print(f"[DEBUG] Parsed model_info from JSON: {parsed_model_info}") + except: + print(f"[DEBUG] Failed to parse model_info as JSON") + parsed_model_info = None + else: + parsed_model_info = model_info + + if parsed_model_info and isinstance(parsed_model_info, dict) and parsed_model_info.get('name'): + model_display_name = parsed_model_info['name'] + print(f"[DEBUG] Using parsed model_info name: {model_display_name}") + elif model_info and isinstance(model_info, dict) and model_info.get('name'): + model_display_name = model_info['name'] + print(f"[DEBUG] Using model_info name: {model_display_name}") + else: + model_display_name = target_model_id.split('/')[-1] if target_model_id else "Unknown Model" + print(f"[DEBUG] Using parsed name: {model_display_name}") + + print(f"[DEBUG] Final model_display_name: {model_display_name}") + + response_data = { + "download_id": download_id, + "huggingface_model_id": target_model_id, + "huggingface_model_name": model_display_name, # Add model name + "huggingface_filename": target_filename, + "huggingface_model_info": model_info, + "save_path": save_path, + "filename": final_filename, + "status": "queued" # Changed from "started" to "queued" to match frontend expectation + } + + print(f"[DEBUG] Response data huggingface_model_name: {response_data['huggingface_model_name']}") + + return web.json_response(response_data) + + except web.HTTPException: + raise + except Exception as e: + print(f"--- Unhandled Error in /huggingface/download ---") + traceback.print_exc() + raise web.HTTPInternalServerError(reason=str(e)) diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/GetBaseModels.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetBaseModels.py new file mode 100644 index 0000000000000000000000000000000000000000..1937cb5995a5cfdd0f47046663d0b980e35fc6b5 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetBaseModels.py @@ -0,0 +1,51 @@ +# ================================================ +# File: server/routes/GetBaseModels.py +# ================================================ +from aiohttp import web +import yaml +import re +import server # ComfyUI server instance +from ...config import AVAILABLE_MEILI_BASE_MODELS +from ...api.huggingface import HuggingFaceAPI +from ..utils import resolve_huggingface_api_key + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.get("/api/huggingface/base_models") +async def route_get_base_models(request): + """API Endpoint to get the known base model types for filtering.""" + try: + # Try to get base models from HuggingFace API + api_key = resolve_huggingface_api_key({}) + api = HuggingFaceAPI(api_key) + + # Get popular models and extract their base models + popular_models = [ + "stabilityai/stable-diffusion-xl-base-1.0", + "runwayml/stable-diffusion-v1-5", + "stabilityai/stable-diffusion-2-1", + "CompVis/stable-diffusion-v1-4", + "black-forest-labs/FLUX.1-dev", + "black-forest-labs/FLUX.1-schnell" + ] + + base_models = set() + base_models.update(AVAILABLE_MEILI_BASE_MODELS) # Keep hardcoded ones as fallback + + for model_id in popular_models: + try: + model_info = api.get_model_info(model_id) + if model_info and not isinstance(model_info, dict) or "error" not in model_info: + # Add the model itself as a base model + base_models.add(model_id) + except Exception as e: + print(f"[GetBaseModels] Error getting info for {model_id}: {e}") + continue + + return web.json_response({"base_models": sorted(list(base_models))}) + + except Exception as e: + print(f"Error getting base model types: {e}") + # Fallback to hardcoded list + return web.json_response({"base_models": AVAILABLE_MEILI_BASE_MODELS}) + return web.json_response({"error": "Internal Server Error", "details": str(e), "status_code": 500}, status=500) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelDetails.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelDetails.py new file mode 100644 index 0000000000000000000000000000000000000000..50429cbfd0e43fe8d7837c0b0e789486508a7faf --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelDetails.py @@ -0,0 +1,120 @@ +# ================================================ +# File: server/routes/GetModelDetails.py +# ================================================ +import os +import json +import traceback +from aiohttp import web + +import server # ComfyUI server instance +from ..utils import get_request_json, resolve_huggingface_api_key +from ...api.huggingface import HuggingFaceAPI + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/get_model_details") +async def route_get_model_details(request): + """API Endpoint to fetch model details from HuggingFace model card.""" + try: + data = await get_request_json(request) + model_url_or_id = data.get("model_url_or_id") + resolved_api_key = resolve_huggingface_api_key(data) + + if not model_url_or_id: + raise web.HTTPBadRequest(reason="Missing 'model_url_or_id'") + + # Parse model ID + from ...utils.helpers import parse_huggingface_input + parsed_model_id, parsed_filename = parse_huggingface_input(model_url_or_id) + + if not parsed_model_id: + raise web.HTTPBadRequest(reason=f"Could not parse HuggingFace model ID from: {model_url_or_id}") + + # Get model info using huggingface_hub + try: + from huggingface_hub import ModelCard + model_card = ModelCard.load(parsed_model_id, token=resolved_api_key) + + # Extract data from model card + model_name = model_card.data.get("model_name", parsed_model_id.split('/')[-1]) + creator_username = model_card.data.get("author", "Unknown Creator") + description = model_card.text + base_model = model_card.data.get("base_model", []) + tags = model_card.data.get("tags", []) + license_info = model_card.data.get("license", "Unknown") + + # Get model info from API as fallback for stats + try: + api = HuggingFaceAPI(resolved_api_key) + api_info = api.get_model_info(parsed_model_id) + + stats = {} + if api_info and not isinstance(api_info, dict) or "error" not in api_info: + stats = { + "downloads": api_info.get("downloads", 0), + "likes": api_info.get("likes", 0), + "created_at": api_info.get("created_at", ""), + "modified_at": api_info.get("modified_at", "") + } + except Exception as api_error: + print(f"[GetModelDetails] API info failed, using model card only: {api_error}") + stats = {} + + response_data = { + "model_name": model_name, + "creator_username": creator_username, + "description": description, + "base_model": base_model, + "tags": tags, + "license": license_info, + "stats": stats, + "model_id": parsed_model_id, + "huggingface_model_name": model_name + } + + return web.json_response(response_data) + + except Exception as e: + print(f"[GetModelDetails] Error loading model card: {e}") + # Try to get at least basic info from API + try: + api = HuggingFaceAPI(resolved_api_key) + api_info = api.get_model_info(parsed_model_id) + + if api_info and not isinstance(api_info, dict) or "error" not in api_info: + return web.json_response({ + "model_name": api_info.get("id", parsed_model_id.split('/')[-1]), + "creator_username": api_info.get("author", "Unknown Creator"), + "description": api_info.get("description", "No description available"), + "base_model": [], + "tags": api_info.get("tags", []), + "license": api_info.get("license", "Unknown"), + "stats": { + "downloads": api_info.get("downloads", 0), + "likes": api_info.get("likes", 0), + "created_at": api_info.get("created_at", ""), + "modified_at": api_info.get("modified_at", "") + }, + "model_id": parsed_model_id, + "huggingface_model_name": api_info.get("id", parsed_model_id.split('/')[-1]) + }) + except Exception as api_error: + print(f"[GetModelDetails] API fallback also failed: {api_error}") + + # Final fallback to minimal info + return web.json_response({ + "model_name": parsed_model_id.split('/')[-1], + "creator_username": "Unknown", + "description": "Model details not available", + "base_model": [], + "tags": [], + "license": "Unknown", + "stats": {}, + "model_id": parsed_model_id, + "huggingface_model_name": parsed_model_id.split('/')[-1] + }) + + except Exception as e: + print(f"Error in get_model_details: {e}") + traceback.print_exc() + return web.json_response({"error": "Internal Server Error", "details": str(e)}, status=500) diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelDirs.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelDirs.py new file mode 100644 index 0000000000000000000000000000000000000000..4c4e1223876e4b69d34dae955f6d8ae388d77250 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelDirs.py @@ -0,0 +1,337 @@ +# ================================================ +# File: server/routes/GetModelDirs.py +# ================================================ +import os +import json +from aiohttp import web + +import server # ComfyUI server instance +from ...utils.helpers import ( + get_model_dir, + get_model_folder_paths, + get_model_type_folder_name, + sanitize_filename, +) +from ...config import PLUGIN_ROOT +import folder_paths + +prompt_server = server.PromptServer.instance + +CUSTOM_ROOTS_FILE = os.path.join(PLUGIN_ROOT, "custom_roots.json") +ROOT_SETTINGS_FILE = os.path.join(PLUGIN_ROOT, "root_settings.json") +GLOBAL_ROOT_KEY = "global_default_root" + +def _load_custom_roots(): + try: + if os.path.exists(CUSTOM_ROOTS_FILE): + with open(CUSTOM_ROOTS_FILE, 'r', encoding='utf-8') as f: + data = json.load(f) + if isinstance(data, dict): + # Normalize values to lists of strings + return {k: [str(p) for p in (v or []) if isinstance(p, str)] for k, v in data.items()} + except Exception as e: + print(f"[HuggingFace] Warning: Failed to load custom roots: {e}") + return {} + +def _save_custom_roots(data): + try: + with open(CUSTOM_ROOTS_FILE, 'w', encoding='utf-8') as f: + json.dump(data, f, indent=2) + return True + except Exception as e: + print(f"[HuggingFace] Error writing custom roots file: {e}") + return False + +def _load_root_settings(): + try: + if os.path.exists(ROOT_SETTINGS_FILE): + with open(ROOT_SETTINGS_FILE, 'r', encoding='utf-8') as f: + data = json.load(f) + if isinstance(data, dict): + return data + except Exception as e: + print(f"[HuggingFace] Warning: Failed to load root settings: {e}") + return {} + +def _save_root_settings(data): + try: + with open(ROOT_SETTINGS_FILE, 'w', encoding='utf-8') as f: + json.dump(data, f, indent=2) + return True + except Exception as e: + print(f"[HuggingFace] Error writing root settings file: {e}") + return False + +def get_global_default_root(): + settings = _load_root_settings() + raw = settings.get(GLOBAL_ROOT_KEY) + if isinstance(raw, str): + raw = raw.strip() + if raw: + return os.path.abspath(raw) + return None + +def _set_global_default_root(path): + settings = _load_root_settings() + if path and str(path).strip(): + settings[GLOBAL_ROOT_KEY] = os.path.abspath(str(path).strip()) + else: + settings.pop(GLOBAL_ROOT_KEY, None) + return _save_root_settings(settings) + +def _get_global_root_for_type(model_type: str): + global_root = get_global_default_root() + if not global_root: + return None + model_subfolder = get_model_type_folder_name(model_type) + return os.path.abspath(os.path.join(global_root, model_subfolder)) + +def _get_effective_base_dir(model_type: str): + global_type_root = _get_global_root_for_type(model_type) + if global_type_root: + return global_type_root + return get_model_dir(model_type) + +def _get_custom_roots_for_type(model_type: str): + roots_map = _load_custom_roots() + model_type_raw = (model_type or '').lower().strip() + canonical_type = get_model_type_folder_name(model_type_raw) + custom = [] + for key in [model_type_raw, canonical_type]: + for p in roots_map.get(key, []): + ap = os.path.abspath(p) + if ap not in custom: + custom.append(ap) + return custom + +def _get_all_roots_for_type(model_type: str): + model_type = (model_type or '').lower().strip() + roots = [] + + # Prefer global default root for this type when configured + global_type_root = _get_global_root_for_type(model_type) + if global_type_root and global_type_root not in roots: + roots.append(global_type_root) + + # Include all registered paths from ComfyUI (respects extra_model_paths.yaml) + for p in get_model_folder_paths(model_type): + if p not in roots: + roots.append(p) + + # Include plugin custom roots + for p in _get_custom_roots_for_type(model_type): + if p not in roots: + roots.append(p) + + # Ensure at least one sane fallback + if not roots: + d = get_model_dir(model_type) + if d: + roots.append(os.path.abspath(d)) + + # Include all immediate subdirectories inside the main ComfyUI models folder + try: + models_dir = getattr(folder_paths, 'models_dir', None) + if not models_dir: + base = getattr(folder_paths, 'base_path', os.getcwd()) + models_dir = os.path.join(base, 'models') + if os.path.isdir(models_dir): + for name in os.listdir(models_dir): + p = os.path.join(models_dir, name) + if os.path.isdir(p): + ap = os.path.abspath(p) + if ap not in roots: + roots.append(ap) + except Exception as e: + print(f"[HuggingFace] Warning: Failed to enumerate models dir subfolders: {e}") + return roots + +def _list_subdirs(root_dir: str, max_entries: int = 5000): + """Return a sorted list of relative subdirectory paths under root_dir, including nested.""" + rel_dirs = set() + root_dir = os.path.abspath(root_dir) + count = 0 + for current, dirs, _files in os.walk(root_dir): + # Avoid following symlinks to reduce risk + abs_current = os.path.abspath(current) + try: + rel = os.path.relpath(abs_current, root_dir) + except Exception: + continue + if rel == ".": + rel = "" # represent root as empty + rel_dirs.add(rel) + count += 1 + if count >= max_entries: + break + return sorted(rel_dirs) + +@prompt_server.routes.get("/api/huggingface/model_dirs") +async def route_get_model_dirs(request): + """List the base directory (or provided root) and all subdirectories for a given model type.""" + model_type = request.query.get("type", "checkpoints").lower().strip() + root = (request.query.get("root") or "").strip() + try: + base_dir = os.path.abspath(root) if root else _get_effective_base_dir(model_type) + os.makedirs(base_dir, exist_ok=True) + subdirs = _list_subdirs(base_dir) + global_root = get_global_default_root() + return web.json_response({ + "model_type": model_type, + "base_dir": base_dir, + "subdirs": subdirs, # relative paths, "" represents the base root + "global_root": global_root or "", + "using_global_root": bool(global_root and not root), + }) + except Exception as e: + return web.json_response({"error": "Failed to list directories", "details": str(e)}, status=500) + +@prompt_server.routes.post("/api/huggingface/create_dir") +async def route_create_model_dir(request): + """Create a new subdirectory under a model type's base directory.""" + try: + data = await request.json() + model_type = (data.get("model_type") or "checkpoints").lower().strip() + new_dir = (data.get("new_dir") or "").strip() + if not new_dir: + return web.json_response({"error": "Missing 'new_dir'"}, status=400) + + # If client provided an explicit root, prefer it + base_dir = (data.get("root") or "").strip() + base_dir = os.path.abspath(base_dir) if base_dir else _get_effective_base_dir(model_type) + os.makedirs(base_dir, exist_ok=True) + + # Normalize and sanitize each part; disallow absolute and traversal + norm = os.path.normpath(new_dir.replace("\\", "/")) + parts = [p for p in norm.split("/") if p and p not in (".", "..")] + safe_parts = [sanitize_filename(p) for p in parts] + rel_path = os.path.join(*safe_parts) if safe_parts else "" + if not rel_path: + return web.json_response({"error": "Invalid folder name"}, status=400) + + abs_path = os.path.abspath(os.path.join(base_dir, rel_path)) + # Ensure it remains inside base_dir + if os.path.commonpath([abs_path, os.path.abspath(base_dir)]) != os.path.abspath(base_dir): + return web.json_response({"error": "Invalid path"}, status=400) + + os.makedirs(abs_path, exist_ok=True) + return web.json_response({ + "success": True, + "created": rel_path, + "abs_path": abs_path, + }) + except Exception as e: + return web.json_response({"error": "Failed to create directory", "details": str(e)}, status=500) + +@prompt_server.routes.post("/api/huggingface/create_model_type") +async def route_create_model_type(request): + """Create a new first-level folder under the main models directory.""" + try: + data = await request.json() + name = (data.get("name") or "").strip() + if not name: + return web.json_response({"error": "Missing 'name'"}, status=400) + + # Sanitize folder name to a single path component + from ...utils.helpers import sanitize_filename + safe = sanitize_filename(name) + if not safe: + return web.json_response({"error": "Invalid folder name"}, status=400) + + # Resolve models directory + models_dir = getattr(folder_paths, 'models_dir', None) + if not models_dir: + base = getattr(folder_paths, 'base_path', os.getcwd()) + models_dir = os.path.join(base, 'models') + + abs_path = os.path.abspath(os.path.join(models_dir, safe)) + # Ensure it remains inside models_dir + if os.path.commonpath([abs_path, os.path.abspath(models_dir)]) != os.path.abspath(models_dir): + return web.json_response({"error": "Invalid path"}, status=400) + + os.makedirs(abs_path, exist_ok=True) + return web.json_response({"success": True, "name": safe, "path": abs_path}) + except Exception as e: + return web.json_response({"error": "Failed to create model type folder", "details": str(e)}, status=500) + +@prompt_server.routes.get("/api/huggingface/model_roots") +async def route_get_model_roots(request): + """Return all known root directories for a model type (ComfyUI + plugin custom roots).""" + model_type = request.query.get("type", "checkpoints").lower().strip() + roots = _get_all_roots_for_type(model_type) + global_root = get_global_default_root() + return web.json_response({ + "model_type": model_type, + "roots": roots, + "global_root": global_root or "", + "effective_root": _get_effective_base_dir(model_type), + }) + +@prompt_server.routes.post("/api/huggingface/create_root") +async def route_create_model_root(request): + """Create a new root directory for a model type and register it in plugin config. + Note: ComfyUI may require restart to recognize this root globally; the plugin uses it immediately. + """ + try: + data = await request.json() + model_type = (data.get("model_type") or "checkpoints").lower().strip() + abs_path = os.path.expanduser((data.get("path") or "").strip()) + if not abs_path: + return web.json_response({"error": "Missing 'path'"}, status=400) + if not os.path.isabs(abs_path): + return web.json_response({"error": "Path must be absolute"}, status=400) + # Normalize to absolute path + abs_path = os.path.abspath(abs_path) + # Create directory if missing + os.makedirs(abs_path, exist_ok=True) + + type_key = get_model_type_folder_name(model_type) + roots = _load_custom_roots() + current = roots.get(type_key, []) + if abs_path not in current: + current.append(abs_path) + roots[type_key] = current + if not _save_custom_roots(roots): + return web.json_response({"error": "Failed to persist custom root"}, status=500) + return web.json_response({"success": True, "path": abs_path}) + except Exception as e: + return web.json_response({"error": "Failed to create root", "details": str(e)}, status=500) + +@prompt_server.routes.get("/api/huggingface/global_root") +async def route_get_global_root(request): + """Return the persisted global download root (if configured).""" + global_root = get_global_default_root() + return web.json_response({ + "global_root": global_root or "", + "enabled": bool(global_root), + }) + +@prompt_server.routes.post("/api/huggingface/global_root") +async def route_set_global_root(request): + """Persist a global download root used as: /.""" + try: + data = await request.json() + path = os.path.expanduser((data.get("path") or "").strip()) + if not path: + return web.json_response({"error": "Missing 'path'"}, status=400) + if not os.path.isabs(path): + return web.json_response({"error": "Path must be absolute"}, status=400) + + abs_path = os.path.abspath(path) + os.makedirs(abs_path, exist_ok=True) + if not _set_global_default_root(abs_path): + return web.json_response({"error": "Failed to persist global root"}, status=500) + + return web.json_response({"success": True, "global_root": abs_path}) + except Exception as e: + return web.json_response({"error": "Failed to set global root", "details": str(e)}, status=500) + +@prompt_server.routes.post("/api/huggingface/global_root/clear") +async def route_clear_global_root(request): + """Clear the persisted global download root.""" + try: + if not _set_global_default_root(None): + return web.json_response({"error": "Failed to clear global root"}, status=500) + return web.json_response({"success": True}) + except Exception as e: + return web.json_response({"error": "Failed to clear global root", "details": str(e)}, status=500) diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelTypes.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelTypes.py new file mode 100644 index 0000000000000000000000000000000000000000..8b4fc67f37ec2aea5670d21169cf88abaaac49de --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetModelTypes.py @@ -0,0 +1,31 @@ +# ================================================ +# File: server/routes/GetModelTypes.py +# ================================================ +import os +from aiohttp import web +import server # ComfyUI server instance +import folder_paths + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.get("/api/huggingface/model_types") +async def route_get_model_types(request): + """API Endpoint to get the known model types and their mapping.""" + try: + # Dynamically list all first-level folders under the main models directory + models_dir = getattr(folder_paths, 'models_dir', None) + if not models_dir: + base = getattr(folder_paths, 'base_path', os.getcwd()) + models_dir = os.path.join(base, 'models') + if not os.path.isdir(models_dir): + return web.json_response({}) + + entries = {} + for name in sorted(os.listdir(models_dir)): + p = os.path.join(models_dir, name) + if os.path.isdir(p): + entries[name] = name + return web.json_response(entries) + except Exception as e: + print(f"Error getting model types: {e}") + return web.json_response({"error": "Internal Server Error", "details": str(e), "status_code": 500}, status=500) diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/GetStatus.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetStatus.py new file mode 100644 index 0000000000000000000000000000000000000000..35e88d973e1598c23ce03f1bf7a1b86cd8b36188 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/GetStatus.py @@ -0,0 +1,19 @@ +# ================================================ +# File: server/routes/GetStatus.py +# ================================================ +from aiohttp import web +import server # ComfyUI server instance +from ...downloader.manager import manager as download_manager + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.get("/api/huggingface/status") +async def route_get_status(request): + """API Endpoint to get the status of downloads.""" + try: + status = download_manager.get_status() + return web.json_response(status) + except Exception as e: + print(f"Error getting download status: {e}") + # Format error response consistently + return web.json_response({"error": "Internal Server Error", "details": f"Failed to get status: {str(e)}", "status_code": 500}, status=500) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/OpenPath.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/OpenPath.py new file mode 100644 index 0000000000000000000000000000000000000000..6845881b61b233664198f6b41ee1101abb111881 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/OpenPath.py @@ -0,0 +1,58 @@ +# ================================================ +# File: server/routes/OpenPath.py +# ================================================ +import asyncio +import json +from aiohttp import web + +import server # ComfyUI server instance +from ...downloader.manager import manager as download_manager + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/open_path") +async def route_open_path(request): + """API Endpoint to open the containing folder of a completed download.""" + if not download_manager: + return web.json_response({"error": "Download Manager not initialized"}, status=500) + + try: + data = await request.json() + download_id = data.get("download_id") + + if not download_id: + return web.json_response({"error": "Missing 'download_id'", "details": "The request body must contain the 'download_id' of the completed item."}, status=400) + + print(f"[API Route /huggingface/open_path] Received open path request for ID: {download_id}") + # Call manager method in thread + result = await asyncio.to_thread(download_manager.open_containing_folder, download_id) + + status_code = 200 if result.get("success") else 404 if "not found" in result.get("error", "").lower() else 400 # Use 400 for OS error, security etc + # Check for specific errors to return better codes + if not result.get("success"): + error_lower = result.get("error", "").lower() + if "directory does not exist" in error_lower or "id not found" in error_lower: + status_code = 404 + elif "cannot open path" in error_lower or "unsupported os" in error_lower or "failed to open" in error_lower or "xdg-open" in error_lower: + status_code = 501 # Not Implemented / Failed on server side + elif "must be 'completed'" in error_lower: + status_code = 409 # Conflict - wrong state + else: + status_code = 400 # Bad Request / general failure + + # Prevent sensitive path info leakage in error messages by default + if not result.get("success") and "error" in result and status_code != 200: + print(f"[API Route /huggingface/open_path] Error for ID {download_id}: {result['error']}") # Log full error on server + # Optionally sanitize error sent to client + # if "Directory:" in result["error"] or "Path:" in result["error"]: + # result["error"] = "Server failed to open the specified directory." + + return web.json_response(result, status=status_code) + + except json.JSONDecodeError: + return web.json_response({"error": "Invalid JSON body"}, status=400) + except Exception as e: + import traceback + print(f"Error handling /huggingface/open_path request for ID '{data.get('download_id', 'N/A')}': {e}") + # traceback.print_exc() + return web.json_response({"error": "Internal Server Error", "details": f"An unexpected error occurred: {str(e)}"}, status=500) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/RetryDownload.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/RetryDownload.py new file mode 100644 index 0000000000000000000000000000000000000000..6e0d41b96c8148ae8b635c91a1296c29071a07d1 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/RetryDownload.py @@ -0,0 +1,39 @@ +# ================================================ +# File: server/routes/RetryDownload.py +# ================================================ +import asyncio +import json +from aiohttp import web + +import server # ComfyUI server instance +from ...downloader.manager import manager as download_manager + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/retry") +async def route_retry_download(request): + """API Endpoint to retry a failed/cancelled download.""" + if not download_manager: + return web.json_response({"error": "Download Manager not initialized"}, status=500) + + try: + data = await request.json() + download_id = data.get("download_id") + + if not download_id: + return web.json_response({"error": "Missing 'download_id'", "details": "The request body must contain the 'download_id' of the item to retry."}, status=400) + + print(f"[API Route /huggingface/retry] Received retry request for ID: {download_id}") + # Call manager method (which handles locking) + result = await asyncio.to_thread(download_manager.retry_download, download_id) # Run sync manager method in thread + + status_code = 200 if result.get("success") else 404 if "not found" in result.get("error", "").lower() else 400 + return web.json_response(result, status=status_code) + + except json.JSONDecodeError: + return web.json_response({"error": "Invalid JSON body"}, status=400) + except Exception as e: + import traceback + print(f"Error handling /huggingface/retry request for ID '{data.get('download_id', 'N/A')}': {e}") + # traceback.print_exc() # Uncomment for detailed logs + return web.json_response({"error": "Internal Server Error", "details": f"An unexpected error occurred: {str(e)}"}, status=500) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/SearchModels.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/SearchModels.py new file mode 100644 index 0000000000000000000000000000000000000000..fd64b1e9c76b8b482751b8119df570f57d9ff3a6 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/SearchModels.py @@ -0,0 +1,121 @@ +# ================================================ +# File: server/routes/SearchModels.py +# ================================================ +import json +import traceback +from aiohttp import web + +import server # ComfyUI server instance +from ..utils import get_request_json, resolve_huggingface_api_key +from ...config import HUGGINGFACE_API_TYPE_MAP + +prompt_server = server.PromptServer.instance + +@prompt_server.routes.post("/api/huggingface/search") +async def route_search_models(request): + """API Endpoint for searching models using huggingface_hub.""" + try: + data = await get_request_json(request) + + query = data.get("query", "").strip() + model_type_keys = data.get("model_types", []) # e.g., ["lora", "checkpoint"] + base_model_filters = data.get("base_models", []) # e.g., ["SD 1.5", "Pony"] + sort = data.get("sort", "Most Downloaded") + limit = int(data.get("limit", 20)) + page = int(data.get("page", 1)) + resolved_api_key = resolve_huggingface_api_key(data) + nsfw = data.get("nsfw", True) # Default to True (enabled) + + if not query and not model_type_keys and not base_model_filters: + raise web.HTTPBadRequest(reason="Search requires a query or at least one filter (type or base model).") + + # Use huggingface_hub for search + try: + from huggingface_hub import HfApi + hf_api = HfApi(token=resolved_api_key) + + # Prepare search parameters + search_params = { + "limit": limit + } + + # Map sort values to valid huggingface_hub values + sort_mapping = { + "Most Downloaded": "downloads", + "Most Liked": "likes", + "Newest": "created_at", + "Relevancy": None, # Remove sort for relevancy (default) + "Alphabetical": "modelId" + } + + sort_value = sort_mapping.get(sort, None) + if sort_value: + search_params["sort"] = sort_value + + if query: + search_params["search"] = query + + # Map model types to tags + if model_type_keys: + tags = [] + for key in model_type_keys: + api_type = HUGGINGFACE_API_TYPE_MAP.get(key.lower()) + if api_type: + tags.append(api_type.lower()) + if tags: + search_params["tags"] = tags + + # Add base model filters as part of query if specified + if base_model_filters: + base_model_query = " ".join([f'base_model:{bm}' for bm in base_model_filters]) + if query: + search_params["search"] = f"{query} {base_model_query}" + else: + search_params["search"] = base_model_query + + print(f"[Server Search] huggingface_hub: search='{search_params.get('search', '')}', tags={search_params.get('tags', 'Any')}, sort={sort}, limit={limit}") + + # Perform search + models = hf_api.list_models(**search_params) + + # Format results for frontend + formatted_models = [] + for model in models: + formatted_model = { + "id": model.id, + "name": model.id.split('/')[-1], + "description": model.tags or "", + "creator": {"username": model.author or ""}, + "modelId": model.id, + "downloads": model.downloads or 0, + "likes": model.likes or 0, + "tags": model.tags or [], + "modelType": "Unknown", # Will be determined by frontend + "baseModel": [], # Will be determined by frontend + "stats": { + "downloadCount": model.downloads or 0, + "thumbsUpCount": model.likes or 0 + } + } + formatted_models.append(formatted_model) + + response_data = { + "items": formatted_models, + "metadata": { + "currentPage": page, + "pageSize": limit, + "totalItems": len(formatted_models), + "totalPages": 1 # huggingface_hub doesn't provide pagination info + } + } + + return web.json_response(response_data) + + except Exception as e: + print(f"[Server Search] huggingface_hub search failed: {e}") + return web.json_response({"error": "Search failed", "details": str(e)}, status=500) + + except Exception as e: + print(f"Error in search_models: {e}") + traceback.print_exc() + return web.json_response({"error": "Internal Server Error", "details": str(e)}, status=500) diff --git a/custom_nodes/ComfyUI-HuggingFace/server/routes/__init__.py b/custom_nodes/ComfyUI-HuggingFace/server/routes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..201a4a7650198f41767e1a2a8ec40dfce9242fea --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/routes/__init__.py @@ -0,0 +1,21 @@ +# ================================================ +# File: server/routes/__init__.py +# ================================================ +# This file imports all the individual route modules. +# When the `routes` package is imported by `server/__init__.py`, +# these imports will be executed, registering the routes with the +# ComfyUI server instance. + +from . import CancelDownload +from . import ClearHistory +from . import DownloadModel +from . import GetBaseModels +from . import GetModelDetails +from . import GetModelTypes +from . import GetModelDirs +from . import GetStatus +from . import OpenPath +from . import RetryDownload +from . import SearchModels + +print("[HuggingFace] All server route modules loaded.") diff --git a/custom_nodes/ComfyUI-HuggingFace/server/utils.py b/custom_nodes/ComfyUI-HuggingFace/server/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7c19e117802387f5c8c0212019be676ea7f2eb1f --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/server/utils.py @@ -0,0 +1,187 @@ +# ================================================ +# File: server/utils.py +# ================================================ +import json +import os +from typing import Any, Dict, Optional +from aiohttp import web + +# Import necessary components from our modules +from ..api.huggingface import HuggingFaceAPI +from ..utils.helpers import parse_huggingface_input + +async def get_request_json(request): + """Safely get JSON data from request.""" + try: + return await request.json() + except Exception as e: + print(f"Error parsing request JSON: {e}") + raise web.HTTPBadRequest(reason=f"Invalid JSON format: {e}") + +def resolve_huggingface_api_key(payload: Optional[Dict[str, Any]] = None) -> Optional[str]: + """ + Resolve API key priority: + 1) Explicit key from request payload (`api_key`) + 2) `HUGGINGFACE_TOKEN` environment variable + """ + request_key = "" + if isinstance(payload, dict): + raw_key = payload.get("api_key", "") + if isinstance(raw_key, str): + request_key = raw_key.strip() + + if request_key: + return request_key + + env_key = os.getenv("HUGGINGFACE_TOKEN", "").strip() + if env_key: + return env_key + + return None + +async def get_huggingface_model_and_version_details(api: HuggingFaceAPI, model_url_or_id: str, req_version_id: Optional[int]) -> Dict[str, Any]: + """ + Helper to fetch HuggingFace details. + Prioritizes fetching model info based on resolved Model ID. + Fetches specific version info if version ID is provided/resolved, otherwise latest. + Returns a dict with 'model_info', 'version_info', 'primary_file', and resolved IDs. + Raises HTTP exceptions on critical failures. + """ + target_model_id = None + target_version_id = None + potential_version_id_from_input = None + model_info = {} + version_info_to_use = {} # The version (specific or latest) whose file we'll use + primary_file = None + + # --- 1. Parse Input to get potential IDs --- + parsed_model_id, parsed_version_id = parse_huggingface_input(model_url_or_id) + + # Determine the initial target model ID (input URL/ID takes precedence) + target_model_id = parsed_model_id + + # Determine the specific version requested (explicit param > URL param) + if req_version_id and str(req_version_id).isdigit(): + try: + potential_version_id_from_input = int(req_version_id) + except (ValueError, TypeError): + print(f"[API Helper] Warning: Invalid req_version_id: {req_version_id}. Ignoring.") + elif parsed_version_id: + potential_version_id_from_input = parsed_version_id + + # --- 2. Ensure we have a Model ID --- + # If we only got a version ID from the input (e.g., huggingface.com/model-versions/456), + # we need to fetch that version *first* just to find the model ID. + if not target_model_id and potential_version_id_from_input: + print(f"[API Helper] Input requires fetching version {potential_version_id_from_input} first to find model ID.") + temp_version_info = api.get_model_version_info(potential_version_id_from_input) + if temp_version_info and "error" not in temp_version_info and temp_version_info.get('modelId'): + target_model_id = temp_version_info['modelId'] + print(f"[API Helper] Found Model ID {target_model_id} from Version ID {potential_version_id_from_input}.") + # We might reuse temp_version_info later if this was the specifically requested version + else: + err = temp_version_info.get('details', 'Could not find model ID from version') if isinstance(temp_version_info, dict) else 'API error' + raise web.HTTPNotFound(reason=f"Could not determine Model ID from Version ID {potential_version_id_from_input}", body=json.dumps({"error": f"Version {potential_version_id_from_input} not found or missing modelId", "details": err})) + + # If still no model ID after potential lookup, fail + if not target_model_id: + raise web.HTTPBadRequest(reason="Could not determine a valid Model ID from the input.") + + # --- 3. Fetch Core Model Information (Always based on target_model_id) --- + print(f"[API Helper] Fetching core model info for Model ID: {target_model_id}") + model_info_result = api.get_model_info(target_model_id) + if not model_info_result or "error" in model_info_result: + err_details = model_info_result.get('details', 'Unknown API error') if isinstance(model_info_result, dict) else 'Unknown API error' + raise web.HTTPNotFound(reason=f"Model {target_model_id} not found or API error", body=json.dumps({"error": f"Model {target_model_id} not found or API error", "details": err_details})) + model_info = model_info_result # Store the successfully fetched model info + + # --- 4. Determine and Fetch Version Info for File Details --- + if potential_version_id_from_input: + # User specified a version explicitly, fetch its details + print(f"[API Helper] Fetching specific version info for Version ID: {potential_version_id_from_input}") + target_version_id = potential_version_id_from_input # This is the version we need info for + # Check if we already fetched this during Model ID lookup + if 'temp_version_info' in locals() and temp_version_info.get('id') == target_version_id: + print("[API Helper] Reusing version info fetched earlier.") + version_info_to_use = temp_version_info + else: + version_info_result = api.get_model_version_info(target_version_id) + if not version_info_result or "error" in version_info_result: + err_details = version_info_result.get('details', 'Unknown API error') if isinstance(version_info_result, dict) else 'Unknown API error' + raise web.HTTPNotFound(reason=f"Specified Version {target_version_id} not found or API error", body=json.dumps({"error": f"Version {target_version_id} not found or API error", "details": err_details})) + version_info_to_use = version_info_result + else: + # No specific version requested, find latest/default from model_info + print(f"[API Helper] Finding latest/default version for Model ID: {target_model_id}") + versions = model_info.get("modelVersions") + if not versions or not isinstance(versions, list) or len(versions) == 0: + raise web.HTTPNotFound(reason=f"Model {target_model_id} has no listed model versions.") + + # Find the 'best' default version (usually first published) + default_version_summary = next((v for v in versions if v.get('status') == 'Published'), versions[0]) + target_version_id = default_version_summary.get('id') + if not target_version_id: + raise web.HTTPNotFound(reason=f"Model {target_model_id}'s latest version has no ID.") + + print(f"[API Helper] Using latest/default Version ID: {target_version_id}. Fetching its full details.") + # Fetch full details for this latest version + version_info_result = api.get_model_version_info(target_version_id) + if not version_info_result or "error" in version_info_result: + # Log error, but maybe try to proceed with summary data if desperate? Risky. + err_details = version_info_result.get('details', 'Unknown error getting full version') if isinstance(version_info_result, dict) else 'Error' + print(f"[API Helper] Warning: Could not fetch full details for latest version {target_version_id}. Details: {err_details}. Falling back to summary.") + # Use summary data from model_info as fallback - file info might be missing! + version_info_to_use = default_version_summary + # Ensure minimal structure for file finding later + version_info_to_use['files'] = version_info_to_use.get('files', []) + version_info_to_use['images'] = version_info_to_use.get('images', []) + version_info_to_use['modelId'] = version_info_to_use.get('modelId', target_model_id) # Ensure modelId is present + version_info_to_use['model'] = version_info_to_use.get('model', {'name': model_info.get('name', 'Unknown')}) # Add fallback model name + + else: + version_info_to_use = version_info_result + + # --- 5. Find Primary File from the Determined Version (version_info_to_use) --- + print(f"[API Helper] Finding primary file for Version ID: {target_version_id}") + files = version_info_to_use.get("files", []) + if not isinstance(files, list): files = [] + + # Handle fallback downloadUrl at version level if 'files' is empty/missing + if not files and 'downloadUrl' in version_info_to_use and version_info_to_use['downloadUrl']: + print("[API Helper] Warning: No 'files' array found, using version-level 'downloadUrl'.") + files = [{ + "id": None, "name": version_info_to_use.get('name', f"version_{target_version_id}_file"), + "primary": True, "type": "Model", "sizeKB": version_info_to_use.get('fileSizeKB'), + "downloadUrl": version_info_to_use['downloadUrl'], "hashes": {}, "metadata": {} + }] + + if not files: + raise web.HTTPNotFound(reason=f"Version {target_version_id} (Name: {version_info_to_use.get('name', 'N/A')}) has no files listed.") + + # For HuggingFace, just pick the first file or largest .safetensors + primary_file = None + if files and isinstance(files, list): + # Try to find .safetensors first + for file_info in files: + if (file_info.get("type") == "file" and + file_info.get("path", "").endswith(".safetensors")): + primary_file = file_info + break + + # If no .safetensors, pick the largest file + if not primary_file: + primary_file = max(files, key=lambda x: x.get("size", 0), default=None) + + if not primary_file: + raise web.HTTPNotFound(reason=f"Could not find any usable file with a download URL for version {target_version_id}.") + + print(f"[API Helper] Selected file: Name='{primary_file.get('name', 'N/A')}', SizeKB={primary_file.get('sizeKB')}") + + # --- 6. Return Results --- + return { + "model_info": model_info, # Always the full model info + "version_info": version_info_to_use, # Info for the specific/latest version + "primary_file": primary_file, # The file from that version + "target_model_id": target_model_id, # Resolved model ID + "target_version_id": target_version_id, # Resolved version ID (specific or latest) + } diff --git a/custom_nodes/ComfyUI-HuggingFace/utils/__init__.py b/custom_nodes/ComfyUI-HuggingFace/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3a431f5ad358b0bc5ba310102d196fa2b0832607 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/utils/__init__.py @@ -0,0 +1,15 @@ +# ================================================ +# File: utils/__init__.py +# ================================================ + +# Import utility functions for easy access +from .helpers import parse_huggingface_input, get_model_dir, sanitize_filename +from . import helpers + +# Make functions available at package level +__all__ = [ + 'parse_huggingface_input', + 'get_model_dir', + 'sanitize_filename', + 'helpers' +] \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/utils/helpers.py b/custom_nodes/ComfyUI-HuggingFace/utils/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..55e908d80325083e5d7a171de42f64e7663175a5 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/utils/helpers.py @@ -0,0 +1,224 @@ +# ================================================ +# File: utils/helpers.py +# ================================================ +import os +import urllib.parse +import re +from pathlib import Path +from typing import Optional, List, Dict, Any + +import folder_paths + +# Import config values needed here +from ..config import PLUGIN_ROOT, MODEL_TYPE_DIRS + +# Canonical aliases for model type/folder names. Values are preferred folder names. +MODEL_TYPE_ALIASES = { + "checkpoint": "checkpoints", + "checkpoints": "checkpoints", + "diffusionmodel": "diffusion_models", + "diffusionmodels": "diffusion_models", + "diffusion_model": "diffusion_models", + "diffusion_models": "diffusion_models", # Wan 2.2 and similar go here + "diffusers": "diffusers", + "unet": "unet", # GGUF models go here + "lora": "loras", + "loras": "loras", + "locon": "loras", + "lycoris": "loras", + "vae": "vae", + "embedding": "embeddings", + "embeddings": "embeddings", + "textualinversion": "embeddings", + "hypernetwork": "hypernetworks", + "hypernetworks": "hypernetworks", + "controlnet": "controlnet", + "upscaler": "upscale_models", + "upscalers": "upscale_models", + "upscale_model": "upscale_models", + "upscale_models": "upscale_models", + "motionmodule": "motion_models", + "motionmodules": "motion_models", + "motion_model": "motion_models", + "motion_models": "motion_models", +} + +MODEL_TYPE_ALIASES_COMPACT = { + re.sub(r'[^a-z0-9]', '', k): v for k, v in MODEL_TYPE_ALIASES.items() +} + +def _normalize_model_type(model_type: str) -> str: + """Normalize model type string to canonical form.""" + if not model_type: + return "other" + + normalized = model_type.lower().strip() + + # Try exact match first + if normalized in MODEL_TYPE_ALIASES: + return MODEL_TYPE_ALIASES[normalized] + + # Try compact match (remove non-alphanumeric chars) + compact = re.sub(r'[^a-z0-9]', '', normalized) + if compact in MODEL_TYPE_ALIASES_COMPACT: + return MODEL_TYPE_ALIASES_COMPACT[compact] + + return "other" + +def get_model_dir(model_type: str, explicit_save_root: str = "", selected_subdir: str = "") -> Optional[str]: + """Get the directory path for a given model type.""" + try: + # Normalize the model type + normalized_type = _normalize_model_type(model_type) + + # Use explicit root if provided + if explicit_save_root: + base_path = explicit_save_root + # If selected_subdir is provided, append it + if selected_subdir: + base_path = os.path.join(base_path, selected_subdir) + return base_path + + # Try ComfyUI's folder_paths first + try: + if normalized_type in ["checkpoints", "loras", "vae", "embeddings", "hypernetworks", "controlnet", "upscale_models"]: + return folder_paths.get_folder_paths(normalized_type)[0] + elif normalized_type in ["diffusion_models", "motion_models", "unet", "diffusers"]: + # These might not be standard ComfyUI types, try to get them + try: + return folder_paths.get_folder_paths(normalized_type)[0] + except: + # Fallback: try unet for diffusion_models + if normalized_type == "diffusion_models": + try: + return folder_paths.get_folder_paths("diffusion_models")[0] + except: + # If diffusion_models doesn't exist, try unet + try: + return folder_paths.get_folder_paths("unet")[0] + except: + # Fallback to checkpoints + return folder_paths.get_folder_paths("checkpoints")[0] + elif normalized_type == "unet": + try: + return folder_paths.get_folder_paths("unet")[0] + except: + # Fallback to checkpoints + return folder_paths.get_folder_paths("checkpoints")[0] + else: + # For other types, try diffusion_models first + try: + return folder_paths.get_folder_paths("diffusion_models")[0] + except: + # Fallback to checkpoints + return folder_paths.get_folder_paths("checkpoints")[0] + else: + # For other types, use our extension directory instead of custom_nodes + from ..config import PLUGIN_ROOT + return os.path.join(PLUGIN_ROOT, "other_models") + + except: + # Fallback to base_path + type + return os.path.join(folder_paths.base_path, normalized_type) + + except Exception as e: + print(f"Error getting model directory for {model_type}: {e}") + return None + +def sanitize_filename(filename: str) -> str: + """Sanitize filename for safe file system usage.""" + if not filename: + return "unnamed_file" + + # Remove invalid characters + sanitized = re.sub(r'[<>:"/\\|?*]', '_', filename) + + # Remove control characters + sanitized = re.sub(r'[\x00-\x1f\x7f]', '', sanitized) + + # Limit length + if len(sanitized) > 255: + name, ext = os.path.splitext(sanitized) + sanitized = name[:255-len(ext)] + ext + + return sanitized.strip() + +def parse_huggingface_input(url_or_id: str) -> tuple[str | None, str | None]: + """ + Parses HuggingFace URL or ID string. + Returns: (model_id, filename) tuple. Both can be None. + Handles URLs like: + - https://huggingface.co/FX-FeiHou/wan2.2-Remix/resolve/main/NSFW/Wan2.2_Remix_NSFW_i2v_14b_high_lighting_fp8_e4m3fn_v2.1.safetensors + - FX-FeiHou/wan2.2-Remix + """ + if not url_or_id: + return None, None + + url_or_id = str(url_or_id).strip() + model_id: str | None = None + filename: str | None = None + + # Check if it's a direct download URL + if "/resolve/main/" in url_or_id: + try: + parsed_url = urllib.parse.urlparse(url_or_id) + print(f"[DEBUG] Parsed URL: {parsed_url}") + print(f"[DEBUG] Path parts: {parsed_url.path.split('/')}") + + # Extract model ID from path + path_parts = parsed_url.path.split('/') + print(f"[DEBUG] Path parts length: {len(path_parts)}") + print(f"[DEBUG] Path parts[2]: {path_parts[2] if len(path_parts) > 2 else 'None'}") + + if len(path_parts) >= 4 and path_parts[3] == "resolve": + # URL format: /FX-FeiHou/wan2.2-Remix/resolve/main/NSFW/file.safetensors + model_id = "/".join(path_parts[1:3]) # FX-FeiHou/wan2.2-Remix + filename = "/".join(path_parts[4:]) # main/NSFW/file.safetensors + # Remove "main/" from filename if it's duplicated + if filename.startswith("main/"): + filename = filename[5:] # Remove "main/" prefix + print(f"[DEBUG] Extracted model_id: {model_id}") + print(f"[DEBUG] Extracted filename: {filename}") + print(f"Parsed HF download URL - Model: {model_id}, File: {filename}") + return model_id, filename + else: + print(f"[DEBUG] Path does not match expected format") + except Exception as e: + print(f"Warning: Could not parse HF download URL '{url_or_id}': {e}") + return None, None + + # Check if it's a simple model ID (username/repo format) + if "/" in url_or_id and not url_or_id.startswith("http"): + parts = url_or_id.split("/") + if len(parts) >= 2: + model_id = "/".join(parts[0:2]) + print(f"Parsed HF model ID: {model_id}") + return model_id, None + + # Try parsing as full URL + try: + parsed_url = urllib.parse.urlparse(url_or_id) + + if "huggingface.co" in parsed_url.netloc.lower(): + path_parts = parsed_url.path.split('/') + if len(path_parts) >= 3: + model_id = "/".join(path_parts[1:3]) # Extract username/repo + print(f"Parsed HF URL - Model: {model_id}") + return model_id, None + except Exception as e: + print(f"Warning: Could not parse HF URL '{url_or_id}': {e}") + + print(f"Input '{url_or_id}' is not a recognizable HF model ID or URL.") + return None, None + +def get_model_folder_paths(model_type: str) -> List[str]: + """Get all folder paths for a given model type.""" + try: + normalized_type = _normalize_model_type(model_type) + return folder_paths.get_folder_paths(normalized_type) + except: + return [] + +def get_model_type_folder_name(model_type: str) -> str: + """Get the standard folder name for a model type.""" + return _normalize_model_type(model_type) diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/api/huggingface.js b/custom_nodes/ComfyUI-HuggingFace/web/js/api/huggingface.js new file mode 100644 index 0000000000000000000000000000000000000000..ab4744b4f3c8842efdc41a1f77127b8fad9c37d6 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/api/huggingface.js @@ -0,0 +1,164 @@ +// API client for HuggingFace UI +// Wraps ComfyUI's fetchApi with consistent error handling + +import { api } from "../../../../scripts/api.js"; + +export class HuggingFaceDownloaderAPI { + static async _request(endpoint, options = {}) { + try { + const url = endpoint.startsWith("/") ? endpoint : `/${endpoint}`; + const response = await api.fetchApi(url, options); + + if (!response.ok) { + let errorData; + const status = response.status; + const statusText = response.statusText; + try { + errorData = await response.json(); + if (typeof errorData !== "object" || errorData === null) { + errorData = { detail: String(errorData) }; + } + } catch (_) { + const detailText = await response.text().catch(() => `Status ${status} - Could not read error text`); + errorData = { + error: `HTTP error ${status}`, + details: String(detailText).substring(0, 500), + }; + } + const err = new Error(errorData.error || errorData.reason || `HTTP Error: ${status} ${statusText}`); + err.details = errorData.details || errorData.detail || errorData.error || "No details provided."; + err.status = status; + throw err; + } + + if (response.status === 204 || response.headers.get("Content-Length") === "0") { + return null; + } + return await response.json(); + } catch (error) { + if (!error.details) error.details = error.message; + throw error; + } + } + + static async downloadModel(params) { + return await this._request("/huggingface/download", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(params), + }); + } + + static async getModelDetails(params) { + return await this._request("/huggingface/get_model_details", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(params), + }); + } + + static async getStatus() { + return await this._request("/huggingface/status"); + } + + static async cancelDownload(downloadId) { + return await this._request("/huggingface/cancel", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ download_id: downloadId }), + }); + } + + static async searchModels(params) { + return await this._request("/huggingface/search", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(params), + }); + } + + static async getBaseModels() { + return await this._request("/huggingface/base_models"); + } + + static async getModelTypes() { + return await this._request("/huggingface/model_types"); + } + + static async getModelDirs(modelType) { + const q = encodeURIComponent(modelType || 'checkpoints'); + return await this._request(`/huggingface/model_dirs?type=${q}`); + } + + static async createModelDir(modelType, newDir, root = "") { + return await this._request("/huggingface/create_dir", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model_type: modelType, new_dir: newDir, root }), + }); + } + + static async createModelType(name) { + return await this._request("/huggingface/create_model_type", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name }), + }); + } + + static async getModelRoots(modelType) { + const q = encodeURIComponent(modelType || 'checkpoints'); + return await this._request(`/huggingface/model_roots?type=${q}`); + } + + static async createModelRoot(modelType, absPath) { + return await this._request("/huggingface/create_root", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model_type: modelType, path: absPath }), + }); + } + + static async getGlobalRoot() { + return await this._request("/huggingface/global_root"); + } + + static async setGlobalRoot(path) { + return await this._request("/huggingface/global_root", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ path }), + }); + } + + static async clearGlobalRoot() { + return await this._request("/huggingface/global_root/clear", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({}), + }); + } + + static async retryDownload(downloadId) { + return await this._request("/huggingface/retry", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ download_id: downloadId }), + }); + } + + static async openPath(downloadId) { + return await this._request("/huggingface/open_path", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ download_id: downloadId }), + }); + } + + static async clearHistory() { + return await this._request("/huggingface/clear_history", { + method: "POST", + headers: { "Content-Type": "application/json" }, + }); + } +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/huggingfaceDownloader.css b/custom_nodes/ComfyUI-HuggingFace/web/js/huggingfaceDownloader.css new file mode 100644 index 0000000000000000000000000000000000000000..469916386d543323e87acfc7e261b39a302f8f0e --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/huggingfaceDownloader.css @@ -0,0 +1,753 @@ +/* HuggingFace Button Styling - Distinct from Civicomfy */ +#huggingface-downloader-button { + background: linear-gradient(135deg, #FFD700 0%, #FFA500 100%) !important; + color: #000 !important; + border: 2px solid #FF8C00 !important; + font-weight: bold !important; + padding: 4px 12px !important; + border-radius: 6px !important; + transition: all 0.3s ease !important; +} + +#huggingface-downloader-button:hover { + background: linear-gradient(135deg, #FFA500 0%, #FF8C00 100%) !important; + transform: translateY(-1px) !important; + box-shadow: 0 4px 8px rgba(255, 140, 0, 0.3) !important; +} + +#huggingface-downloader-button:active { + transform: translateY(0) !important; + box-shadow: 0 2px 4px rgba(255, 140, 0, 0.3) !important; +} + +.huggingface-downloader-modal { + position: fixed; + z-index: 1001; /* Above ComfyUI elements */ + left: 0; + top: 0; + width: 100%; + height: 100%; + overflow: hidden; /* Prevent body scroll */ + background-color: rgba(0, 0, 0, 0.6); + display: flex; + justify-content: center; + align-items: center; + opacity: 0; + visibility: hidden; + transition: opacity 0.3s ease, visibility 0s linear 0.3s; +} + +.huggingface-downloader-modal.open { + opacity: 1; + visibility: visible; + transition: opacity 0.3s ease, visibility 0s linear 0s; +} + +.huggingface-downloader-modal-content { + background-color: var(--comfy-menu-bg); + color: var(--comfy-text-color); + margin: auto; + padding: 0; /* Remove padding, handle internally */ + border-radius: 8px; + box-shadow: 0 5px 15px rgba(0, 0, 0, 0.3); + width: 900px; + max-width: 95%; + height: 700px; /* Fixed height */ + max-height: 90vh; + display: flex; /* Use flexbox for layout */ + flex-direction: column; /* Stack header, tabs, content */ + overflow: hidden; /* Prevent content overflow */ +} + +.huggingface-downloader-header { + display: flex; + justify-content: space-between; + align-items: center; + padding: 15px 20px; + border-bottom: 1px solid var(--border-color, #444); + flex-shrink: 0; /* Prevent header from shrinking */ +} + +.huggingface-downloader-header h2 { + margin: 0; + font-size: 1.3em; +} + +.huggingface-close-button { + background: none; + border: none; + color: var(--comfy-text-color); + font-size: 28px; + cursor: pointer; + padding: 0 5px; + line-height: 1; +} +.huggingface-close-button:hover { + color: #aaa; +} + +.huggingface-downloader-body { + display: flex; + flex-direction: column; + flex-grow: 1; /* Allow body to take remaining space */ + overflow: hidden; /* Manage internal scrolling */ +} + +.huggingface-downloader-tabs { + display: flex; + border-bottom: 1px solid var(--border-color, #444); + padding: 0 15px; + flex-shrink: 0; /* Prevent tabs from shrinking */ +} + +.huggingface-downloader-tab { + padding: 10px 18px; + cursor: pointer; + border: none; + background: none; + color: var(--comfy-text-color); + opacity: 0.7; + position: relative; + top: 1px; /* Align with bottom border */ + margin-bottom: -1px; /* Overlap border */ +} + +.huggingface-downloader-tab.active { + opacity: 1; + border-bottom: 3px solid var(--accent-color, #5c8aff); + font-weight: bold; +} +.huggingface-downloader-tab:hover { + opacity: 1; +} + +.huggingface-downloader-tab-content { + display: none; + padding: 20px; + flex-grow: 1; /* Allow content to take space */ + overflow-y: auto; /* Enable scrolling ONLY for the content area */ +} + +.huggingface-downloader-tab-content.active { + display: block; /* Show active tab */ +} + +.huggingface-form-group { + margin-bottom: 18px; +} +.huggingface-form-group:last-child { + margin-bottom: 0; +} + +.huggingface-form-group label { + display: block; + margin-bottom: 6px; + font-weight: bold; + font-size: 0.95em; +} + +.huggingface-input, +.huggingface-select { + width: 100%; + padding: 10px; + background-color: var(--comfy-input-bg); + color: var(--comfy-text-color); + border: 1px solid var(--border-color, #555); + border-radius: 4px; + box-sizing: border-box; /* Include padding/border in width */ +} +.huggingface-input:focus, +.huggingface-select:focus { + outline: none; + border-color: var(--accent-color, #5c8aff); + box-shadow: 0 0 0 2px rgba(92, 138, 255, 0.3); +} + +.huggingface-form-row { + display: flex; + gap: 15px; + flex-wrap: wrap; /* Allow wrapping on smaller screens */ +} +.huggingface-form-row > .huggingface-form-group { + flex: 1; /* Distribute space */ + min-width: 180px; /* Minimum width before wrapping */ +} +/* Adjust for checkbox/inline elements */ +.huggingface-form-group.inline { + display: flex; + align-items: center; + gap: 8px; + margin-bottom: 10px; /* Less margin for checkboxes */ +} +.huggingface-form-group.inline label { + margin-bottom: 0; /* Remove default bottom margin */ + order: 1; /* Put label after checkbox */ +} +.huggingface-checkbox { + width: auto; + order: 0; /* Put checkbox before label */ + accent-color: var(--accent-color, #5c8aff); + transform: scale(1.1); +} + +.huggingface-button { + background-color: var(--comfy-input-bg); + color: var(--comfy-text-color); + border: 1px solid var(--border-color, #555); + padding: 10px 18px; + border-radius: 4px; + cursor: pointer; + font-size: 1em; + transition: background-color 0.2s ease, border-color 0.2s ease; +} +.huggingface-button:hover { + background-color: var(--comfy-input-bg-hover); + border-color: #777; +} +.huggingface-button:disabled { + opacity: 0.5; + cursor: not-allowed; +} + +.huggingface-button.primary { + background-color: var(--accent-color, #5c8aff); + border-color: var(--accent-color, #5c8aff); + color: white; +} +.huggingface-button.primary:hover { + background-color: #4a7ee0; /* Slightly darker blue */ + border-color: #4a7ee0; +} +.huggingface-button.primary:disabled { + background-color: var(--accent-color, #5c8aff); /* Keep color but lower opacity */ + border-color: var(--accent-color, #5c8aff); +} + +.huggingface-button.danger { + background-color: #e64b4b; + border-color: #e64b4b; + color: white; +} +.huggingface-button.danger:hover { + background-color: #d93a3a; + border-color: #d93a3a; +} +.huggingface-button.small { + padding: 5px 10px; + font-size: 0.85em; +} + +/* Status Tab Styling */ +.huggingface-status-section { + margin-bottom: 25px; +} +.huggingface-status-section h3 { + margin-top: 0; + margin-bottom: 15px; + border-bottom: 1px solid var(--border-color, #444); + padding-bottom: 8px; + font-size: 1.1em; +} + +.huggingface-download-list { + display: flex; + flex-direction: column; + gap: 15px; +} + +.huggingface-download-item { + display: flex; + gap: 15px; + align-items: flex-start; /* Align items top */ + padding: 15px; + border-radius: 6px; + background-color: var(--comfy-input-bg); /* Slightly different bg for items */ + border: 1px solid var(--border-color, #555); +} + +.huggingface-download-thumbnail { + width: 80px; + height: 80px; + object-fit: cover; + border-radius: 4px; + background-color: #333; /* Placeholder color */ + flex-shrink: 0; /* Don't shrink thumbnail */ + margin-top: 3px; /* Align better with text */ +} + +.huggingface-download-info { + flex-grow: 1; /* Allow info to take available space */ + display: flex; + flex-direction: column; + gap: 5px; + overflow: hidden; /* Prevent long text overflow */ +} +.huggingface-download-info strong { + font-size: 1.05em; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} +.huggingface-download-info p { + margin: 0; + font-size: 0.9em; + color: #ccc; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} +.huggingface-download-info .filename { + font-style: italic; + color: #bbb; +} +.huggingface-download-info .error-message { + color: #ff6b6b; + font-size: 0.85em; + white-space: normal; /* Allow error message to wrap */ + word-break: break-word; +} + +.huggingface-download-actions { + display: flex; + align-items: center; + gap: 8px; + margin-left: auto; /* Push actions to the right */ + flex-shrink: 0; /* Don't shrink action buttons */ +} + +.huggingface-progress-container { + margin-top: 8px; + width: 100%; + background-color: var(--comfy-menu-bg); /* Darker background for contrast */ + border-radius: 4px; + overflow: hidden; + height: 20px; + border: 1px solid var(--border-color, #555); +} + +.huggingface-progress-bar { + height: 100%; + background-color: var(--accent-color, #5c8aff); + text-align: center; + color: white; + line-height: 20px; /* Center text vertically */ + font-size: 12px; + font-weight: bold; + transition: width 0.3s ease-out; + min-width: 20px; /* Show something even for 0% */ + width: 0%; /* Start at 0 */ + white-space: nowrap; + overflow: hidden; +} +.huggingface-progress-bar.completed { + background-color: #4caf50; /* Green for completed */ +} +.huggingface-progress-bar.failed { + background-color: #f44336; /* Red for failed */ +} +.huggingface-progress-bar.cancelled { + background-color: #9e9e9e; /* Grey for cancelled */ +} + +.huggingface-speed-indicator { + font-size: 0.85em; + color: #ccc; + margin-top: 4px; + text-align: right; +} + +/* Search Tab Styling */ +.huggingface-search-controls { + display: flex; + gap: 15px; + margin-bottom: 20px; + flex-wrap: wrap; +} +.huggingface-search-controls .huggingface-input { + flex-grow: 1; /* Let search input take more space */ +} +.huggingface-search-controls .huggingface-select { + min-width: 150px; /* Min width for dropdowns */ +} +.huggingface-search-controls .huggingface-button { + align-self: flex-end; /* Align button with bottom of inputs */ +} + +.huggingface-search-results { + margin-top: 20px; +} +.huggingface-search-item { + display: flex; + gap: 15px; + align-items: flex-start; + margin-bottom: 15px; + padding: 15px; + border-radius: 6px; + background-color: var(--comfy-input-bg); + border: 1px solid var(--border-color, #555); + transition: background-color 0.2s ease; +} + +.huggingface-thumbnail-container { + position: relative; /* Required for absolute positioning of the child badge */ + display: block; + line-height: 0; + width: 120px; + height: 170px; + flex-shrink: 0; +} + +.huggingface-search-thumbnail { + width: 120px; /* Slightly larger for search */ + height: 170px; + object-fit: cover; + border-radius: 4px; + background-color: #333; + flex-shrink: 0; +} + +.huggingface-search-meta-info { + font-size: 0.85em; + color: #aaa; + margin-bottom: 5px; + display: flex; + gap: 15px; + flex-wrap: wrap; +} + +.huggingface-search-meta-info span { + display: inline-flex; + align-items: center; + gap: 5px; +} + +.huggingface-search-meta-info i { + color: #ccc; +} +.huggingface-type-badge { + position: absolute; + bottom: 5px; /* Adjust spacing from bottom */ + left: 5px; /* Adjust spacing from left */ + background-color: rgb(48, 94, 70); /* Black background with transparency */ + color: white; + padding: 10px 8px; /* Adjust padding (top/bottom, left/right) */ + border-radius: 10px; /* Adjust for desired roundness */ + font-size: 0.75em; /* Adjust font size */ + font-weight: bold; + z-index: 2; /* Ensure it's above the image */ + white-space: nowrap; /* Prevent text wrapping */ +} + +.huggingface-search-info { + flex-grow: 1; + display: flex; + flex-direction: column; + gap: 5px; + overflow: hidden; +} +.huggingface-search-info h4 { + margin: 0; + font-size: 1.1em; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} +.huggingface-search-info p { + margin: 0; + font-size: 0.9em; + color: #ccc; +} +.huggingface-search-tags { + display: flex; + flex-wrap: wrap; + gap: 5px; + margin-top: 8px; +} +.huggingface-search-tag { + background-color: rgba(255, 255, 255, 0.1); + color: #eee; + padding: 3px 8px; + border-radius: 10px; + font-size: 0.8em; +} +.huggingface-search-stats span { + margin-right: 15px; + font-size: 0.85em; + color: #bbb; +} +.huggingface-search-stats i { /* If using icons */ + margin-right: 4px; +} + +.huggingface-search-actions { + display: flex; + flex-direction: column; /* Stack buttons vertically */ + gap: 8px; + align-items: flex-end; /* Align buttons to the right */ + margin-left: auto; + flex-shrink: 0; +} + +.version-buttons-container { + display: flex; + flex-direction: column; + align-items: flex-end; + gap: 5px; +} + +.all-versions-container { + display: flex; + flex-direction: column; + gap: 5px; + padding-bottom: 5px; + white-space: nowrap; /* Prevent wrapping */ + max-width: 100%; /* Prevent excessive width */ + align-items: flex-end; +} + +.show-all-versions-button { + align-self: flex-end; /* Align button to left */ + margin-top: auto; /* Push to the bottom */ +} +.base-model-badge { + display: inline-block; /* Allows padding and background */ + background-color: #3a3a3a; /* Or a specific purple hex code like #800080 */ + color: white; + font-weight: bold; + padding: 2px 6px; /* Adjust padding for desired size (vertical horizontal) */ + border-radius: 3px; /* Adjust for desired roundness */ + margin-right: 4px; /* Optional: Adds some space between badge and hyphen */ + line-height: 1; /* Optional: Can help contain the background better */ + vertical-align: middle; /* Optional: Helps align the badge nicely with text */ +} + + +/* Settings Tab Styling */ +.huggingface-settings-container { + display: grid; + grid-template-columns: 1fr; /* Single column */ + gap: 20px; +} + +@media (min-width: 768px) { + .huggingface-settings-container { + grid-template-columns: repeat(auto-fit, minmax(300px, 1fr)); /* Responsive columns */ + } +} + +.huggingface-settings-section { + background-color: var(--comfy-input-bg); + padding: 15px; + border-radius: 6px; + border: 1px solid var(--border-color, #555); +} +.huggingface-settings-section h4 { + margin-top: 0; + margin-bottom: 15px; + border-bottom: 1px solid var(--border-color, #444); + padding-bottom: 8px; +} + +.huggingface-toast { + position: fixed; + bottom: 20px; + left: 50%; + transform: translateX(-50%); + background-color: rgba(0, 0, 0, 0.8); + color: white; + padding: 12px 20px; + border-radius: 6px; + z-index: 1005; /* Above modal */ + font-size: 0.95em; + opacity: 0; + transition: opacity 0.3s ease, bottom 0.3s ease; + pointer-events: none; /* Don't block clicks */ +} + +.huggingface-toast.show { + opacity: 1; + bottom: 30px; +} + +.huggingface-toast.success { + background-color: rgba(76, 175, 80, 0.9); /* Green */ +} + +.huggingface-toast.error { + background-color: rgba(244, 67, 54, 0.9); /* Red */ +} + +.huggingface-toast.info { + background-color: rgba(33, 150, 243, 0.9); /* Blue */ +} + +/* Tooltip Styles */ +[data-tooltip] { + position: relative; + cursor: help; +} +[data-tooltip]::after { + content: attr(data-tooltip); + position: absolute; + left: 50%; + transform: translateX(-50%); + bottom: 100%; + margin-bottom: 5px; + background-color: rgba(0, 0, 0, 0.85); + color: white; + padding: 5px 10px; + border-radius: 4px; + font-size: 0.85em; + white-space: nowrap; + opacity: 0; + visibility: hidden; + transition: opacity 0.2s ease, visibility 0.2s ease; + z-index: 1010; /* Ensure tooltip is on top */ +} +[data-tooltip]:hover::after { + opacity: 1; + visibility: visible; +} + +.huggingface-download-preview-area { + /* Add specific margins/padding if needed */ + min-height: 50px; /* Ensure it has some height for the loading spinner */ +} + +/* Style for the description box */ +.model-description-content { + max-height: 200px; /* Limit height */ + overflow-y: auto; /* Allow scrolling */ + background-color: var(--comfy-input-bg); /* Match input background */ + padding: 10px; + border-radius: 4px; + font-size: 0.9em; /* Slightly smaller text */ + border: 1px solid var(--border-color, #555); + color: var(--comfy-text-color); /* Ensure text color matches */ + line-height: 1.5; /* Improve readability */ +} + +/* Style links within the description */ +.model-description-content a { + color: var(--accent-color, #5c8aff); /* Make links stand out */ + text-decoration: underline; +} +.model-description-content a:hover { + color: #8bb0ff; /* Lighter color on hover */ +} + +/* Ensure reused search item styles look okay in this context */ +.huggingface-download-preview-area .huggingface-search-item { + margin-bottom: 15px; /* Add space below the item */ + background-color: var(--comfy-input-bg); /* Explicitly set background */ +} + +/* --- Confirmation Modal Styles --- */ +.huggingface-confirmation-modal { + display: none; /* Hidden by default */ + position: fixed; /* Stay in place */ + z-index: 1001; /* Ensure it's above the main modal */ + left: 0; + top: 0; + width: 100%; + height: 100%; + overflow: auto; /* Enable scroll if needed */ + background-color: rgba(0,0,0,0.6); /* Dim background */ + /* Use flexbox for easy centering */ + justify-content: center; + align-items: center; +} + +.huggingface-confirmation-modal-content { + background-color: var(--comfy-menu-bg, #282828); + margin: auto; + padding: 25px; + border: 1px solid var(--border-color, #444); + border-radius: 8px; + width: 80%; + max-width: 450px; /* Limit width */ + box-shadow: 0 5px 15px rgba(0,0,0,0.3); + color: var(--input-text, #ddd); +} + +.huggingface-confirmation-modal-content h4 { + margin-top: 0; + margin-bottom: 15px; + color: var(--desc-text, #eee); + border-bottom: 1px solid var(--border-color, #444); + padding-bottom: 10px; +} + +.huggingface-confirmation-modal-content p { + margin-bottom: 25px; + line-height: 1.5; + font-size: 0.95em; +} + +.huggingface-confirmation-modal-actions { + display: flex; + justify-content: flex-end; /* Align buttons to the right */ + gap: 10px; /* Space between buttons */ +} + +/* Style buttons within the confirmation modal */ +.huggingface-confirmation-modal-actions .huggingface-button { + padding: 8px 15px; /* Adjust padding */ + min-width: 80px; /* Ensure minimum width */ +} + +/* Ensure both img and video within the thumbnail container behave similarly */ +.huggingface-thumbnail-container img.huggingface-search-thumbnail, +.huggingface-thumbnail-container video.huggingface-search-thumbnail, +.huggingface-thumbnail-container img.huggingface-download-thumbnail, +.huggingface-thumbnail-container video.huggingface-download-thumbnail { + display: block; /* Prevent extra space below */ + width: 100%; /* Fill the container width */ + height: 100%; /* Fill the container height */ + object-fit: cover; /* Scale while maintaining aspect ratio, cropping if needed */ + background-color: #333; /* Background for loading/error states */ +} + +/* You might already have rules for the container itself, ensure they define a size */ +.huggingface-thumbnail-container { + width: 100px; /* Example size */ + height: 140px; /* Example size */ + overflow: hidden; /* Hide parts of the video/image outside the cover area */ + position: relative; /* For positioning badges etc. */ + flex-shrink: 0; /* Prevent shrinking in flex layouts */ + border-radius: 4px; /* Match styling */ +} + +/* Blur/Hide for R-rated thumbnails */ +.huggingface-thumbnail-container.blurred img.huggingface-search-thumbnail, +.huggingface-thumbnail-container.blurred video.huggingface-search-thumbnail, +.huggingface-thumbnail-container.blurred img.huggingface-download-thumbnail, +.huggingface-thumbnail-container.blurred video.huggingface-download-thumbnail { + filter: blur(14px) brightness(0.6) saturate(0.6); + pointer-events: none; /* Let clicks hit the container */ +} +.huggingface-thumbnail-container .huggingface-nsfw-overlay { + position: absolute; + inset: 0; + display: flex; + align-items: center; + justify-content: center; + color: #fff; + font-weight: 800; + font-size: 20px; + letter-spacing: 1px; + background: rgba(0,0,0,0.35); + z-index: 3; + user-select: none; +} +/* Make it obvious it’s clickable */ +.huggingface-thumbnail-container.blurred { cursor: pointer; } + +/* Style for the type badge (likely exists) */ +.huggingface-type-badge { + /* ... existing styles ... */ + position: absolute; + bottom: 5px; + right: 5px; + /* ... */ +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/huggingfaceDownloader.js b/custom_nodes/ComfyUI-HuggingFace/web/js/huggingfaceDownloader.js new file mode 100644 index 0000000000000000000000000000000000000000..34aae2f63dd29aebf2bc10b7f7f518941e188994 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/huggingfaceDownloader.js @@ -0,0 +1,89 @@ +import { app } from "../../../scripts/app.js"; +import { addCssLink } from "./utils/dom.js"; +import { HuggingFaceDownloaderUI } from "./ui/UI.js"; + +console.log("Loading HuggingFace UI..."); + +// --- Configuration --- +const EXTENSION_NAME = "HuggingFace"; +const CSS_URL = `../huggingfaceDownloader.css`; +const PLACEHOLDER_IMAGE_URL = `/extensions/ComfyUI-HuggingFace/images/placeholder.jpg`; + +// Add Menu Button to ComfyUI +function addMenuButton() { + const buttonGroup = document.querySelector(".comfyui-button-group"); + + if (!buttonGroup) { + console.warn(`[${EXTENSION_NAME}] ComfyUI button group not found. Retrying...`); + setTimeout(addMenuButton, 500); + return; + } + + if (document.getElementById("huggingface-downloader-button")) { + console.log(`[${EXTENSION_NAME}] Button already exists.`); + return; + } + + const huggingfaceButton = document.createElement("button"); + huggingfaceButton.innerHTML = "🤗 HuggingFace"; + huggingfaceButton.id = "huggingface-downloader-button"; + huggingfaceButton.title = "Open HuggingFace Model Downloader"; + + huggingfaceButton.onclick = async () => { + if (!window.huggingfaceDownloaderUI) { + console.info(`[${EXTENSION_NAME}] Creating HuggingFaceDownloaderUI instance...`); + window.huggingfaceDownloaderUI = new HuggingFaceDownloaderUI(); + document.body.appendChild(window.huggingfaceDownloaderUI.modal); + + try { + await window.huggingfaceDownloaderUI.initializeUI(); + console.info(`[${EXTENSION_NAME}] UI Initialization complete.`); + } catch (error) { + console.error(`[${EXTENSION_NAME}] Error during UI initialization:`, error); + window.huggingfaceDownloaderUI?.showToast("Error initializing UI components. Check console.", "error", 5000); + } + } + + if (window.huggingfaceDownloaderUI) { + window.huggingfaceDownloaderUI.openModal(); + } else { + console.error(`[${EXTENSION_NAME}] Cannot open modal: UI instance not available.`); + alert("HuggingFace failed to initialize. Please check the browser console for errors."); + } + }; + + buttonGroup.appendChild(huggingfaceButton); + console.log(`[${EXTENSION_NAME}] HuggingFace button added to .comfyui-button-group.`); + + const menu = document.querySelector(".comfy-menu"); + if (!buttonGroup.contains(huggingfaceButton) && menu && !menu.contains(huggingfaceButton)) { + console.warn(`[${EXTENSION_NAME}] Failed to append button to group, falling back to menu.`); + const settingsButton = menu.querySelector("#comfy-settings-button"); + if (settingsButton) { + settingsButton.insertAdjacentElement("beforebegin", huggingfaceButton); + } else { + menu.appendChild(huggingfaceButton); + } + } +} + +// --- Initialization --- +app.registerExtension({ + name: "ComfyUI-HuggingFace.HuggingFaceDownloader", + async setup(appInstance) { + console.log(`[${EXTENSION_NAME}] Setting up HuggingFace Extension...`); + addCssLink(CSS_URL); + addMenuButton(); + + // Optional: Pre-check placeholder image + fetch(PLACEHOLDER_IMAGE_URL) + .then(res => { + if (!res.ok) { + console.warn(`[${EXTENSION_NAME}] Placeholder image not found at ${PLACEHOLDER_IMAGE_URL}.`); + } + }) + .catch(err => console.warn(`[${EXTENSION_NAME}] Error checking for placeholder image:`, err)); + + console.log(`[${EXTENSION_NAME}] Extension setup complete. UI will initialize on first click.`); + }, +}); diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/UI.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/UI.js new file mode 100644 index 0000000000000000000000000000000000000000..a5a8dbe3bcd1ff8792858cf2311f7546904183d7 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/UI.js @@ -0,0 +1,465 @@ +import { Feedback } from "./feedback.js"; +import { setupEventListeners } from "./handlers/eventListeners.js"; +import { handleDownloadSubmit, fetchAndDisplayDownloadPreview, debounceFetchDownloadPreview } from "./handlers/downloadHandler.js"; +import { handleSearchSubmit } from "./handlers/searchHandler.js"; +import { + handleSettingsSave, + loadAndApplySettings, + loadSettingsFromCookie, + saveSettingsToCookie, + applySettings, + getDefaultSettings, + loadGlobalRootSetting, + handleSetGlobalRoot, + handleClearGlobalRoot, +} from "./handlers/settingsHandler.js"; +import { startStatusUpdates, stopStatusUpdates, updateStatus, handleCancelDownload, handleRetryDownload, handleOpenPath, handleClearHistory } from "./handlers/statusHandler.js"; +import { renderSearchResults } from "./searchRenderer.js"; +import { renderDownloadList } from "./statusRenderer.js"; +import { renderDownloadPreview } from "./previewRenderer.js"; +import { modalTemplate } from "./templates.js"; +import { HuggingFaceDownloaderAPI } from "../api/huggingface.js"; + +export class HuggingFaceDownloaderUI { + constructor() { + this.modal = null; + this.tabs = {}; + this.tabContents = {}; + this.activeTab = 'download'; + this.modelTypes = {}; + this.statusInterval = null; + this.statusData = { queue: [], active: [], history: [] }; + this.baseModels = []; + this.searchPagination = { currentPage: 1, totalPages: 1, limit: 20 }; + this.settings = this.getDefaultSettings(); + this.toastTimeout = null; + this.modelPreviewDebounceTimeout = null; + + this.updateStatus(); + this.buildModalHTML(); + this.cacheDOMElements(); + this.setupEventListeners(); + this.feedback = new Feedback(this.modal.querySelector('#huggingface-toast')); + // Ensure icon stylesheet is loaded so buttons render icons immediately + this.ensureFontAwesome(); + } + + // --- Core UI Methods --- + buildModalHTML() { + this.modal = document.createElement('div'); + this.modal.className = 'huggingface-downloader-modal'; + this.modal.id = 'huggingface-downloader-modal'; + this.modal.innerHTML = modalTemplate(this.settings); + } + + cacheDOMElements() { + this.closeButton = this.modal.querySelector('#huggingface-close-modal'); + this.tabContainer = this.modal.querySelector('.huggingface-downloader-tabs'); + + // Download Tab + this.downloadForm = this.modal.querySelector('#huggingface-download-form'); + this.downloadPreviewArea = this.modal.querySelector('#huggingface-download-preview-area'); + this.modelUrlInput = this.modal.querySelector('#huggingface-model-url'); + this.modelVersionIdInput = this.modal.querySelector('#huggingface-model-version-id'); + this.downloadModelTypeSelect = this.modal.querySelector('#huggingface-model-type'); + this.createModelTypeButton = this.modal.querySelector('#huggingface-create-model-type'); + this.customFilenameInput = this.modal.querySelector('#huggingface-custom-filename'); + this.subdirSelect = this.modal.querySelector('#huggingface-subdir-select'); + this.createSubdirButton = this.modal.querySelector('#huggingface-create-subdir'); + this.saveBasePathHint = this.modal.querySelector('#huggingface-save-base-path'); + this.downloadConnectionsInput = this.modal.querySelector('#huggingface-connections'); + this.forceRedownloadCheckbox = this.modal.querySelector('#huggingface-force-redownload'); + this.downloadSubmitButton = this.modal.querySelector('#huggingface-download-submit'); + + // Search Tab + this.searchForm = this.modal.querySelector('#huggingface-search-form'); + this.searchQueryInput = this.modal.querySelector('#huggingface-search-query'); + this.searchTypeSelect = this.modal.querySelector('#huggingface-search-type'); + this.searchBaseModelSelect = this.modal.querySelector('#huggingface-search-base-model'); + this.searchSortSelect = this.modal.querySelector('#huggingface-search-sort'); + this.searchPeriodSelect = this.modal.querySelector('#huggingface-search-period'); + this.searchSubmitButton = this.modal.querySelector('#huggingface-search-submit'); + this.searchResultsContainer = this.modal.querySelector('#huggingface-search-results'); + this.searchPaginationContainer = this.modal.querySelector('#huggingface-search-pagination'); + + // Status Tab + this.statusContent = this.modal.querySelector('#huggingface-status-content'); + this.activeListContainer = this.modal.querySelector('#huggingface-active-list'); + this.queuedListContainer = this.modal.querySelector('#huggingface-queued-list'); + this.historyListContainer = this.modal.querySelector('#huggingface-history-list'); + this.statusIndicator = this.modal.querySelector('#huggingface-status-indicator'); + this.activeCountSpan = this.modal.querySelector('#huggingface-active-count'); + this.clearHistoryButton = this.modal.querySelector('#huggingface-clear-history-button'); + this.confirmClearModal = this.modal.querySelector('#huggingface-confirm-clear-modal'); + this.confirmClearYesButton = this.modal.querySelector('#huggingface-confirm-clear-yes'); + this.confirmClearNoButton = this.modal.querySelector('#huggingface-confirm-clear-no'); + + // Settings Tab + this.settingsForm = this.modal.querySelector('#huggingface-settings-form'); + this.settingsApiKeyInput = this.modal.querySelector('#huggingface-settings-api-key'); + this.settingsGlobalRootInput = this.modal.querySelector('#huggingface-settings-global-root'); + this.settingsSetGlobalRootButton = this.modal.querySelector('#huggingface-settings-set-global-root'); + this.settingsClearGlobalRootButton = this.modal.querySelector('#huggingface-settings-clear-global-root'); + this.settingsConnectionsInput = this.modal.querySelector('#huggingface-settings-connections'); + this.settingsDefaultTypeSelect = this.modal.querySelector('#huggingface-settings-default-type'); + this.settingsAutoOpenCheckbox = this.modal.querySelector('#huggingface-settings-auto-open-status'); + this.settingsHideMatureCheckbox = this.modal.querySelector('#huggingface-settings-hide-mature'); + this.settingsNsfwThresholdInput = this.modal.querySelector('#huggingface-settings-nsfw-threshold'); + this.settingsSaveButton = this.modal.querySelector('#huggingface-settings-save'); + + // Toast Notification + this.toastElement = this.modal.querySelector('#huggingface-toast'); + + // Collect tabs and contents + this.tabs = {}; + this.modal.querySelectorAll('.huggingface-downloader-tab').forEach(tab => { + this.tabs[tab.dataset.tab] = tab; + }); + this.tabContents = {}; + this.modal.querySelectorAll('.huggingface-downloader-tab-content').forEach(content => { + const tabName = content.id.replace('huggingface-tab-', ''); + if (tabName) this.tabContents[tabName] = content; + }); + } + + async initializeUI() { + console.info("[HuggingFace] Initializing UI components..."); + await this.populateModelTypes(); + await this.populateBaseModels(); + this.loadAndApplySettings(); + await this.loadGlobalRootSetting(); + if (this.downloadModelTypeSelect) { + await this.loadAndPopulateSubdirs(this.downloadModelTypeSelect.value); + } + } + + async populateModelTypes() { + console.log("[HuggingFace] Populating model types..."); + try { + const types = await HuggingFaceDownloaderAPI.getModelTypes(); + if (!types || typeof types !== 'object' || Object.keys(types).length === 0) { + throw new Error("Received invalid model types data format."); + } + this.modelTypes = types; + const sortedTypes = Object.entries(this.modelTypes).sort((a, b) => a[1].localeCompare(b[1])); + + this.downloadModelTypeSelect.innerHTML = ''; + this.searchTypeSelect.innerHTML = ''; + this.settingsDefaultTypeSelect.innerHTML = ''; + + sortedTypes.forEach(([key, displayName]) => { + const option = document.createElement('option'); + option.value = key; + option.textContent = displayName; + this.downloadModelTypeSelect.appendChild(option.cloneNode(true)); + this.settingsDefaultTypeSelect.appendChild(option.cloneNode(true)); + this.searchTypeSelect.appendChild(option.cloneNode(true)); + }); + // After types are populated, load subdirs for the current selection + await this.loadAndPopulateSubdirs(this.downloadModelTypeSelect.value); + } catch (error) { + console.error("[HuggingFace] Failed to get or populate model types:", error); + this.showToast('Failed to load model types', 'error'); + this.downloadModelTypeSelect.innerHTML = ''; + this.modelTypes = { "checkpoints": "Checkpoints (Default)" }; + } + } + + async loadAndPopulateSubdirs(modelType) { + try { + const res = await HuggingFaceDownloaderAPI.getModelDirs(modelType); + const select = this.subdirSelect; + if (!select) return; + const current = select.value; + select.innerHTML = ''; + const optRoot = document.createElement('option'); + optRoot.value = ''; + optRoot.textContent = '(root)'; + select.appendChild(optRoot); + if (res && Array.isArray(res.subdirs)) { + // res.subdirs contains '' for root; skip empty since we added (root) + res.subdirs.filter(p => p && typeof p === 'string').forEach(rel => { + const opt = document.createElement('option'); + opt.value = rel; + opt.textContent = rel; + select.appendChild(opt); + }); + } + // Restore selection if still present + if (Array.from(select.options).some(o => o.value === current)) { + select.value = current; + } + if (this.saveBasePathHint) { + const basePath = (res && typeof res.base_dir === 'string') ? res.base_dir : ''; + this.saveBasePathHint.textContent = basePath ? `Base path: ${basePath}` : ''; + } + } catch (e) { + console.error('[HuggingFace] Failed to load subdirectories:', e); + if (this.subdirSelect) { + this.subdirSelect.innerHTML = ''; + } + if (this.saveBasePathHint) { + this.saveBasePathHint.textContent = ''; + } + } + } + + // (loadAndPopulateRoots removed; dynamic types already reflect models/ subfolders) + + async populateBaseModels() { + console.log("[HuggingFace] Populating base models..."); + try { + const result = await HuggingFaceDownloaderAPI.getBaseModels(); + if (!result || !Array.isArray(result.base_models)) { + throw new Error("Invalid base models data format received."); + } + this.baseModels = result.base_models.sort(); + const existingOptions = Array.from(this.searchBaseModelSelect.options); + existingOptions.slice(1).forEach(opt => opt.remove()); + this.baseModels.forEach(baseModelName => { + const option = document.createElement('option'); + option.value = baseModelName; + option.textContent = baseModelName; + this.searchBaseModelSelect.appendChild(option); + }); + } catch (error) { + console.error("[HuggingFace] Failed to get or populate base models:", error); + this.showToast('Failed to load base models list', 'error'); + } + } + + switchTab(tabId) { + if (this.activeTab === tabId || !this.tabs[tabId] || !this.tabContents[tabId]) return; + + this.tabs[this.activeTab]?.classList.remove('active'); + this.tabContents[this.activeTab]?.classList.remove('active'); + + this.tabs[tabId].classList.add('active'); + this.tabContents[tabId].classList.add('active'); + this.tabContents[tabId].scrollTop = 0; + this.activeTab = tabId; + + if (tabId === 'status') this.updateStatus(); + else if (tabId === 'settings') this.applySettings(); + else if(tabId === 'download') { + this.downloadConnectionsInput.value = this.settings.numConnections; + if (Object.keys(this.modelTypes).length > 0) { + this.downloadModelTypeSelect.value = this.settings.defaultModelType; + } + } + } + + // --- Modal Control --- + openModal() { + this.modal?.classList.add('open'); + document.body.style.setProperty('overflow', 'hidden', 'important'); + this.startStatusUpdates(); + if (this.activeTab === 'status') this.updateStatus(); + if (!this.settings.apiKey) this.switchTab('settings'); + } + + closeModal() { + this.modal?.classList.remove('open'); + document.body.style.removeProperty('overflow'); + this.stopStatusUpdates(); + } + + // --- Utility Methods --- + formatBytes(bytes, decimals = 2) { + if (bytes === null || bytes === undefined || isNaN(bytes)) return 'N/A'; + if (bytes === 0) return '0 Bytes'; + const k = 1024; + const dm = decimals < 0 ? 0 : decimals; + const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB']; + const i = Math.floor(Math.log(Math.abs(bytes)) / Math.log(k)); + return parseFloat((bytes / Math.pow(k, i)).toFixed(dm)) + ' ' + sizes[i]; + } + + formatSpeed(bytesPerSecond) { + if (!isFinite(bytesPerSecond) || bytesPerSecond <= 0) return ''; + return this.formatBytes(bytesPerSecond) + '/s'; + } + + formatDuration(isoStart, isoEnd) { + try { + const diffSeconds = Math.round((new Date(isoEnd) - new Date(isoStart)) / 1000); + if (isNaN(diffSeconds) || diffSeconds < 0) return 'N/A'; + if (diffSeconds < 60) return `${diffSeconds}s`; + const diffMinutes = Math.floor(diffSeconds / 60); + const remainingSeconds = diffSeconds % 60; + return `${diffMinutes}m ${remainingSeconds}s`; + } catch (e) { + return 'N/A'; + } + } + + showToast(message, type = 'info', duration = 3000) { + this.feedback?.show(message, type, duration); + } + + ensureFontAwesome() { + this.feedback?.ensureFontAwesome(); + } + + // --- Rendering (delegated to external renderers) --- + renderDownloadList = (items, container, emptyMessage) => renderDownloadList(this, items, container, emptyMessage); + renderSearchResults = (items) => renderSearchResults(this, items); + renderDownloadPreview = (data) => renderDownloadPreview(this, data); + + // --- Auto-select model type based on HuggingFace model type --- + inferFolderFromHuggingFaceType(huggingfaceType) { + if (!huggingfaceType || typeof huggingfaceType !== 'string') return null; + const t = huggingfaceType.trim().toLowerCase(); + const keys = Object.keys(this.modelTypes || {}); + if (keys.length === 0) return null; + + const exists = (k) => keys.includes(k); + const findBy = (pred) => keys.find(pred); + + // Direct matches first + if (exists(t)) return t; + if (exists(`${t}s`)) return `${t}s`; + + // Common mappings from HuggingFace types to ComfyUI folders + const candidates = []; + const addIfExists = (k) => { if (exists(k)) candidates.push(k); }; + + switch (t) { + case 'checkpoint': + addIfExists('checkpoints'); + addIfExists('models'); + break; + case 'lora': case 'locon': case 'lycoris': + addIfExists('loras'); + break; + case 'vae': + addIfExists('vae'); + break; + case 'textualinversion': case 'embedding': case 'embeddings': + addIfExists('embeddings'); + break; + case 'hypernetwork': + addIfExists('hypernetworks'); + break; + case 'controlnet': + addIfExists('controlnet'); + break; + case 'unet': case 'unet2': + addIfExists('unet'); + break; + case 'diffusers': case 'diffusionmodels': case 'diffusion_models': case 'diffusion': + addIfExists('diffusers'); + addIfExists('diffusion_models'); + break; + case 'upscaler': case 'upscalers': + addIfExists('upscale_models'); + addIfExists('upscalers'); + break; + case 'motionmodule': + addIfExists('motion_models'); + break; + case 'poses': + addIfExists('poses'); + break; + case 'wildcards': + addIfExists('wildcards'); + break; + case 'onnx': + addIfExists('onnx'); + break; + } + if (candidates.length > 0) return candidates[0]; + + // Relaxed match: name contains type + const contains = findBy(k => k.toLowerCase().includes(t)); + if (contains) return contains; + + return null; + } + + async autoSelectModelTypeFromHuggingFace(huggingfaceType) { + try { + const folder = this.inferFolderFromHuggingFaceType(huggingfaceType); + if (!folder) return; + if (this.downloadModelTypeSelect && this.downloadModelTypeSelect.value !== folder) { + this.downloadModelTypeSelect.value = folder; + await this.loadAndPopulateSubdirs(folder); + // Reset subdir to root after auto-switch + if (this.subdirSelect) this.subdirSelect.value = ''; + } + } catch (e) { + console.warn('[HuggingFace] Auto-select model type failed:', e); + } + } + + renderSearchPagination(metadata) { + if (!this.searchPaginationContainer) return; + if (!metadata || metadata.totalPages <= 1) { + this.searchPaginationContainer.innerHTML = ''; + this.searchPagination = { ...this.searchPagination, ...metadata }; + return; + } + + this.searchPagination = { ...this.searchPagination, ...metadata }; + const { currentPage, totalPages, totalItems } = this.searchPagination; + + const createButton = (text, page, isDisabled = false, isCurrent = false) => { + const button = document.createElement('button'); + button.className = `huggingface-button small huggingface-page-button ${isCurrent ? 'primary active' : ''}`; + button.dataset.page = page; + button.disabled = isDisabled; + button.innerHTML = text; + button.type = 'button'; + return button; + }; + + const fragment = document.createDocumentFragment(); + fragment.appendChild(createButton('« Prev', currentPage - 1, currentPage === 1)); + + let startPage = Math.max(1, currentPage - 2); + let endPage = Math.min(totalPages, currentPage + 2); + + if (startPage > 1) fragment.appendChild(createButton('1', 1)); + if (startPage > 2) fragment.appendChild(document.createElement('span')).textContent = '...'; + + for (let i = startPage; i <= endPage; i++) { + fragment.appendChild(createButton(i, i, false, i === currentPage)); + } + + if (endPage < totalPages - 1) fragment.appendChild(document.createElement('span')).textContent = '...'; + if (endPage < totalPages) fragment.appendChild(createButton(totalPages, totalPages)); + + fragment.appendChild(createButton('Next »', currentPage + 1, currentPage === totalPages)); + + const info = document.createElement('div'); + info.className = 'huggingface-pagination-info'; + info.textContent = `Page ${currentPage} of ${totalPages} (${totalItems.toLocaleString()} models)`; + fragment.appendChild(info); + + this.searchPaginationContainer.innerHTML = ''; + this.searchPaginationContainer.appendChild(fragment); + } + + // --- Event Handlers and State Management (delegated to handlers) --- + setupEventListeners = () => setupEventListeners(this); + getDefaultSettings = () => getDefaultSettings(); + loadAndApplySettings = () => loadAndApplySettings(this); + loadSettingsFromCookie = () => loadSettingsFromCookie(this); + saveSettingsToCookie = () => saveSettingsToCookie(this); + applySettings = () => applySettings(this); + handleSettingsSave = () => handleSettingsSave(this); + loadGlobalRootSetting = () => loadGlobalRootSetting(this); + handleSetGlobalRoot = () => handleSetGlobalRoot(this); + handleClearGlobalRoot = () => handleClearGlobalRoot(this); + handleDownloadSubmit = () => handleDownloadSubmit(this); + handleSearchSubmit = () => handleSearchSubmit(this); + fetchAndDisplayDownloadPreview = () => fetchAndDisplayDownloadPreview(this); + debounceFetchDownloadPreview = (delay) => debounceFetchDownloadPreview(this, delay); + startStatusUpdates = () => startStatusUpdates(this); + stopStatusUpdates = () => stopStatusUpdates(this); + updateStatus = () => updateStatus(this); + handleCancelDownload = (downloadId) => handleCancelDownload(this, downloadId); + handleRetryDownload = (downloadId, button) => handleRetryDownload(this, downloadId, button); + handleOpenPath = (downloadId, button) => handleOpenPath(this, downloadId, button); + handleClearHistory = () => handleClearHistory(this); +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/feedback.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/feedback.js new file mode 100644 index 0000000000000000000000000000000000000000..727bad4c256ae4ccb9d4077f062526de8d6922e5 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/feedback.js @@ -0,0 +1,41 @@ +// Centralized feedback utilities: toasts and icon CSS + +export class Feedback { + constructor(toastElement) { + this.toastElement = toastElement || null; + this.toastTimeout = null; + } + + ensureFontAwesome() { + if (!document.getElementById('huggingface-fontawesome-link')) { + const faLink = document.createElement('link'); + faLink.id = 'huggingface-fontawesome-link'; + faLink.rel = 'stylesheet'; + faLink.href = 'https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/css/all.min.css'; + faLink.integrity = 'sha512-1ycn6IcaQQ40/MKBW2W4Rhis/DbILU74C1vSrLJxCq57o941Ym01SwNsOMqvEBFlcgUa6xLiPY/NS5R+E6ztJQ=='; + faLink.crossOrigin = 'anonymous'; + faLink.referrerPolicy = 'no-referrer'; + document.head.appendChild(faLink); + } + } + + show(message, type = 'info', duration = 3000) { + if (!this.toastElement) return; + if (this.toastTimeout) { + clearTimeout(this.toastTimeout); + this.toastTimeout = null; + } + const valid = ['info', 'success', 'error', 'warning']; + const toastType = valid.includes(type) ? type : 'info'; + + this.toastElement.textContent = message; + this.toastElement.className = 'huggingface-toast'; + this.toastElement.classList.add(toastType); + requestAnimationFrame(() => this.toastElement.classList.add('show')); + this.toastTimeout = setTimeout(() => { + this.toastElement.classList.remove('show'); + this.toastTimeout = null; + }, duration); + } +} + diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/downloadHandler.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/downloadHandler.js new file mode 100644 index 0000000000000000000000000000000000000000..08f42a667b48b1af4411f3cbe58af6bf4b1e5418 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/downloadHandler.js @@ -0,0 +1,106 @@ +import { HuggingFaceDownloaderAPI } from "../../api/huggingface.js"; + +export function debounceFetchDownloadPreview(ui, delay = 500) { + clearTimeout(ui.modelPreviewDebounceTimeout); + ui.modelPreviewDebounceTimeout = setTimeout(() => { + fetchAndDisplayDownloadPreview(ui); + }, delay); +} + +export async function fetchAndDisplayDownloadPreview(ui) { + const modelUrlOrId = ui.modelUrlInput.value.trim(); + const versionId = ui.modelVersionIdInput.value.trim(); + + if (!modelUrlOrId) { + ui.downloadPreviewArea.innerHTML = ''; + return; + } + + ui.downloadPreviewArea.innerHTML = '

Loading model details...

'; + ui.ensureFontAwesome(); + + const params = { + model_url_or_id: modelUrlOrId, + model_version_id: versionId ? parseInt(versionId, 10) : null, + api_key: ui.settings.apiKey + }; + + try { + const result = await HuggingFaceDownloaderAPI.getModelDetails(params); + if (result && result.success) { + ui.renderDownloadPreview(result); + // Auto-select model type save location based on HuggingFace model type + if (result.model_type) { + await ui.autoSelectModelTypeFromHuggingFace(result.model_type); + } + } else { + // Don't show error for missing details - just show neutral message + const message = result.details || result.error || 'Model details not available'; + ui.downloadPreviewArea.innerHTML = `

${message}

`; + } + } catch (error) { + // Don't show scary error messages - just neutral info + const message = 'Model details not available'; + console.info("Download Preview - details not available:", error); + ui.downloadPreviewArea.innerHTML = `

${message}

`; + } +} + +export async function handleDownloadSubmit(ui) { + ui.downloadSubmitButton.disabled = true; + ui.downloadSubmitButton.textContent = 'Starting...'; + + const modelUrlOrId = ui.modelUrlInput.value.trim(); + if (!modelUrlOrId) { + ui.showToast("Model URL or ID cannot be empty.", "error"); + ui.downloadSubmitButton.disabled = false; + ui.downloadSubmitButton.textContent = 'Start Download'; + return; + } + + // Subfolder comes from dropdown; filename is base name only + const selectedSubdir = ui.subdirSelect ? ui.subdirSelect.value.trim() : ''; + const userFilename = ui.customFilenameInput.value.trim(); + + const params = { + model_url_or_id: modelUrlOrId, + model_type: ui.downloadModelTypeSelect.value, + model_version_id: ui.modelVersionIdInput.value ? parseInt(ui.modelVersionIdInput.value, 10) : null, + custom_filename: userFilename, + subdir: selectedSubdir, + num_connections: parseInt(ui.downloadConnectionsInput.value, 10), + force_redownload: ui.forceRedownloadCheckbox.checked, + api_key: ui.settings.apiKey + }; + + const fileSelectEl = ui.modal.querySelector('#huggingface-file-select'); + if (fileSelectEl && fileSelectEl.value) { + const fid = parseInt(fileSelectEl.value, 10); + if (!Number.isNaN(fid)) params.file_id = fid; + } + + try { + const result = await HuggingFaceDownloaderAPI.downloadModel(params); + + if (result.status === 'queued') { + ui.showToast(`Download queued: ${result.details?.filename || 'Model'}`, 'success'); + if (ui.settings.autoOpenStatusTab) { + ui.switchTab('status'); + } else { + ui.updateStatus(); + } + } else if (result.status === 'exists' || result.status === 'exists_size_mismatch') { + ui.showToast(`${result.message}`, 'info', 4000); + } else { + console.warn("Unexpected success response from /huggingface/download:", result); + ui.showToast(`Unexpected status: ${result.status} - ${result.message || ''}`, 'info'); + } + } catch (error) { + const message = `Download failed: ${error.details || error.message || 'Unknown error'}`; + console.error("Download Submit Error:", error); + ui.showToast(message, 'error', 6000); + } finally { + ui.downloadSubmitButton.disabled = false; + ui.downloadSubmitButton.textContent = 'Start Download'; + } +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/eventListeners.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/eventListeners.js new file mode 100644 index 0000000000000000000000000000000000000000..70ac7fc699a54152c7d0967c602f6fe012f6b1be --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/eventListeners.js @@ -0,0 +1,248 @@ +import { HuggingFaceDownloaderAPI } from "../../api/huggingface.js"; +export function setupEventListeners(ui) { + // Modal close + ui.closeButton.addEventListener('click', () => ui.closeModal()); + ui.modal.addEventListener('click', (event) => { + if (event.target === ui.modal) ui.closeModal(); + }); + + // Tab switching + ui.tabContainer.addEventListener('click', (event) => { + if (event.target.matches('.huggingface-downloader-tab')) { + ui.switchTab(event.target.dataset.tab); + } + }); + + // --- FORMS --- + ui.downloadForm.addEventListener('submit', (event) => { + event.preventDefault(); + ui.handleDownloadSubmit(); + }); + + // Change of model type should refresh subdir list + ui.downloadModelTypeSelect.addEventListener('change', async () => { + await ui.loadAndPopulateSubdirs(ui.downloadModelTypeSelect.value); + }); + + // Create new model type folder (first-level under models/) + ui.createModelTypeButton.addEventListener('click', async () => { + const name = prompt('Enter new model type folder name (will be created under models/)'); + if (!name) return; + try { + const res = await HuggingFaceDownloaderAPI.createModelType(name); + if (res && res.success) { + await ui.populateModelTypes(); + ui.downloadModelTypeSelect.value = res.name; + await ui.loadAndPopulateSubdirs(res.name); + ui.showToast(`Created model type folder: ${res.name}`, 'success'); + } else { + ui.showToast(res?.error || 'Failed to create model type folder', 'error'); + } + } catch (e) { + ui.showToast(e.details || e.message || 'Error creating model type folder', 'error'); + } + }); + + // Create new subfolder under current model type + ui.createSubdirButton.addEventListener('click', async () => { + const type = ui.downloadModelTypeSelect.value; + const name = prompt('Enter new subfolder name (you can include nested paths like A/B):'); + if (!name) return; + try { + const res = await HuggingFaceDownloaderAPI.createModelDir(type, name); + if (res && res.success) { + await ui.loadAndPopulateSubdirs(type); + if (ui.subdirSelect) ui.subdirSelect.value = res.created || ''; + ui.showToast(`Created folder: ${res.created}`, 'success'); + } else { + ui.showToast(res?.error || 'Failed to create folder', 'error'); + } + } catch (e) { + ui.showToast(e.details || e.message || 'Error creating folder', 'error'); + } + }); + + ui.searchForm.addEventListener('submit', (event) => { + event.preventDefault(); + if (!ui.searchQueryInput.value.trim() && ui.searchTypeSelect.value === 'any' && ui.searchBaseModelSelect.value === 'any') { + ui.showToast("Please enter a search query or select a filter.", "error"); + if (ui.searchResultsContainer) ui.searchResultsContainer.innerHTML = '

Please enter a search query or select a filter.

'; + if (ui.searchPaginationContainer) ui.searchPaginationContainer.innerHTML = ''; + return; + } + ui.searchPagination.currentPage = 1; + ui.handleSearchSubmit(); + }); + + ui.settingsForm.addEventListener('submit', (event) => { + event.preventDefault(); + ui.handleSettingsSave(); + }); + if (ui.settingsSetGlobalRootButton) { + ui.settingsSetGlobalRootButton.addEventListener('click', () => { + ui.handleSetGlobalRoot(); + }); + } + if (ui.settingsClearGlobalRootButton) { + ui.settingsClearGlobalRootButton.addEventListener('click', () => { + ui.handleClearGlobalRoot(); + }); + } + + // Download form inputs + ui.modelUrlInput.addEventListener('input', () => ui.debounceFetchDownloadPreview()); + ui.modelUrlInput.addEventListener('paste', () => ui.debounceFetchDownloadPreview(0)); + ui.modelVersionIdInput.addEventListener('blur', () => ui.fetchAndDisplayDownloadPreview()); + + // --- DYNAMIC CONTENT LISTENERS (Event Delegation) --- + + // Status tab actions (Cancel/Retry/Open/Clear) and click-to-toggle blur on thumbs + ui.statusContent.addEventListener('click', (event) => { + const thumbContainer = event.target.closest('.huggingface-thumbnail-container'); + if (thumbContainer) { + const nsfwLevel = Number(thumbContainer.dataset.nsfwLevel ?? thumbContainer.getAttribute('data-nsfw-level')); + const threshold = Number(ui.settings?.nsfwBlurMinLevel ?? 4); + const enabled = ui.settings?.hideMatureInSearch === true; + if (enabled && Number.isFinite(nsfwLevel) && nsfwLevel >= threshold) { + if (thumbContainer.classList.contains('blurred')) { + thumbContainer.classList.remove('blurred'); + const overlay = thumbContainer.querySelector('.huggingface-nsfw-overlay'); + if (overlay) overlay.remove(); + } else { + thumbContainer.classList.add('blurred'); + if (!thumbContainer.querySelector('.huggingface-nsfw-overlay')) { + const ov = document.createElement('div'); + ov.className = 'huggingface-nsfw-overlay'; + ov.title = 'R-rated: click to reveal'; + ov.textContent = 'R'; + thumbContainer.appendChild(ov); + } + } + return; // consume + } + } + + const button = event.target.closest('button'); + if (!button) return; + + const downloadId = button.dataset.id; + if (downloadId) { + if (button.classList.contains('huggingface-cancel-button')) ui.handleCancelDownload(downloadId); + else if (button.classList.contains('huggingface-retry-button')) ui.handleRetryDownload(downloadId, button); + else if (button.classList.contains('huggingface-openpath-button')) ui.handleOpenPath(downloadId, button); + } else if (button.id === 'huggingface-clear-history-button') { + ui.confirmClearModal.style.display = 'flex'; + } + }); + + // Download preview click-to-toggle blur + ui.downloadPreviewArea.addEventListener('click', (event) => { + const thumbContainer = event.target.closest('.huggingface-thumbnail-container'); + if (thumbContainer) { + const nsfwLevel = Number(thumbContainer.dataset.nsfwLevel ?? thumbContainer.getAttribute('data-nsfw-level')); + const threshold = Number(ui.settings?.nsfwBlurMinLevel ?? 4); + const enabled = ui.settings?.hideMatureInSearch === true; + if (enabled && Number.isFinite(nsfwLevel) && nsfwLevel >= threshold) { + if (thumbContainer.classList.contains('blurred')) { + thumbContainer.classList.remove('blurred'); + const overlay = thumbContainer.querySelector('.huggingface-nsfw-overlay'); + if (overlay) overlay.remove(); + } else { + thumbContainer.classList.add('blurred'); + if (!thumbContainer.querySelector('.huggingface-nsfw-overlay')) { + const ov = document.createElement('div'); + ov.className = 'huggingface-nsfw-overlay'; + ov.title = 'R-rated: click to reveal'; + ov.textContent = 'R'; + thumbContainer.appendChild(ov); + } + } + } + } + }); + + // Search results actions, including click-to-toggle blur + ui.searchResultsContainer.addEventListener('click', (event) => { + const thumbContainer = event.target.closest('.huggingface-thumbnail-container'); + if (thumbContainer) { + const nsfwLevel = Number(thumbContainer.dataset.nsfwLevel ?? thumbContainer.getAttribute('data-nsfw-level')); + const threshold = Number(ui.settings?.nsfwBlurMinLevel ?? 4); + const enabled = ui.settings?.hideMatureInSearch === true; + if (enabled && Number.isFinite(nsfwLevel) && nsfwLevel >= threshold) { + if (thumbContainer.classList.contains('blurred')) { + thumbContainer.classList.remove('blurred'); + const overlay = thumbContainer.querySelector('.huggingface-nsfw-overlay'); + if (overlay) overlay.remove(); + } else { + thumbContainer.classList.add('blurred'); + if (!thumbContainer.querySelector('.huggingface-nsfw-overlay')) { + const ov = document.createElement('div'); + ov.className = 'huggingface-nsfw-overlay'; + ov.title = 'R-rated: click to reveal'; + ov.textContent = 'R'; + thumbContainer.appendChild(ov); + } + } + return; // Don't trigger other actions on this click + } + } + + const downloadButton = event.target.closest('.huggingface-search-download-button'); + if (downloadButton) { + event.preventDefault(); + const { modelId, versionId, modelType, creator, modelName } = downloadButton.dataset; + if (!modelId) { + ui.showToast("Error: Missing model ID for download.", "error"); + return; + } + const modelTypeInternalKey = Object.keys(ui.modelTypes).find(key => ui.modelTypes[key]?.toLowerCase() === modelType?.toLowerCase()) || ui.settings.defaultModelType; + + ui.modelUrlInput.value = modelId; + ui.modelVersionIdInput.value = versionId; + ui.customFilenameInput.value = modelName ? `${modelName.replace(/[^a-zA-Z0-9_-]/g, '_')}` : ''; + ui.forceRedownloadCheckbox.checked = false; + ui.downloadModelTypeSelect.value = modelTypeInternalKey; + + ui.switchTab('download'); + ui.showToast(`Filled download form for "${modelName || modelId}" by ${creator || 'Unknown'}.`, 'info', 4000); + ui.fetchAndDisplayDownloadPreview(); + return; + } + + const viewAllButton = event.target.closest('.show-all-versions-button'); + if (viewAllButton) { + const modelId = viewAllButton.dataset.modelId; + const versionsContainer = ui.searchResultsContainer.querySelector(`#all-versions-${modelId}`); + if (versionsContainer) { + const currentlyVisible = versionsContainer.style.display !== 'none'; + versionsContainer.style.display = currentlyVisible ? 'none' : 'flex'; + viewAllButton.innerHTML = currentlyVisible + ? `All versions (${viewAllButton.dataset.totalVersions}) ` + : `Show less `; + } + } + }); + + // Pagination + ui.searchPaginationContainer.addEventListener('click', (event) => { + const button = event.target.closest('.huggingface-page-button'); + if (button && !button.disabled) { + const page = parseInt(button.dataset.page, 10); + if (page && page !== ui.searchPagination.currentPage) { + ui.searchPagination.currentPage = page; + ui.handleSearchSubmit(); + } + } + }); + + // Confirmation Modal + ui.confirmClearYesButton.addEventListener('click', () => ui.handleClearHistory()); + ui.confirmClearNoButton.addEventListener('click', () => { + ui.confirmClearModal.style.display = 'none'; + }); + ui.confirmClearModal.addEventListener('click', (event) => { + if (event.target === ui.confirmClearModal) { + ui.confirmClearModal.style.display = 'none'; + } + }); +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/searchHandler.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/searchHandler.js new file mode 100644 index 0000000000000000000000000000000000000000..96d0a1573702fb50bdf924488d41ffe27502b475 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/searchHandler.js @@ -0,0 +1,39 @@ +import { HuggingFaceDownloaderAPI } from "../../api/huggingface.js"; + +export async function handleSearchSubmit(ui) { + ui.searchSubmitButton.disabled = true; + ui.searchSubmitButton.textContent = 'Searching...'; + ui.searchResultsContainer.innerHTML = '

Searching...

'; + ui.searchPaginationContainer.innerHTML = ''; + ui.ensureFontAwesome(); + + const params = { + query: ui.searchQueryInput.value.trim(), + model_types: ui.searchTypeSelect.value === 'any' ? [] : [ui.searchTypeSelect.value], + base_models: ui.searchBaseModelSelect.value === 'any' ? [] : [ui.searchBaseModelSelect.value], + sort: ui.searchSortSelect.value, + limit: ui.searchPagination.limit, + page: ui.searchPagination.currentPage, + api_key: ui.settings.apiKey, + }; + + try { + const response = await HuggingFaceDownloaderAPI.searchModels(params); + if (!response || !response.metadata || !Array.isArray(response.items)) { + console.error("Invalid search response structure:", response); + throw new Error("Received invalid data from search API."); + } + + ui.renderSearchResults(response.items); + ui.renderSearchPagination(response.metadata); + + } catch (error) { + const message = `Search failed: ${error.details || error.message || 'Unknown error'}`; + console.error("Search Submit Error:", error); + ui.searchResultsContainer.innerHTML = `

${message}

`; + ui.showToast(message, 'error'); + } finally { + ui.searchSubmitButton.disabled = false; + ui.searchSubmitButton.textContent = 'Search'; + } +} \ No newline at end of file diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/settingsHandler.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/settingsHandler.js new file mode 100644 index 0000000000000000000000000000000000000000..86a47173b53d502871f6ca93990a990250e2a37b --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/settingsHandler.js @@ -0,0 +1,160 @@ +import { setCookie, getCookie } from "../../utils/cookies.js"; +import { HuggingFaceDownloaderAPI } from "../../api/huggingface.js"; + +const SETTINGS_COOKIE_NAME = 'huggingfaceDownloaderSettings'; + +export function getDefaultSettings() { + return { + apiKey: '', + numConnections: 1, + defaultModelType: 'checkpoints', + autoOpenStatusTab: true, + searchResultLimit: 20, + hideMatureInSearch: true, + nsfwBlurMinLevel: 4, // Blur thumbnails with nsfwLevel >= this value + }; +} + +export function loadAndApplySettings(ui) { + ui.settings = ui.loadSettingsFromCookie(); + ui.applySettings(); +} + +export function loadSettingsFromCookie(ui) { + const defaults = ui.getDefaultSettings(); + const cookieValue = getCookie(SETTINGS_COOKIE_NAME); + + if (cookieValue) { + try { + const loadedSettings = JSON.parse(cookieValue); + return { ...defaults, ...loadedSettings }; + } catch (e) { + console.error("Failed to parse settings cookie:", e); + return defaults; + } + } + return defaults; +} + +export function saveSettingsToCookie(ui) { + try { + const settingsString = JSON.stringify(ui.settings); + setCookie(SETTINGS_COOKIE_NAME, settingsString, 365); + ui.showToast('Settings saved successfully!', 'success'); + } catch (e) { + console.error("Failed to save settings to cookie:", e); + ui.showToast('Error saving settings', 'error'); + } +} + +export function applySettings(ui) { + if (ui.settingsApiKeyInput) { + ui.settingsApiKeyInput.value = ui.settings.apiKey || ''; + } + if (ui.settingsConnectionsInput) { + ui.settingsConnectionsInput.value = Math.max(1, Math.min(16, ui.settings.numConnections || 1)); + } + if (ui.settingsDefaultTypeSelect) { + const desired = ui.settings.defaultModelType || 'checkpoints'; + ui.settingsDefaultTypeSelect.value = desired; + if (!ui.settingsDefaultTypeSelect.querySelector(`option[value="${ui.settingsDefaultTypeSelect.value}"]`)) { + const first = ui.settingsDefaultTypeSelect.querySelector('option'); + if (first) ui.settingsDefaultTypeSelect.value = first.value; + } + } + if (ui.settingsAutoOpenCheckbox) { + ui.settingsAutoOpenCheckbox.checked = ui.settings.autoOpenStatusTab === true; + } + if (ui.settingsHideMatureCheckbox) { + ui.settingsHideMatureCheckbox.checked = ui.settings.hideMatureInSearch === true; + } + if (ui.settingsNsfwThresholdInput) { + const val = Number(ui.settings.nsfwBlurMinLevel); + ui.settingsNsfwThresholdInput.value = Number.isFinite(val) ? val : 4; + } + if (ui.downloadConnectionsInput) { + ui.downloadConnectionsInput.value = Math.max(1, Math.min(16, ui.settings.numConnections || 1)); + } + if (ui.downloadModelTypeSelect && Object.keys(ui.modelTypes).length > 0) { + const desired = ui.settings.defaultModelType || 'checkpoints'; + ui.downloadModelTypeSelect.value = desired; + if (!ui.downloadModelTypeSelect.querySelector(`option[value="${ui.downloadModelTypeSelect.value}"]`)) { + const first = ui.downloadModelTypeSelect.querySelector('option'); + if (first) ui.downloadModelTypeSelect.value = first.value; + } + } + ui.searchPagination.limit = ui.settings.searchResultLimit || 20; +} + +export async function loadGlobalRootSetting(ui) { + if (!ui.settingsGlobalRootInput) return; + try { + const result = await HuggingFaceDownloaderAPI.getGlobalRoot(); + const globalRoot = (result && typeof result.global_root === 'string') ? result.global_root : ''; + ui.settingsGlobalRootInput.value = globalRoot; + } catch (e) { + console.warn("[HuggingFace] Failed to load global root setting:", e); + } +} + +export async function handleSetGlobalRoot(ui) { + if (!ui.settingsGlobalRootInput) return; + const path = ui.settingsGlobalRootInput.value.trim(); + if (!path) { + ui.showToast("Please enter a global root path first.", "error"); + return; + } + try { + const result = await HuggingFaceDownloaderAPI.setGlobalRoot(path); + const saved = (result && typeof result.global_root === 'string') ? result.global_root : path; + ui.settingsGlobalRootInput.value = saved; + ui.showToast("Global root updated.", "success"); + if (ui.downloadModelTypeSelect) { + await ui.loadAndPopulateSubdirs(ui.downloadModelTypeSelect.value); + } + } catch (e) { + ui.showToast(e.details || e.message || "Failed to set global root.", "error", 6000); + } +} + +export async function handleClearGlobalRoot(ui) { + if (!ui.settingsGlobalRootInput) return; + try { + await HuggingFaceDownloaderAPI.clearGlobalRoot(); + ui.settingsGlobalRootInput.value = ""; + ui.showToast("Global root cleared. Using default ComfyUI paths.", "success"); + if (ui.downloadModelTypeSelect) { + await ui.loadAndPopulateSubdirs(ui.downloadModelTypeSelect.value); + } + } catch (e) { + ui.showToast(e.details || e.message || "Failed to clear global root.", "error", 6000); + } +} + +export function handleSettingsSave(ui) { + const apiKey = ui.settingsApiKeyInput.value.trim(); + const numConnections = parseInt(ui.settingsConnectionsInput.value, 10); + const defaultModelType = ui.settingsDefaultTypeSelect.value; + const autoOpenStatusTab = ui.settingsAutoOpenCheckbox.checked; + const hideMatureInSearch = ui.settingsHideMatureCheckbox.checked; + const nsfwBlurMinLevel = Number(ui.settingsNsfwThresholdInput.value); + + if (isNaN(numConnections) || numConnections < 1 || numConnections > 16) { + ui.showToast("Invalid Default Connections (must be 1-16).", "error"); + return; + } + if (!ui.settingsDefaultTypeSelect.querySelector(`option[value="${defaultModelType}"]`)) { + ui.showToast("Invalid Default Model Type selected.", "error"); + return; + } + + ui.settings.apiKey = apiKey; + ui.settings.numConnections = numConnections; + ui.settings.defaultModelType = defaultModelType; + ui.settings.autoOpenStatusTab = autoOpenStatusTab; + ui.settings.hideMatureInSearch = hideMatureInSearch; + ui.settings.nsfwBlurMinLevel = (Number.isFinite(nsfwBlurMinLevel) && nsfwBlurMinLevel >= 0) ? Math.min(128, Math.round(nsfwBlurMinLevel)) : 4; + + ui.saveSettingsToCookie(); + ui.applySettings(); +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/statusHandler.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/statusHandler.js new file mode 100644 index 0000000000000000000000000000000000000000..9ebc554a0eb86196e880e5c58e7951c9e4bc6717 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/handlers/statusHandler.js @@ -0,0 +1,154 @@ +import { HuggingFaceDownloaderAPI } from "../../api/huggingface.js"; + +export function startStatusUpdates(ui) { + if (!ui.statusInterval) { + console.log("[HuggingFace] Starting status updates (every 3s)..."); + ui.updateStatus(); + ui.statusInterval = setInterval(() => ui.updateStatus(), 3000); + } +} + +export function stopStatusUpdates(ui) { + if (ui.statusInterval) { + clearInterval(ui.statusInterval); + ui.statusInterval = null; + console.log("[HuggingFace] Stopped status updates."); + } +} + +export async function updateStatus(ui) { + if (!ui.modal || !ui.modal.classList.contains('open')) return; + + try { + const newStatusData = await HuggingFaceDownloaderAPI.getStatus(); + if (!newStatusData || !Array.isArray(newStatusData.active) || !Array.isArray(newStatusData.queue) || !Array.isArray(newStatusData.history)) { + throw new Error("Invalid status data structure received from server."); + } + + const oldStateString = JSON.stringify(ui.statusData); + const newStateString = JSON.stringify(newStatusData); + + // Cache new state if it differs + if (oldStateString !== newStateString) { + ui.statusData = newStatusData; + } + + // Always keep counters in sync + const activeCount = ui.statusData.active.length + ui.statusData.queue.length; + ui.activeCountSpan.textContent = activeCount; + ui.statusIndicator.style.display = activeCount > 0 ? 'inline' : 'none'; + + // Always render when Status tab is active, even if data hasn't changed + if (ui.activeTab === 'status') { + ui.renderDownloadList(ui.statusData.active, ui.activeListContainer, 'No active downloads.'); + ui.renderDownloadList(ui.statusData.queue, ui.queuedListContainer, 'Download queue is empty.'); + ui.renderDownloadList(ui.statusData.history, ui.historyListContainer, 'No download history yet.'); + } + } catch (error) { + console.error("[HuggingFace] Failed to update status:", error); + if (ui.activeTab === 'status') { + const errorHtml = `

${error.details || error.message}

`; + if (ui.activeListContainer) ui.activeListContainer.innerHTML = errorHtml; + if (ui.queuedListContainer) ui.queuedListContainer.innerHTML = ''; + if (ui.historyListContainer) ui.historyListContainer.innerHTML = ''; + } + } +} + +export async function handleCancelDownload(ui, downloadId) { + const button = ui.modal.querySelector(`.huggingface-cancel-button[data-id="${downloadId}"]`); + if (button) { + button.disabled = true; + button.innerHTML = ''; + button.title = "Cancelling..."; + } + try { + const result = await HuggingFaceDownloaderAPI.cancelDownload(downloadId); + ui.showToast(result.message || `Cancellation requested for ${downloadId}`, 'info'); + ui.updateStatus(); + } catch (error) { + const message = `Cancel failed: ${error.details || error.message}`; + console.error("Cancel Download Error:", error); + ui.showToast(message, 'error'); + if (button) { + button.disabled = false; + button.innerHTML = ''; + button.title = "Cancel Download"; + } + } +} + +export async function handleRetryDownload(ui, downloadId, button) { + button.disabled = true; + button.innerHTML = ''; + button.title = "Retrying..."; + try { + const result = await HuggingFaceDownloaderAPI.retryDownload(downloadId); + if (result.success) { + ui.showToast(result.message || `Retry queued successfully!`, 'success'); + if (ui.settings.autoOpenStatusTab) ui.switchTab('status'); + else ui.updateStatus(); + } else { + ui.showToast(`Retry failed: ${result.details || result.error}`, 'error', 5000); + button.disabled = false; + button.innerHTML = ''; + button.title = "Retry Download"; + } + } catch (error) { + const message = `Retry failed: ${error.details || error.message}`; + console.error("Retry Download UI Error:", error); + ui.showToast(message, 'error', 5000); + button.disabled = false; + button.innerHTML = ''; + button.title = "Retry Download"; + } +} + +export async function handleOpenPath(ui, downloadId, button) { + const originalIcon = button.innerHTML; + button.disabled = true; + button.innerHTML = ''; + button.title = "Opening..."; + try { + const result = await HuggingFaceDownloaderAPI.openPath(downloadId); + if (result.success) { + ui.showToast(result.message || `Opened path successfully!`, 'success'); + } else { + ui.showToast(`Open path failed: ${result.details || result.error}`, 'error', 5000); + } + } catch (error) { + const message = `Open path failed: ${error.details || error.message}`; + console.error("Open Path UI Error:", error); + ui.showToast(message, 'error', 5000); + } finally { + button.disabled = false; + button.innerHTML = originalIcon; + button.title = "Open Containing Folder"; + } +} + +export async function handleClearHistory(ui) { + ui.confirmClearYesButton.disabled = true; + ui.confirmClearNoButton.disabled = true; + ui.confirmClearYesButton.textContent = 'Clearing...'; + + try { + const result = await HuggingFaceDownloaderAPI.clearHistory(); + if (result.success) { + ui.showToast(result.message || 'History cleared successfully!', 'success'); + ui.statusData.history = []; + ui.renderDownloadList(ui.statusData.history, ui.historyListContainer, 'No download history yet.'); + ui.confirmClearModal.style.display = 'none'; + } else { + ui.showToast(`Clear history failed: ${result.details || result.error}`, 'error', 5000); + } + } catch (error) { + const message = `Clear history failed: ${error.details || error.message}`; + console.error("Clear History UI Error:", error); + ui.showToast(message, 'error', 5000); + } finally { + ui.confirmClearYesButton.disabled = false; + ui.confirmClearNoButton.disabled = false; + ui.confirmClearYesButton.textContent = 'Confirm Clear'; + } +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/previewRenderer.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/previewRenderer.js new file mode 100644 index 0000000000000000000000000000000000000000..bb33018c184fdef0d986bd4f774f9843bac43bd4 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/previewRenderer.js @@ -0,0 +1,96 @@ +// Renders the download preview panel + +const PLACEHOLDER_IMAGE_URL = `/extensions/ComfyUI-HuggingFace/images/placeholder.jpg`; + +export function renderDownloadPreview(ui, data) { + if (!ui.downloadPreviewArea) return; + ui.ensureFontAwesome(); + + const modelId = data.model_id; + const modelName = data.model_name || 'Untitled Model'; + const creator = data.creator_username || 'Unknown Creator'; + const modelType = data.model_type || 'N/A'; + const versionName = data.version_name || 'N/A'; + const baseModel = data.base_model || 'N/A'; + const stats = data.stats || {}; + const descriptionHtml = data.description_html || '

No description.

'; + const version_description_html = data.version_description_html || '

No description.

'; + const fileInfo = data.file_info || {}; + const files = Array.isArray(data.files) ? data.files : []; + const thumbnail = data.thumbnail_url || PLACEHOLDER_IMAGE_URL; + const nsfwLevel = Number(data.nsfw_level ?? 0); + const blurMinLevel = Number(ui.settings?.nsfwBlurMinLevel ?? 4); + const shouldBlur = ui.settings?.hideMatureInSearch === true && nsfwLevel >= blurMinLevel; + const huggingfaceLink = `https://huggingface.com/models/${modelId}${data.version_id ? '?modelVersionId=' + data.version_id : ''}`; + + const onErrorScript = `this.onerror=null; this.src='${PLACEHOLDER_IMAGE_URL}'; this.style.backgroundColor='#444';`; + + const overlayHtml = shouldBlur ? `
R
` : ''; + const containerClasses = `huggingface-thumbnail-container${shouldBlur ? ' blurred' : ''}`; + + const previewHtml = ` +
+
+ ${modelName} thumbnail + ${overlayHtml} +
${modelType}
+
+
+

${modelName} by ${creator}

+

Version: ${versionName} ${baseModel}

+
+ ${stats.downloads?.toLocaleString() || 0} + ${stats.likes?.toLocaleString(0) || 0} + ${stats.dislikes?.toLocaleString() || 0} + ${stats.buzz?.toLocaleString() || 0} +
+

Primary File:

+

+ Name: ${fileInfo.name || 'N/A'}
+ Size: ${ui.formatBytes(fileInfo.size_kb * 1024) || 'N/A'}
+ Format: ${fileInfo.format || 'N/A'}
+ Precision: ${fileInfo.precision || 'N/A'}
+ Model Size: ${fileInfo.model_size || 'N/A'} +

+ ${files.length > 0 ? ` +
+ + +

Pick other variants when available.

+
+ ` : ''} + + View on HuggingFace + +
+
+
+
Model Description:
+
+ ${descriptionHtml} +
+
+
+
Version Description:
+
+ ${version_description_html} +
+
+ `; + + ui.downloadPreviewArea.innerHTML = previewHtml; +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/searchRenderer.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/searchRenderer.js new file mode 100644 index 0000000000000000000000000000000000000000..6dde4b27dd7cdde68f5ab30050794b650d6c250a --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/searchRenderer.js @@ -0,0 +1,196 @@ +// Rendering of search results list +// Usage: renderSearchResults(uiInstance, itemsArray) + +const PLACEHOLDER_IMAGE_URL = `/extensions/ComfyUI-HuggingFace/images/placeholder.jpg`; + +export function renderSearchResults(ui, items) { + ui.feedback?.ensureFontAwesome(); + + if (!items || items.length === 0) { + const queryUsed = ui.searchQueryInput && ui.searchQueryInput.value.trim(); + const typeFilterUsed = ui.searchTypeSelect && ui.searchTypeSelect.value !== 'any'; + const baseModelFilterUsed = ui.searchBaseModelSelect && ui.searchBaseModelSelect.value !== 'any'; + const message = (queryUsed || typeFilterUsed || baseModelFilterUsed) + ? 'No models found matching your criteria.' + : 'Enter a query or select filters and click Search.'; + ui.searchResultsContainer.innerHTML = `

${message}

`; + return; + } + + const placeholder = PLACEHOLDER_IMAGE_URL; + const onErrorScript = `this.onerror=null; this.src='${placeholder}'; this.style.backgroundColor='#444';`; + const fragment = document.createDocumentFragment(); + + items.forEach(hit => { + const modelId = hit.id; + if (!modelId) return; + + // Extract base repository ID from full path (remove /tree/main, /blob/main, etc.) + const baseRepoId = modelId.split('/tree/')[0].split('/blob/')[0].split('/raw/')[0]; + + const creator = hit.user?.username || 'Unknown Creator'; + const modelName = hit.name || 'Untitled Model'; + const modelTypeApi = hit.type || 'other'; + console.log('Model type for badge:', modelTypeApi); + const stats = hit.metrics || {}; + const tags = hit.tags?.map(t => t.name) || []; + + const thumbnailUrl = hit.thumbnailUrl || placeholder; + const firstImage = Array.isArray(hit.images) && hit.images.length > 0 ? hit.images[0] : null; + const thumbnailType = firstImage?.type; + const nsfwLevel = Number(firstImage?.nsfwLevel ?? hit.nsfwLevel ?? 0); + const blurMinLevel = Number(ui.settings?.nsfwBlurMinLevel ?? 4); + const shouldBlur = ui.settings?.hideMatureInSearch === true && nsfwLevel >= blurMinLevel; + + const allVersions = hit.versions || []; + const primaryVersion = hit.version || (allVersions.length > 0 ? allVersions[0] : {}); + const primaryVersionId = primaryVersion.id; + const primaryBaseModel = primaryVersion.baseModel || 'N/A'; + + const uniqueBaseModels = allVersions.length > 0 + ? [...new Set(allVersions.map(v => v.baseModel).filter(Boolean))] + : (primaryBaseModel !== 'N/A' ? [primaryBaseModel] : []); + const baseModelsDisplay = uniqueBaseModels.length > 0 ? uniqueBaseModels.join(', ') : 'N/A'; + + const publishedAt = hit.publishedAt; + let lastUpdatedFormatted = 'N/A'; + if (publishedAt) { + try { + const date = new Date(publishedAt); + lastUpdatedFormatted = date.toLocaleDateString(undefined, { year: 'numeric', month: 'short', day: 'numeric' }); + } catch (_) {} + } + + const listItem = document.createElement('div'); + listItem.className = 'huggingface-search-item'; + listItem.dataset.modelId = modelId; + + const MAX_VISIBLE_VERSIONS = 3; + let visibleVersions = []; + if (primaryVersionId) { + visibleVersions.push({ id: primaryVersionId, name: primaryVersion.name || 'Primary Version', baseModel: primaryBaseModel }); + } + allVersions.forEach(v => { + if (v.id !== primaryVersionId && visibleVersions.length < MAX_VISIBLE_VERSIONS) visibleVersions.push(v); + }); + + let versionButtonsHtml = visibleVersions.map(version => { + const versionId = version.id; + const versionName = version.name || 'Unknown Version'; + const baseModel = version.baseModel || 'N/A'; + return ` + + `; + }).join(''); + + const hasMoreVersions = allVersions.length > visibleVersions.length; + const totalVersionCount = allVersions.length; + const moreButtonHtml = hasMoreVersions ? ` + + ` : ''; + + let allVersionsHtml = ''; + if (hasMoreVersions) { + const hiddenVersions = allVersions.filter(v => !visibleVersions.some(vis => vis.id === v.id)); + allVersionsHtml = ` + + `; + } + + let thumbnailHtml = ''; + const videoTitle = `Video preview for ${modelName}`; + const imageAlt = `${modelName} thumbnail`; + if (thumbnailUrl && typeof thumbnailUrl === 'string' && thumbnailType === 'video') { + thumbnailHtml = ` + + `; + } else { + const effective = thumbnailUrl || placeholder; + thumbnailHtml = ` + ${imageAlt} + `; + } + + const overlayHtml = shouldBlur ? `
R
` : ''; + const containerClasses = `huggingface-thumbnail-container${shouldBlur ? ' blurred' : ''}`; + + listItem.innerHTML = ` +
+ ${thumbnailHtml} + ${overlayHtml} +
${modelTypeApi}
+
+
+

${modelName}

+
+ ${creator} + ${baseModelsDisplay} + ${lastUpdatedFormatted} +
+
+ ${stats.downloadCount?.toLocaleString() || 0} + ${stats.thumbsUpCount?.toLocaleString() || 0} + ${stats.collectedCount?.toLocaleString() || 0} + ${stats.tippedAmountCount?.toLocaleString() || 0} +
+ ${tags.length > 0 ? ` +
+ ${tags.slice(0, 5).map(tag => `${tag}`).join('')} + ${tags.length > 5 ? `...` : ''} +
+ ` : ''} +
+
+ + View + +
+ ${versionButtonsHtml} + ${moreButtonHtml} +
+ ${allVersionsHtml} +
+ `; + + fragment.appendChild(listItem); + }); + + ui.searchResultsContainer.innerHTML = ''; + ui.searchResultsContainer.appendChild(fragment); +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/statusRenderer.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/statusRenderer.js new file mode 100644 index 0000000000000000000000000000000000000000..424138405a1889554b77bffa973b9e607252ddf8 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/statusRenderer.js @@ -0,0 +1,124 @@ +// Renders active/queued/history download lists + +const PLACEHOLDER_IMAGE_URL = `/extensions/ComfyUI-HuggingFace/images/placeholder.jpg`; + +export function renderDownloadList(ui, items, container, emptyMessage) { + if (!items || items.length === 0) { + container.innerHTML = `

${emptyMessage}

`; + return; + } + + const fragment = document.createDocumentFragment(); + items.forEach(item => { + const id = item.id || 'unknown-id'; + const progress = item.progress !== undefined ? Math.max(0, Math.min(100, item.progress)) : 0; + const speed = item.speed !== undefined ? Math.max(0, item.speed) : 0; + const status = item.status || 'unknown'; + const size = item.known_size !== undefined && item.known_size !== null ? item.known_size : (item.file_size || 0); + const downloadedBytes = size > 0 ? size * (progress / 100) : 0; + const errorMsg = item.error || null; + + // Debug: log what we receive from backend + console.log(`[DEBUG Frontend] Item data:`, { + huggingface_model_name: item.huggingface_model_name, + model_name: item.model_name, + model: item.model, + version_name: item.version_name, + full_item: item // Log the full item to see all available fields + }); + + const modelName = item.huggingface_model_name || item.model_name || item.model?.name || 'Unknown Model'; + console.log(`[DEBUG Frontend] Final modelName: ${modelName}`); + const versionName = item.version_name || 'Unknown Version'; + const filename = item.filename || 'N/A'; + const addedTime = item.added_time || null; + const startTime = item.start_time || null; + const endTime = item.end_time || null; + const thumbnail = item.thumbnail || PLACEHOLDER_IMAGE_URL; + const nsfwLevel = Number(item.thumbnail_nsfw_level ?? 0); + const blurMinLevel = Number(ui.settings?.nsfwBlurMinLevel ?? 4); + const shouldBlur = ui.settings?.hideMatureInSearch === true && nsfwLevel >= blurMinLevel; + const connectionType = item.connection_type || "N/A"; + + let progressBarClass = ''; + let statusText = status.charAt(0).toUpperCase() + status.slice(1); + switch (status) { + case 'completed': progressBarClass = 'completed'; break; + case 'failed': progressBarClass = 'failed'; statusText = 'Failed'; break; + case 'cancelled': progressBarClass = 'cancelled'; statusText = 'Cancelled'; break; + case 'downloading': case 'queued': case 'starting': default: break; + } + + const listItem = document.createElement('div'); + listItem.className = 'huggingface-download-item'; + listItem.dataset.id = id; + + const onErrorScript = `this.onerror=null; this.src='${PLACEHOLDER_IMAGE_URL}'; this.style.backgroundColor='#444';`; + const addedTooltip = addedTime ? `data-tooltip="Added: ${new Date(addedTime).toLocaleString()}"` : ''; + const startedTooltip = startTime ? `data-tooltip="Started: ${new Date(startTime).toLocaleString()}"` : ''; + const endedTooltip = endTime ? `data-tooltip="Ended: ${new Date(endTime).toLocaleString()}"` : ''; + const durationTooltip = startTime && endTime ? `data-tooltip="Duration: ${ui.formatDuration(startTime, endTime)}"` : ''; + const filenameTooltip = filename !== 'N/A' ? `title="Filename: ${filename}"` : ''; + const errorTooltip = errorMsg ? `title="Error Details: ${String(errorMsg).substring(0, 200)}${String(errorMsg).length > 200 ? '...' : ''}"` : ''; + const connectionInfoHtml = connectionType !== "N/A" ? `(Conn: ${connectionType})` : ''; + + const overlayHtml = shouldBlur ? `
R
` : ''; + const containerClasses = `huggingface-thumbnail-container${shouldBlur ? ' blurred' : ''}`; + + let innerHTML = ` +
+ thumbnail + ${overlayHtml} +
+
+ ${modelName} +

Ver: ${versionName}

+

${filename}

+ ${size > 0 ? `

Size: ${ui.formatBytes(size)}

` : ''} + ${item.file_format ? `

Format: ${item.file_format}

` : ''} + ${item.file_precision || item.file_model_size ? `

${item.file_precision ? 'Precision: ' + String(item.file_precision).toUpperCase() : ''}${item.file_precision && item.file_model_size ? ' • ' : ''}${item.file_model_size ? 'Model Size: ' + item.file_model_size : ''}

` : ''} + ${errorMsg ? `

${String(errorMsg).substring(0, 100)}${String(errorMsg).length > 100 ? '...' : ''}

` : ''} + `; + + if (status === 'downloading' || status === 'starting' || status === 'completed') { + const statusLine = `
Status: ${statusText} ${connectionInfoHtml}
`; + innerHTML += ` +
+
+ ${progress > 15 ? progress.toFixed(0)+'%' : ''} +
+
+ `; + const speedText = (status === 'downloading' && speed > 0) ? ui.formatSpeed(speed) : ''; + const progressText = (status === 'downloading' && size > 0) ? `(${ui.formatBytes(downloadedBytes)} / ${ui.formatBytes(size)})` : ''; + const completedText = status === 'completed' ? '' : ''; + const speedProgLine = `
${speedText} ${progressText} ${completedText}
`; + if (status === 'downloading') { innerHTML += speedProgLine; } + innerHTML += statusLine; + } else if (status === 'failed' || status === 'cancelled' || status === 'queued') { + innerHTML += `
Status: ${statusText} ${connectionInfoHtml}
`; + } else { + innerHTML += `
Status: ${statusText} ${connectionInfoHtml}
`; + } + + innerHTML += `
`; + innerHTML += `
`; + if (status === 'queued' || status === 'downloading' || status === 'starting') { + innerHTML += ``; + } + if (status === 'failed' || status === 'cancelled') { + innerHTML += ``; + } + if (status === 'completed') { + innerHTML += ``; + } + innerHTML += `
`; + + listItem.innerHTML = innerHTML; + fragment.appendChild(listItem); + }); + + container.innerHTML = ''; + container.appendChild(fragment); + ui.ensureFontAwesome(); +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/ui/templates.js b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/templates.js new file mode 100644 index 0000000000000000000000000000000000000000..e74af211e567b0ac98c51270c4d3d7b73335f367 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/ui/templates.js @@ -0,0 +1,193 @@ +// Modal template for HuggingFace UI +// Keep structure identical to the original inline HTML to minimize risk + +export function modalTemplate(settings = {}) { + const numConnections = Number.isFinite(settings.numConnections) ? settings.numConnections : 1; + return ` +
+
+

HuggingFace

+ +
+
+
+ + + + +
+
+
+
+ + +
+

You can optionally specify a version ID using "?modelVersionId=xxxxx" in the URL or in the field below.

+
+
+ +
+ + +
+
+
+ +
+ + +
+

+
+
+ + +
+
+
+
+ + +
+
+ + +

Disabled: Only single connection possible for now

+
+
+
+ + +
+
+ +
+ +
+
+ +
+
+
+

Active Downloads

+
+

No active downloads.

+
+
+
+

Queued Downloads

+
+

Download queue is empty.

+
+
+
+
+

Download History (Recent)

+ +
+
+

No download history yet.

+
+
+
+
+
+
+
+
+

API & Defaults

+
+ + +

Needed for some downloads/features. Leave blank to use server env HUGGINGFACE_TOKEN. Find keys at huggingface.com/user/account

+
+
+ + +

+ When set, downloads use <global_root>/<model_type> (for example /runpod-volume/ComfyUI/checkpoints). +

+
+ + +
+
+
+ + +

Number of parallel connections for downloads (1-16)

+
+
+ + +
+
+
+

Interface & Search

+
+ + +
+
+ + +
+
+ + +

+ Blur thumbnails when an image's nsfwLevel is greater than or equal to this value. + Higher numbers indicate more explicit content. None (Safe/PG): 1, Mild (PG-13): 2, Mature (R): 4, Adult (X): 5, Extra Explicit (R): 8, Explicit (XXX): 16/32+ +

+
+
+
+ +
+
+
+ +
+ +
+
+

Confirm Clear History

+

Are you sure you want to clear the download history? This action cannot be undone.

+
+ + +
+
+
+
+ `; +} diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/utils/cookies.js b/custom_nodes/ComfyUI-HuggingFace/web/js/utils/cookies.js new file mode 100644 index 0000000000000000000000000000000000000000..f8d7c1a2896186707d7b76905ad16e0c969b7357 --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/utils/cookies.js @@ -0,0 +1,24 @@ +// Lightweight cookie helpers for HuggingFace UI +// Exports: setCookie, getCookie + +export function setCookie(name, value, days) { + let expires = ""; + if (days) { + const date = new Date(); + date.setTime(date.getTime() + days * 24 * 60 * 60 * 1000); + expires = "; expires=" + date.toUTCString(); + } + document.cookie = `${name}=${value || ""}${expires}; path=/; SameSite=Lax`; +} + +export function getCookie(name) { + const nameEQ = name + "="; + const parts = document.cookie.split(";"); + for (let i = 0; i < parts.length; i++) { + let c = parts[i]; + while (c.charAt(0) === " ") c = c.substring(1); + if (c.indexOf(nameEQ) === 0) return c.substring(nameEQ.length); + } + return null; +} + diff --git a/custom_nodes/ComfyUI-HuggingFace/web/js/utils/dom.js b/custom_nodes/ComfyUI-HuggingFace/web/js/utils/dom.js new file mode 100644 index 0000000000000000000000000000000000000000..e06f81af9585b6dab855458d5608508f4e1bde3d --- /dev/null +++ b/custom_nodes/ComfyUI-HuggingFace/web/js/utils/dom.js @@ -0,0 +1,32 @@ +// File: web/js/utils/dom.js + +/** + * Dynamically adds a CSS link to the document's head. + * It resolves the path relative to this script's location using import.meta.url, + * making it robust against case-sensitivity issues and different install paths. + * @param {string} relativeHref - Relative path to the CSS file (e.g., '../huggingfaceDownloader.css'). + * @param {string} [id="huggingface-downloader-styles"] - The ID for the link element. + */ +export function addCssLink(relativeHref, id = "huggingface-downloader-styles") { + if (document.getElementById(id)) return; // Prevent duplicates + + try { + const absoluteUrl = new URL(relativeHref, import.meta.url); + + const link = document.createElement("link"); + link.id = id; + link.rel = "stylesheet"; + link.href = absoluteUrl.href; + + link.onload = () => { + console.log("[HuggingFace] CSS loaded successfully:", link.href); + }; + link.onerror = () => { + console.error("[HuggingFace] Critical error: Failed to load CSS from:", link.href); + }; + + document.head.appendChild(link); + } catch (e) { + console.error("[HuggingFace] Error creating CSS link. import.meta.url may be unsupported in this context.", e); + } +} diff --git a/custom_nodes/ComfyUI-KJNodes/.github/workflows/publish.yml b/custom_nodes/ComfyUI-KJNodes/.github/workflows/publish.yml new file mode 100644 index 0000000000000000000000000000000000000000..e155f5f40e46fa83942dee5a9460e8093f3b4208 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/.github/workflows/publish.yml @@ -0,0 +1,25 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'kijai' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/custom_nodes/ComfyUI-KJNodes/docs/images/319121566-05f66385-7568-4b1f-8bbc-11053660b02f.png b/custom_nodes/ComfyUI-KJNodes/docs/images/319121566-05f66385-7568-4b1f-8bbc-11053660b02f.png new file mode 100644 index 0000000000000000000000000000000000000000..e749239c1c4ffd5ab29b51695dd8d8b51ed3597f Binary files /dev/null and b/custom_nodes/ComfyUI-KJNodes/docs/images/319121566-05f66385-7568-4b1f-8bbc-11053660b02f.png differ diff --git a/custom_nodes/ComfyUI-KJNodes/docs/images/319121636-706b5081-9120-4a29-bd76-901691ada688.png b/custom_nodes/ComfyUI-KJNodes/docs/images/319121636-706b5081-9120-4a29-bd76-901691ada688.png new file mode 100644 index 0000000000000000000000000000000000000000..b53ad666ff060d87971f3962e74101f0cb2a5c3f Binary files /dev/null and b/custom_nodes/ComfyUI-KJNodes/docs/images/319121636-706b5081-9120-4a29-bd76-901691ada688.png differ diff --git a/custom_nodes/ComfyUI-KJNodes/example_workflows/leapfusion_hunyuuanvideo_i2v_native_testing.json b/custom_nodes/ComfyUI-KJNodes/example_workflows/leapfusion_hunyuuanvideo_i2v_native_testing.json new file mode 100644 index 0000000000000000000000000000000000000000..134a83788815d04b5574a807db67eb6e45bf9263 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/example_workflows/leapfusion_hunyuuanvideo_i2v_native_testing.json @@ -0,0 +1,1188 @@ +{ + "last_node_id": 86, + "last_link_id": 144, + "nodes": [ + { + "id": 62, + "type": "FluxGuidance", + "pos": [ + -630, + -170 + ], + "size": [ + 317.4000244140625, + 58 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 82 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 83 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 6 + ] + }, + { + "id": 51, + "type": "KSamplerSelect", + "pos": [ + -610, + -480 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 61 + ] + } + ], + "properties": { + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "euler" + ] + }, + { + "id": 57, + "type": "VAEDecodeTiled", + "pos": [ + -200, + 90 + ], + "size": [ + 315, + 150 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 142 + }, + { + "name": "vae", + "type": "VAE", + "link": 74 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 105 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecodeTiled" + }, + "widgets_values": [ + 128, + 64, + 64, + 8 + ] + }, + { + "id": 65, + "type": "LoadImage", + "pos": [ + -2212.498779296875, + -632.4085083007812 + ], + "size": [ + 315, + 314 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 86 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Mona-Lisa-oil-wood-panel-Leonardo-da.webp", + "image" + ] + }, + { + "id": 64, + "type": "VAEEncode", + "pos": [ + -1336.7884521484375, + -492.5806884765625 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 144 + }, + { + "name": "vae", + "type": "VAE", + "link": 88 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 137 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + }, + "widgets_values": [] + }, + { + "id": 44, + "type": "UNETLoader", + "pos": [ + -2373.55029296875, + -193.91510009765625 + ], + "size": [ + 459.56060791015625, + 82 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 135 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "UNETLoader" + }, + "widgets_values": [ + "hyvideo\\hunyuan_video_720_fp8_e4m3fn.safetensors", + "fp8_e4m3fn_fast" + ] + }, + { + "id": 49, + "type": "VAELoader", + "pos": [ + -1876.39306640625, + -35.19633865356445 + ], + "size": [ + 433.7603454589844, + 58.71116256713867 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 74, + 88 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "hyvid\\hunyuan_video_vae_bf16.safetensors" + ] + }, + { + "id": 47, + "type": "DualCLIPLoader", + "pos": [ + -2284.893798828125, + 150.4042205810547 + ], + "size": [ + 343.3958435058594, + 106.86042785644531 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 56 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader" + }, + "widgets_values": [ + "clip_l.safetensors", + "llava_llama3_fp16.safetensors", + "hunyuan_video", + "default" + ] + }, + { + "id": 45, + "type": "CLIPTextEncode", + "pos": [ + -1839.1649169921875, + 143.5203094482422 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 56 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 69, + 82 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "woman puts on sunglasses" + ] + }, + { + "id": 53, + "type": "EmptyHunyuanLatentVideo", + "pos": [ + -1120, + 90 + ], + "size": [ + 315, + 130 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "link": 89, + "widget": { + "name": "width" + } + }, + { + "name": "height", + "type": "INT", + "link": 90, + "widget": { + "name": "height" + } + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 119 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyHunyuanLatentVideo" + }, + "widgets_values": [ + 960, + 544, + 65, + 1 + ] + }, + { + "id": 55, + "type": "ConditioningZeroOut", + "pos": [ + -910, + 300 + ], + "size": [ + 251.14309692382812, + 26 + ], + "flags": { + "collapsed": true + }, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 69 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 70 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ConditioningZeroOut" + }, + "widgets_values": [] + }, + { + "id": 52, + "type": "BasicScheduler", + "pos": [ + -600, + -350 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 78 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 62 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 20, + 1 + ] + }, + { + "id": 42, + "type": "SamplerCustom", + "pos": [ + -640, + 10 + ], + "size": [ + 355.20001220703125, + 467.4666748046875 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 77 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 83 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 70 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 61 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 62 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 119 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": null + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": [ + 141 + ], + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "SamplerCustom" + }, + "widgets_values": [ + true, + 6, + "fixed", + 1, + null + ] + }, + { + "id": 84, + "type": "GetLatentRangeFromBatch", + "pos": [ + -240, + -100 + ], + "size": [ + 340.20001220703125, + 82 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "latents", + "type": "LATENT", + "link": 141 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 142 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetLatentRangeFromBatch" + }, + "widgets_values": [ + 1, + -1 + ] + }, + { + "id": 50, + "type": "VHS_VideoCombine", + "pos": [ + 165.77645874023438, + -619.0606079101562 + ], + "size": [ + 1112.6898193359375, + 1076.4598388671875 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 105 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 24, + "loop_count": 0, + "filename_prefix": "hyvidcomfy", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "hyvidcomfy_00001.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 24, + "workflow": "hyvidcomfy_00001.png", + "fullpath": "N:\\AI\\ComfyUI\\temp\\hyvidcomfy_00001.mp4" + }, + "muted": false + } + } + }, + { + "id": 54, + "type": "ModelSamplingSD3", + "pos": [ + -1079.9112548828125, + -146.69448852539062 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 117 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 77, + 78 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ModelSamplingSD3" + }, + "widgets_values": [ + 9 + ] + }, + { + "id": 80, + "type": "PathchSageAttentionKJ", + "pos": [ + -2273.926513671875, + -36.720542907714844 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 7, + "mode": 4, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 135 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 136 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PathchSageAttentionKJ" + }, + "widgets_values": [ + "auto" + ] + }, + { + "id": 85, + "type": "Note", + "pos": [ + -1838.572265625, + -302.1575927734375 + ], + "size": [ + 408.4594421386719, + 58 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "https://huggingface.co/Kijai/Leapfusion-image2vid-comfy/blob/main/leapfusion_img2vid544p_comfy.safetensors" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 74, + "type": "LeapfusionHunyuanI2VPatcher", + "pos": [ + -1059.552978515625, + -459.34674072265625 + ], + "size": [ + 277.3238525390625, + 150 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 123 + }, + { + "name": "latent", + "type": "LATENT", + "link": 137 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 117 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LeapfusionHunyuanI2VPatcher" + }, + "widgets_values": [ + 0, + 0, + 1, + 0.8 + ] + }, + { + "id": 59, + "type": "LoraLoaderModelOnly", + "pos": [ + -1870.3748779296875, + -194.6091766357422 + ], + "size": [ + 442.8438720703125, + 82 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 136 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 123 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LoraLoaderModelOnly" + }, + "widgets_values": [ + "hyvid\\musubi-tuner\\img2vid544p.safetensors", + 1 + ] + }, + { + "id": 66, + "type": "ImageResizeKJ", + "pos": [ + -1821.1531982421875, + -632.925048828125 + ], + "size": [ + 315, + 266 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 86 + }, + { + "name": "get_image_size", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "width_input", + "type": "INT", + "link": null, + "widget": { + "name": "width_input" + }, + "shape": 7 + }, + { + "name": "height_input", + "type": "INT", + "link": null, + "widget": { + "name": "height_input" + }, + "shape": 7 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 143 + ], + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "links": [ + 89 + ], + "slot_index": 1 + }, + { + "name": "height", + "type": "INT", + "links": [ + 90 + ], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "ImageResizeKJ" + }, + "widgets_values": [ + 960, + 640, + "lanczos", + false, + 2, + 0, + 0, + "center" + ] + }, + { + "id": 86, + "type": "ImageNoiseAugmentation", + "pos": [ + -1361.111572265625, + -667.0104370117188 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 143 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 144 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageNoiseAugmentation" + }, + "widgets_values": [ + 0.05, + 123, + "fixed" + ] + } + ], + "links": [ + [ + 56, + 47, + 0, + 45, + 0, + "CLIP" + ], + [ + 61, + 51, + 0, + 42, + 3, + "SAMPLER" + ], + [ + 62, + 52, + 0, + 42, + 4, + "SIGMAS" + ], + [ + 69, + 45, + 0, + 55, + 0, + "CONDITIONING" + ], + [ + 70, + 55, + 0, + 42, + 2, + "CONDITIONING" + ], + [ + 74, + 49, + 0, + 57, + 1, + "VAE" + ], + [ + 77, + 54, + 0, + 42, + 0, + "MODEL" + ], + [ + 78, + 54, + 0, + 52, + 0, + "MODEL" + ], + [ + 82, + 45, + 0, + 62, + 0, + "CONDITIONING" + ], + [ + 83, + 62, + 0, + 42, + 1, + "CONDITIONING" + ], + [ + 86, + 65, + 0, + 66, + 0, + "IMAGE" + ], + [ + 88, + 49, + 0, + 64, + 1, + "VAE" + ], + [ + 89, + 66, + 1, + 53, + 0, + "INT" + ], + [ + 90, + 66, + 2, + 53, + 1, + "INT" + ], + [ + 105, + 57, + 0, + 50, + 0, + "IMAGE" + ], + [ + 117, + 74, + 0, + 54, + 0, + "MODEL" + ], + [ + 119, + 53, + 0, + 42, + 5, + "LATENT" + ], + [ + 123, + 59, + 0, + 74, + 0, + "MODEL" + ], + [ + 135, + 44, + 0, + 80, + 0, + "MODEL" + ], + [ + 136, + 80, + 0, + 59, + 0, + "MODEL" + ], + [ + 137, + 64, + 0, + 74, + 1, + "LATENT" + ], + [ + 141, + 42, + 1, + 84, + 0, + "LATENT" + ], + [ + 142, + 84, + 0, + 57, + 0, + "LATENT" + ], + [ + 143, + 66, + 0, + 86, + 0, + "IMAGE" + ], + [ + 144, + 86, + 0, + 64, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.740024994425854, + "offset": [ + 2525.036093151529, + 802.59123935694 + ] + }, + "node_versions": { + "comfy-core": "0.3.13", + "ComfyUI-KJNodes": "a8aeef670b3f288303f956bf94385cb87978ea93", + "ComfyUI-VideoHelperSuite": "c47b10ca1798b4925ff5a5f07d80c51ca80a837d" + }, + "VHS_latentpreview": true, + "VHS_latentpreviewrate": 0 + }, + "version": 0.4 +} \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/kjweb_async/marked.min.js b/custom_nodes/ComfyUI-KJNodes/kjweb_async/marked.min.js new file mode 100644 index 0000000000000000000000000000000000000000..2e66c369c388c135cc68d399861a737f4c5e68cd --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/kjweb_async/marked.min.js @@ -0,0 +1,6 @@ +/** + * marked v12.0.1 - a markdown parser + * Copyright (c) 2011-2024, Christopher Jeffrey. (MIT Licensed) + * https://github.com/markedjs/marked + */ +!function(e,t){"object"==typeof exports&&"undefined"!=typeof module?t(exports):"function"==typeof define&&define.amd?define(["exports"],t):t((e="undefined"!=typeof globalThis?globalThis:e||self).marked={})}(this,(function(e){"use strict";function t(){return{async:!1,breaks:!1,extensions:null,gfm:!0,hooks:null,pedantic:!1,renderer:null,silent:!1,tokenizer:null,walkTokens:null}}function n(t){e.defaults=t}e.defaults={async:!1,breaks:!1,extensions:null,gfm:!0,hooks:null,pedantic:!1,renderer:null,silent:!1,tokenizer:null,walkTokens:null};const s=/[&<>"']/,r=new RegExp(s.source,"g"),i=/[<>"']|&(?!(#\d{1,7}|#[Xx][a-fA-F0-9]{1,6}|\w+);)/,l=new RegExp(i.source,"g"),o={"&":"&","<":"<",">":">",'"':""","'":"'"},a=e=>o[e];function c(e,t){if(t){if(s.test(e))return e.replace(r,a)}else if(i.test(e))return e.replace(l,a);return e}const h=/&(#(?:\d+)|(?:#x[0-9A-Fa-f]+)|(?:\w+));?/gi;function p(e){return e.replace(h,((e,t)=>"colon"===(t=t.toLowerCase())?":":"#"===t.charAt(0)?"x"===t.charAt(1)?String.fromCharCode(parseInt(t.substring(2),16)):String.fromCharCode(+t.substring(1)):""))}const u=/(^|[^\[])\^/g;function k(e,t){let n="string"==typeof e?e:e.source;t=t||"";const s={replace:(e,t)=>{let r="string"==typeof t?t:t.source;return r=r.replace(u,"$1"),n=n.replace(e,r),s},getRegex:()=>new RegExp(n,t)};return s}function g(e){try{e=encodeURI(e).replace(/%25/g,"%")}catch(e){return null}return e}const f={exec:()=>null};function d(e,t){const n=e.replace(/\|/g,((e,t,n)=>{let s=!1,r=t;for(;--r>=0&&"\\"===n[r];)s=!s;return s?"|":" |"})).split(/ \|/);let s=0;if(n[0].trim()||n.shift(),n.length>0&&!n[n.length-1].trim()&&n.pop(),t)if(n.length>t)n.splice(t);else for(;n.length0)return{type:"space",raw:t[0]}}code(e){const t=this.rules.block.code.exec(e);if(t){const e=t[0].replace(/^ {1,4}/gm,"");return{type:"code",raw:t[0],codeBlockStyle:"indented",text:this.options.pedantic?e:x(e,"\n")}}}fences(e){const t=this.rules.block.fences.exec(e);if(t){const e=t[0],n=function(e,t){const n=e.match(/^(\s+)(?:```)/);if(null===n)return t;const s=n[1];return t.split("\n").map((e=>{const t=e.match(/^\s+/);if(null===t)return e;const[n]=t;return n.length>=s.length?e.slice(s.length):e})).join("\n")}(e,t[3]||"");return{type:"code",raw:e,lang:t[2]?t[2].trim().replace(this.rules.inline.anyPunctuation,"$1"):t[2],text:n}}}heading(e){const t=this.rules.block.heading.exec(e);if(t){let e=t[2].trim();if(/#$/.test(e)){const t=x(e,"#");this.options.pedantic?e=t.trim():t&&!/ $/.test(t)||(e=t.trim())}return{type:"heading",raw:t[0],depth:t[1].length,text:e,tokens:this.lexer.inline(e)}}}hr(e){const t=this.rules.block.hr.exec(e);if(t)return{type:"hr",raw:t[0]}}blockquote(e){const t=this.rules.block.blockquote.exec(e);if(t){const e=x(t[0].replace(/^ *>[ \t]?/gm,""),"\n"),n=this.lexer.state.top;this.lexer.state.top=!0;const s=this.lexer.blockTokens(e);return this.lexer.state.top=n,{type:"blockquote",raw:t[0],tokens:s,text:e}}}list(e){let t=this.rules.block.list.exec(e);if(t){let n=t[1].trim();const s=n.length>1,r={type:"list",raw:"",ordered:s,start:s?+n.slice(0,-1):"",loose:!1,items:[]};n=s?`\\d{1,9}\\${n.slice(-1)}`:`\\${n}`,this.options.pedantic&&(n=s?n:"[*+-]");const i=new RegExp(`^( {0,3}${n})((?:[\t ][^\\n]*)?(?:\\n|$))`);let l="",o="",a=!1;for(;e;){let n=!1;if(!(t=i.exec(e)))break;if(this.rules.block.hr.test(e))break;l=t[0],e=e.substring(l.length);let s=t[2].split("\n",1)[0].replace(/^\t+/,(e=>" ".repeat(3*e.length))),c=e.split("\n",1)[0],h=0;this.options.pedantic?(h=2,o=s.trimStart()):(h=t[2].search(/[^ ]/),h=h>4?1:h,o=s.slice(h),h+=t[1].length);let p=!1;if(!s&&/^ *$/.test(c)&&(l+=c+"\n",e=e.substring(c.length+1),n=!0),!n){const t=new RegExp(`^ {0,${Math.min(3,h-1)}}(?:[*+-]|\\d{1,9}[.)])((?:[ \t][^\\n]*)?(?:\\n|$))`),n=new RegExp(`^ {0,${Math.min(3,h-1)}}((?:- *){3,}|(?:_ *){3,}|(?:\\* *){3,})(?:\\n+|$)`),r=new RegExp(`^ {0,${Math.min(3,h-1)}}(?:\`\`\`|~~~)`),i=new RegExp(`^ {0,${Math.min(3,h-1)}}#`);for(;e;){const a=e.split("\n",1)[0];if(c=a,this.options.pedantic&&(c=c.replace(/^ {1,4}(?=( {4})*[^ ])/g," ")),r.test(c))break;if(i.test(c))break;if(t.test(c))break;if(n.test(e))break;if(c.search(/[^ ]/)>=h||!c.trim())o+="\n"+c.slice(h);else{if(p)break;if(s.search(/[^ ]/)>=4)break;if(r.test(s))break;if(i.test(s))break;if(n.test(s))break;o+="\n"+c}p||c.trim()||(p=!0),l+=a+"\n",e=e.substring(a.length+1),s=c.slice(h)}}r.loose||(a?r.loose=!0:/\n *\n *$/.test(l)&&(a=!0));let u,k=null;this.options.gfm&&(k=/^\[[ xX]\] /.exec(o),k&&(u="[ ] "!==k[0],o=o.replace(/^\[[ xX]\] +/,""))),r.items.push({type:"list_item",raw:l,task:!!k,checked:u,loose:!1,text:o,tokens:[]}),r.raw+=l}r.items[r.items.length-1].raw=l.trimEnd(),r.items[r.items.length-1].text=o.trimEnd(),r.raw=r.raw.trimEnd();for(let e=0;e"space"===e.type)),n=t.length>0&&t.some((e=>/\n.*\n/.test(e.raw)));r.loose=n}if(r.loose)for(let e=0;e$/,"$1").replace(this.rules.inline.anyPunctuation,"$1"):"",s=t[3]?t[3].substring(1,t[3].length-1).replace(this.rules.inline.anyPunctuation,"$1"):t[3];return{type:"def",tag:e,raw:t[0],href:n,title:s}}}table(e){const t=this.rules.block.table.exec(e);if(!t)return;if(!/[:|]/.test(t[2]))return;const n=d(t[1]),s=t[2].replace(/^\||\| *$/g,"").split("|"),r=t[3]&&t[3].trim()?t[3].replace(/\n[ \t]*$/,"").split("\n"):[],i={type:"table",raw:t[0],header:[],align:[],rows:[]};if(n.length===s.length){for(const e of s)/^ *-+: *$/.test(e)?i.align.push("right"):/^ *:-+: *$/.test(e)?i.align.push("center"):/^ *:-+ *$/.test(e)?i.align.push("left"):i.align.push(null);for(const e of n)i.header.push({text:e,tokens:this.lexer.inline(e)});for(const e of r)i.rows.push(d(e,i.header.length).map((e=>({text:e,tokens:this.lexer.inline(e)}))));return i}}lheading(e){const t=this.rules.block.lheading.exec(e);if(t)return{type:"heading",raw:t[0],depth:"="===t[2].charAt(0)?1:2,text:t[1],tokens:this.lexer.inline(t[1])}}paragraph(e){const t=this.rules.block.paragraph.exec(e);if(t){const e="\n"===t[1].charAt(t[1].length-1)?t[1].slice(0,-1):t[1];return{type:"paragraph",raw:t[0],text:e,tokens:this.lexer.inline(e)}}}text(e){const t=this.rules.block.text.exec(e);if(t)return{type:"text",raw:t[0],text:t[0],tokens:this.lexer.inline(t[0])}}escape(e){const t=this.rules.inline.escape.exec(e);if(t)return{type:"escape",raw:t[0],text:c(t[1])}}tag(e){const t=this.rules.inline.tag.exec(e);if(t)return!this.lexer.state.inLink&&/^/i.test(t[0])&&(this.lexer.state.inLink=!1),!this.lexer.state.inRawBlock&&/^<(pre|code|kbd|script)(\s|>)/i.test(t[0])?this.lexer.state.inRawBlock=!0:this.lexer.state.inRawBlock&&/^<\/(pre|code|kbd|script)(\s|>)/i.test(t[0])&&(this.lexer.state.inRawBlock=!1),{type:"html",raw:t[0],inLink:this.lexer.state.inLink,inRawBlock:this.lexer.state.inRawBlock,block:!1,text:t[0]}}link(e){const t=this.rules.inline.link.exec(e);if(t){const e=t[2].trim();if(!this.options.pedantic&&/^$/.test(e))return;const t=x(e.slice(0,-1),"\\");if((e.length-t.length)%2==0)return}else{const e=function(e,t){if(-1===e.indexOf(t[1]))return-1;let n=0;for(let s=0;s-1){const n=(0===t[0].indexOf("!")?5:4)+t[1].length+e;t[2]=t[2].substring(0,e),t[0]=t[0].substring(0,n).trim(),t[3]=""}}let n=t[2],s="";if(this.options.pedantic){const e=/^([^'"]*[^\s])\s+(['"])(.*)\2/.exec(n);e&&(n=e[1],s=e[3])}else s=t[3]?t[3].slice(1,-1):"";return n=n.trim(),/^$/.test(e)?n.slice(1):n.slice(1,-1)),b(t,{href:n?n.replace(this.rules.inline.anyPunctuation,"$1"):n,title:s?s.replace(this.rules.inline.anyPunctuation,"$1"):s},t[0],this.lexer)}}reflink(e,t){let n;if((n=this.rules.inline.reflink.exec(e))||(n=this.rules.inline.nolink.exec(e))){const e=t[(n[2]||n[1]).replace(/\s+/g," ").toLowerCase()];if(!e){const e=n[0].charAt(0);return{type:"text",raw:e,text:e}}return b(n,e,n[0],this.lexer)}}emStrong(e,t,n=""){let s=this.rules.inline.emStrongLDelim.exec(e);if(!s)return;if(s[3]&&n.match(/[\p{L}\p{N}]/u))return;if(!(s[1]||s[2]||"")||!n||this.rules.inline.punctuation.exec(n)){const n=[...s[0]].length-1;let r,i,l=n,o=0;const a="*"===s[0][0]?this.rules.inline.emStrongRDelimAst:this.rules.inline.emStrongRDelimUnd;for(a.lastIndex=0,t=t.slice(-1*e.length+n);null!=(s=a.exec(t));){if(r=s[1]||s[2]||s[3]||s[4]||s[5]||s[6],!r)continue;if(i=[...r].length,s[3]||s[4]){l+=i;continue}if((s[5]||s[6])&&n%3&&!((n+i)%3)){o+=i;continue}if(l-=i,l>0)continue;i=Math.min(i,i+l+o);const t=[...s[0]][0].length,a=e.slice(0,n+s.index+t+i);if(Math.min(n,i)%2){const e=a.slice(1,-1);return{type:"em",raw:a,text:e,tokens:this.lexer.inlineTokens(e)}}const c=a.slice(2,-2);return{type:"strong",raw:a,text:c,tokens:this.lexer.inlineTokens(c)}}}}codespan(e){const t=this.rules.inline.code.exec(e);if(t){let e=t[2].replace(/\n/g," ");const n=/[^ ]/.test(e),s=/^ /.test(e)&&/ $/.test(e);return n&&s&&(e=e.substring(1,e.length-1)),e=c(e,!0),{type:"codespan",raw:t[0],text:e}}}br(e){const t=this.rules.inline.br.exec(e);if(t)return{type:"br",raw:t[0]}}del(e){const t=this.rules.inline.del.exec(e);if(t)return{type:"del",raw:t[0],text:t[2],tokens:this.lexer.inlineTokens(t[2])}}autolink(e){const t=this.rules.inline.autolink.exec(e);if(t){let e,n;return"@"===t[2]?(e=c(t[1]),n="mailto:"+e):(e=c(t[1]),n=e),{type:"link",raw:t[0],text:e,href:n,tokens:[{type:"text",raw:e,text:e}]}}}url(e){let t;if(t=this.rules.inline.url.exec(e)){let e,n;if("@"===t[2])e=c(t[0]),n="mailto:"+e;else{let s;do{s=t[0],t[0]=this.rules.inline._backpedal.exec(t[0])?.[0]??""}while(s!==t[0]);e=c(t[0]),n="www."===t[1]?"http://"+t[0]:t[0]}return{type:"link",raw:t[0],text:e,href:n,tokens:[{type:"text",raw:e,text:e}]}}}inlineText(e){const t=this.rules.inline.text.exec(e);if(t){let e;return e=this.lexer.state.inRawBlock?t[0]:c(t[0]),{type:"text",raw:t[0],text:e}}}}const m=/^ {0,3}((?:-[\t ]*){3,}|(?:_[ \t]*){3,}|(?:\*[ \t]*){3,})(?:\n+|$)/,y=/(?:[*+-]|\d{1,9}[.)])/,$=k(/^(?!bull |blockCode|fences|blockquote|heading|html)((?:.|\n(?!\s*?\n|bull |blockCode|fences|blockquote|heading|html))+?)\n {0,3}(=+|-+) *(?:\n+|$)/).replace(/bull/g,y).replace(/blockCode/g,/ {4}/).replace(/fences/g,/ {0,3}(?:`{3,}|~{3,})/).replace(/blockquote/g,/ {0,3}>/).replace(/heading/g,/ {0,3}#{1,6}/).replace(/html/g,/ {0,3}<[^\n>]+>\n/).getRegex(),z=/^([^\n]+(?:\n(?!hr|heading|lheading|blockquote|fences|list|html|table| +\n)[^\n]+)*)/,T=/(?!\s*\])(?:\\.|[^\[\]\\])+/,R=k(/^ {0,3}\[(label)\]: *(?:\n *)?([^<\s][^\s]*|<.*?>)(?:(?: +(?:\n *)?| *\n *)(title))? *(?:\n+|$)/).replace("label",T).replace("title",/(?:"(?:\\"?|[^"\\])*"|'[^'\n]*(?:\n[^'\n]+)*\n?'|\([^()]*\))/).getRegex(),_=k(/^( {0,3}bull)([ \t][^\n]+?)?(?:\n|$)/).replace(/bull/g,y).getRegex(),A="address|article|aside|base|basefont|blockquote|body|caption|center|col|colgroup|dd|details|dialog|dir|div|dl|dt|fieldset|figcaption|figure|footer|form|frame|frameset|h[1-6]|head|header|hr|html|iframe|legend|li|link|main|menu|menuitem|meta|nav|noframes|ol|optgroup|option|p|param|search|section|summary|table|tbody|td|tfoot|th|thead|title|tr|track|ul",S=/|$))/,I=k("^ {0,3}(?:<(script|pre|style|textarea)[\\s>][\\s\\S]*?(?:[^\\n]*\\n+|$)|comment[^\\n]*(\\n+|$)|<\\?[\\s\\S]*?(?:\\?>\\n*|$)|\\n*|$)|\\n*|$)|)[\\s\\S]*?(?:(?:\\n *)+\\n|$)|<(?!script|pre|style|textarea)([a-z][\\w-]*)(?:attribute)*? */?>(?=[ \\t]*(?:\\n|$))[\\s\\S]*?(?:(?:\\n *)+\\n|$)|(?=[ \\t]*(?:\\n|$))[\\s\\S]*?(?:(?:\\n *)+\\n|$))","i").replace("comment",S).replace("tag",A).replace("attribute",/ +[a-zA-Z:_][\w.:-]*(?: *= *"[^"\n]*"| *= *'[^'\n]*'| *= *[^\s"'=<>`]+)?/).getRegex(),E=k(z).replace("hr",m).replace("heading"," {0,3}#{1,6}(?:\\s|$)").replace("|lheading","").replace("|table","").replace("blockquote"," {0,3}>").replace("fences"," {0,3}(?:`{3,}(?=[^`\\n]*\\n)|~{3,})[^\\n]*\\n").replace("list"," {0,3}(?:[*+-]|1[.)]) ").replace("html",")|<(?:script|pre|style|textarea|!--)").replace("tag",A).getRegex(),q={blockquote:k(/^( {0,3}> ?(paragraph|[^\n]*)(?:\n|$))+/).replace("paragraph",E).getRegex(),code:/^( {4}[^\n]+(?:\n(?: *(?:\n|$))*)?)+/,def:R,fences:/^ {0,3}(`{3,}(?=[^`\n]*(?:\n|$))|~{3,})([^\n]*)(?:\n|$)(?:|([\s\S]*?)(?:\n|$))(?: {0,3}\1[~`]* *(?=\n|$)|$)/,heading:/^ {0,3}(#{1,6})(?=\s|$)(.*)(?:\n+|$)/,hr:m,html:I,lheading:$,list:_,newline:/^(?: *(?:\n|$))+/,paragraph:E,table:f,text:/^[^\n]+/},Z=k("^ *([^\\n ].*)\\n {0,3}((?:\\| *)?:?-+:? *(?:\\| *:?-+:? *)*(?:\\| *)?)(?:\\n((?:(?! *\\n|hr|heading|blockquote|code|fences|list|html).*(?:\\n|$))*)\\n*|$)").replace("hr",m).replace("heading"," {0,3}#{1,6}(?:\\s|$)").replace("blockquote"," {0,3}>").replace("code"," {4}[^\\n]").replace("fences"," {0,3}(?:`{3,}(?=[^`\\n]*\\n)|~{3,})[^\\n]*\\n").replace("list"," {0,3}(?:[*+-]|1[.)]) ").replace("html",")|<(?:script|pre|style|textarea|!--)").replace("tag",A).getRegex(),L={...q,table:Z,paragraph:k(z).replace("hr",m).replace("heading"," {0,3}#{1,6}(?:\\s|$)").replace("|lheading","").replace("table",Z).replace("blockquote"," {0,3}>").replace("fences"," {0,3}(?:`{3,}(?=[^`\\n]*\\n)|~{3,})[^\\n]*\\n").replace("list"," {0,3}(?:[*+-]|1[.)]) ").replace("html",")|<(?:script|pre|style|textarea|!--)").replace("tag",A).getRegex()},P={...q,html:k("^ *(?:comment *(?:\\n|\\s*$)|<(tag)[\\s\\S]+? *(?:\\n{2,}|\\s*$)|\\s]*)*?/?> *(?:\\n{2,}|\\s*$))").replace("comment",S).replace(/tag/g,"(?!(?:a|em|strong|small|s|cite|q|dfn|abbr|data|time|code|var|samp|kbd|sub|sup|i|b|u|mark|ruby|rt|rp|bdi|bdo|span|br|wbr|ins|del|img)\\b)\\w+(?!:|[^\\w\\s@]*@)\\b").getRegex(),def:/^ *\[([^\]]+)\]: *]+)>?(?: +(["(][^\n]+[")]))? *(?:\n+|$)/,heading:/^(#{1,6})(.*)(?:\n+|$)/,fences:f,lheading:/^(.+?)\n {0,3}(=+|-+) *(?:\n+|$)/,paragraph:k(z).replace("hr",m).replace("heading"," *#{1,6} *[^\n]").replace("lheading",$).replace("|table","").replace("blockquote"," {0,3}>").replace("|fences","").replace("|list","").replace("|html","").replace("|tag","").getRegex()},Q=/^\\([!"#$%&'()*+,\-./:;<=>?@\[\]\\^_`{|}~])/,v=/^( {2,}|\\)\n(?!\s*$)/,B="\\p{P}\\p{S}",C=k(/^((?![*_])[\spunctuation])/,"u").replace(/punctuation/g,B).getRegex(),M=k(/^(?:\*+(?:((?!\*)[punct])|[^\s*]))|^_+(?:((?!_)[punct])|([^\s_]))/,"u").replace(/punct/g,B).getRegex(),O=k("^[^_*]*?__[^_*]*?\\*[^_*]*?(?=__)|[^*]+(?=[^*])|(?!\\*)[punct](\\*+)(?=[\\s]|$)|[^punct\\s](\\*+)(?!\\*)(?=[punct\\s]|$)|(?!\\*)[punct\\s](\\*+)(?=[^punct\\s])|[\\s](\\*+)(?!\\*)(?=[punct])|(?!\\*)[punct](\\*+)(?!\\*)(?=[punct])|[^punct\\s](\\*+)(?=[^punct\\s])","gu").replace(/punct/g,B).getRegex(),D=k("^[^_*]*?\\*\\*[^_*]*?_[^_*]*?(?=\\*\\*)|[^_]+(?=[^_])|(?!_)[punct](_+)(?=[\\s]|$)|[^punct\\s](_+)(?!_)(?=[punct\\s]|$)|(?!_)[punct\\s](_+)(?=[^punct\\s])|[\\s](_+)(?!_)(?=[punct])|(?!_)[punct](_+)(?!_)(?=[punct])","gu").replace(/punct/g,B).getRegex(),j=k(/\\([punct])/,"gu").replace(/punct/g,B).getRegex(),H=k(/^<(scheme:[^\s\x00-\x1f<>]*|email)>/).replace("scheme",/[a-zA-Z][a-zA-Z0-9+.-]{1,31}/).replace("email",/[a-zA-Z0-9.!#$%&'*+/=?^_`{|}~-]+(@)[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)+(?![-_])/).getRegex(),U=k(S).replace("(?:--\x3e|$)","--\x3e").getRegex(),X=k("^comment|^|^<[a-zA-Z][\\w-]*(?:attribute)*?\\s*/?>|^<\\?[\\s\\S]*?\\?>|^|^").replace("comment",U).replace("attribute",/\s+[a-zA-Z:_][\w.:-]*(?:\s*=\s*"[^"]*"|\s*=\s*'[^']*'|\s*=\s*[^\s"'=<>`]+)?/).getRegex(),F=/(?:\[(?:\\.|[^\[\]\\])*\]|\\.|`[^`]*`|[^\[\]\\`])*?/,N=k(/^!?\[(label)\]\(\s*(href)(?:\s+(title))?\s*\)/).replace("label",F).replace("href",/<(?:\\.|[^\n<>\\])+>|[^\s\x00-\x1f]*/).replace("title",/"(?:\\"?|[^"\\])*"|'(?:\\'?|[^'\\])*'|\((?:\\\)?|[^)\\])*\)/).getRegex(),G=k(/^!?\[(label)\]\[(ref)\]/).replace("label",F).replace("ref",T).getRegex(),J=k(/^!?\[(ref)\](?:\[\])?/).replace("ref",T).getRegex(),K={_backpedal:f,anyPunctuation:j,autolink:H,blockSkip:/\[[^[\]]*?\]\([^\(\)]*?\)|`[^`]*?`|<[^<>]*?>/g,br:v,code:/^(`+)([^`]|[^`][\s\S]*?[^`])\1(?!`)/,del:f,emStrongLDelim:M,emStrongRDelimAst:O,emStrongRDelimUnd:D,escape:Q,link:N,nolink:J,punctuation:C,reflink:G,reflinkSearch:k("reflink|nolink(?!\\()","g").replace("reflink",G).replace("nolink",J).getRegex(),tag:X,text:/^(`+|[^`])(?:(?= {2,}\n)|[\s\S]*?(?:(?=[\\t+" ".repeat(n.length)));e;)if(!(this.options.extensions&&this.options.extensions.block&&this.options.extensions.block.some((s=>!!(n=s.call({lexer:this},e,t))&&(e=e.substring(n.raw.length),t.push(n),!0)))))if(n=this.tokenizer.space(e))e=e.substring(n.raw.length),1===n.raw.length&&t.length>0?t[t.length-1].raw+="\n":t.push(n);else if(n=this.tokenizer.code(e))e=e.substring(n.raw.length),s=t[t.length-1],!s||"paragraph"!==s.type&&"text"!==s.type?t.push(n):(s.raw+="\n"+n.raw,s.text+="\n"+n.text,this.inlineQueue[this.inlineQueue.length-1].src=s.text);else if(n=this.tokenizer.fences(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.heading(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.hr(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.blockquote(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.list(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.html(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.def(e))e=e.substring(n.raw.length),s=t[t.length-1],!s||"paragraph"!==s.type&&"text"!==s.type?this.tokens.links[n.tag]||(this.tokens.links[n.tag]={href:n.href,title:n.title}):(s.raw+="\n"+n.raw,s.text+="\n"+n.raw,this.inlineQueue[this.inlineQueue.length-1].src=s.text);else if(n=this.tokenizer.table(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.lheading(e))e=e.substring(n.raw.length),t.push(n);else{if(r=e,this.options.extensions&&this.options.extensions.startBlock){let t=1/0;const n=e.slice(1);let s;this.options.extensions.startBlock.forEach((e=>{s=e.call({lexer:this},n),"number"==typeof s&&s>=0&&(t=Math.min(t,s))})),t<1/0&&t>=0&&(r=e.substring(0,t+1))}if(this.state.top&&(n=this.tokenizer.paragraph(r)))s=t[t.length-1],i&&"paragraph"===s.type?(s.raw+="\n"+n.raw,s.text+="\n"+n.text,this.inlineQueue.pop(),this.inlineQueue[this.inlineQueue.length-1].src=s.text):t.push(n),i=r.length!==e.length,e=e.substring(n.raw.length);else if(n=this.tokenizer.text(e))e=e.substring(n.raw.length),s=t[t.length-1],s&&"text"===s.type?(s.raw+="\n"+n.raw,s.text+="\n"+n.text,this.inlineQueue.pop(),this.inlineQueue[this.inlineQueue.length-1].src=s.text):t.push(n);else if(e){const t="Infinite loop on byte: "+e.charCodeAt(0);if(this.options.silent){console.error(t);break}throw new Error(t)}}return this.state.top=!0,t}inline(e,t=[]){return this.inlineQueue.push({src:e,tokens:t}),t}inlineTokens(e,t=[]){let n,s,r,i,l,o,a=e;if(this.tokens.links){const e=Object.keys(this.tokens.links);if(e.length>0)for(;null!=(i=this.tokenizer.rules.inline.reflinkSearch.exec(a));)e.includes(i[0].slice(i[0].lastIndexOf("[")+1,-1))&&(a=a.slice(0,i.index)+"["+"a".repeat(i[0].length-2)+"]"+a.slice(this.tokenizer.rules.inline.reflinkSearch.lastIndex))}for(;null!=(i=this.tokenizer.rules.inline.blockSkip.exec(a));)a=a.slice(0,i.index)+"["+"a".repeat(i[0].length-2)+"]"+a.slice(this.tokenizer.rules.inline.blockSkip.lastIndex);for(;null!=(i=this.tokenizer.rules.inline.anyPunctuation.exec(a));)a=a.slice(0,i.index)+"++"+a.slice(this.tokenizer.rules.inline.anyPunctuation.lastIndex);for(;e;)if(l||(o=""),l=!1,!(this.options.extensions&&this.options.extensions.inline&&this.options.extensions.inline.some((s=>!!(n=s.call({lexer:this},e,t))&&(e=e.substring(n.raw.length),t.push(n),!0)))))if(n=this.tokenizer.escape(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.tag(e))e=e.substring(n.raw.length),s=t[t.length-1],s&&"text"===n.type&&"text"===s.type?(s.raw+=n.raw,s.text+=n.text):t.push(n);else if(n=this.tokenizer.link(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.reflink(e,this.tokens.links))e=e.substring(n.raw.length),s=t[t.length-1],s&&"text"===n.type&&"text"===s.type?(s.raw+=n.raw,s.text+=n.text):t.push(n);else if(n=this.tokenizer.emStrong(e,a,o))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.codespan(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.br(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.del(e))e=e.substring(n.raw.length),t.push(n);else if(n=this.tokenizer.autolink(e))e=e.substring(n.raw.length),t.push(n);else if(this.state.inLink||!(n=this.tokenizer.url(e))){if(r=e,this.options.extensions&&this.options.extensions.startInline){let t=1/0;const n=e.slice(1);let s;this.options.extensions.startInline.forEach((e=>{s=e.call({lexer:this},n),"number"==typeof s&&s>=0&&(t=Math.min(t,s))})),t<1/0&&t>=0&&(r=e.substring(0,t+1))}if(n=this.tokenizer.inlineText(r))e=e.substring(n.raw.length),"_"!==n.raw.slice(-1)&&(o=n.raw.slice(-1)),l=!0,s=t[t.length-1],s&&"text"===s.type?(s.raw+=n.raw,s.text+=n.text):t.push(n);else if(e){const t="Infinite loop on byte: "+e.charCodeAt(0);if(this.options.silent){console.error(t);break}throw new Error(t)}}else e=e.substring(n.raw.length),t.push(n);return t}}class se{options;constructor(t){this.options=t||e.defaults}code(e,t,n){const s=(t||"").match(/^\S*/)?.[0];return e=e.replace(/\n$/,"")+"\n",s?'
'+(n?e:c(e,!0))+"
\n":"
"+(n?e:c(e,!0))+"
\n"}blockquote(e){return`
\n${e}
\n`}html(e,t){return e}heading(e,t,n){return`${e}\n`}hr(){return"
\n"}list(e,t,n){const s=t?"ol":"ul";return"<"+s+(t&&1!==n?' start="'+n+'"':"")+">\n"+e+"\n"}listitem(e,t,n){return`
  • ${e}
  • \n`}checkbox(e){return"'}paragraph(e){return`

    ${e}

    \n`}table(e,t){return t&&(t=`${t}`),"\n\n"+e+"\n"+t+"
    \n"}tablerow(e){return`\n${e}\n`}tablecell(e,t){const n=t.header?"th":"td";return(t.align?`<${n} align="${t.align}">`:`<${n}>`)+e+`\n`}strong(e){return`${e}`}em(e){return`${e}`}codespan(e){return`${e}`}br(){return"
    "}del(e){return`${e}`}link(e,t,n){const s=g(e);if(null===s)return n;let r='
    ",r}image(e,t,n){const s=g(e);if(null===s)return n;let r=`${n}0&&"paragraph"===n.tokens[0].type?(n.tokens[0].text=e+" "+n.tokens[0].text,n.tokens[0].tokens&&n.tokens[0].tokens.length>0&&"text"===n.tokens[0].tokens[0].type&&(n.tokens[0].tokens[0].text=e+" "+n.tokens[0].tokens[0].text)):n.tokens.unshift({type:"text",text:e+" "}):o+=e+" "}o+=this.parse(n.tokens,i),l+=this.renderer.listitem(o,r,!!s)}n+=this.renderer.list(l,t,s);continue}case"html":{const e=r;n+=this.renderer.html(e.text,e.block);continue}case"paragraph":{const e=r;n+=this.renderer.paragraph(this.parseInline(e.tokens));continue}case"text":{let i=r,l=i.tokens?this.parseInline(i.tokens):i.text;for(;s+1{const r=e[s].flat(1/0);n=n.concat(this.walkTokens(r,t))})):e.tokens&&(n=n.concat(this.walkTokens(e.tokens,t)))}}return n}use(...e){const t=this.defaults.extensions||{renderers:{},childTokens:{}};return e.forEach((e=>{const n={...e};if(n.async=this.defaults.async||n.async||!1,e.extensions&&(e.extensions.forEach((e=>{if(!e.name)throw new Error("extension name required");if("renderer"in e){const n=t.renderers[e.name];t.renderers[e.name]=n?function(...t){let s=e.renderer.apply(this,t);return!1===s&&(s=n.apply(this,t)),s}:e.renderer}if("tokenizer"in e){if(!e.level||"block"!==e.level&&"inline"!==e.level)throw new Error("extension level must be 'block' or 'inline'");const n=t[e.level];n?n.unshift(e.tokenizer):t[e.level]=[e.tokenizer],e.start&&("block"===e.level?t.startBlock?t.startBlock.push(e.start):t.startBlock=[e.start]:"inline"===e.level&&(t.startInline?t.startInline.push(e.start):t.startInline=[e.start]))}"childTokens"in e&&e.childTokens&&(t.childTokens[e.name]=e.childTokens)})),n.extensions=t),e.renderer){const t=this.defaults.renderer||new se(this.defaults);for(const n in e.renderer){if(!(n in t))throw new Error(`renderer '${n}' does not exist`);if("options"===n)continue;const s=n,r=e.renderer[s],i=t[s];t[s]=(...e)=>{let n=r.apply(t,e);return!1===n&&(n=i.apply(t,e)),n||""}}n.renderer=t}if(e.tokenizer){const t=this.defaults.tokenizer||new w(this.defaults);for(const n in e.tokenizer){if(!(n in t))throw new Error(`tokenizer '${n}' does not exist`);if(["options","rules","lexer"].includes(n))continue;const s=n,r=e.tokenizer[s],i=t[s];t[s]=(...e)=>{let n=r.apply(t,e);return!1===n&&(n=i.apply(t,e)),n}}n.tokenizer=t}if(e.hooks){const t=this.defaults.hooks||new le;for(const n in e.hooks){if(!(n in t))throw new Error(`hook '${n}' does not exist`);if("options"===n)continue;const s=n,r=e.hooks[s],i=t[s];le.passThroughHooks.has(n)?t[s]=e=>{if(this.defaults.async)return Promise.resolve(r.call(t,e)).then((e=>i.call(t,e)));const n=r.call(t,e);return i.call(t,n)}:t[s]=(...e)=>{let n=r.apply(t,e);return!1===n&&(n=i.apply(t,e)),n}}n.hooks=t}if(e.walkTokens){const t=this.defaults.walkTokens,s=e.walkTokens;n.walkTokens=function(e){let n=[];return n.push(s.call(this,e)),t&&(n=n.concat(t.call(this,e))),n}}this.defaults={...this.defaults,...n}})),this}setOptions(e){return this.defaults={...this.defaults,...e},this}lexer(e,t){return ne.lex(e,t??this.defaults)}parser(e,t){return ie.parse(e,t??this.defaults)}#e(e,t){return(n,s)=>{const r={...s},i={...this.defaults,...r};!0===this.defaults.async&&!1===r.async&&(i.silent||console.warn("marked(): The async option was set to true by an extension. The async: false option sent to parse will be ignored."),i.async=!0);const l=this.#t(!!i.silent,!!i.async);if(null==n)return l(new Error("marked(): input parameter is undefined or null"));if("string"!=typeof n)return l(new Error("marked(): input parameter is of type "+Object.prototype.toString.call(n)+", string expected"));if(i.hooks&&(i.hooks.options=i),i.async)return Promise.resolve(i.hooks?i.hooks.preprocess(n):n).then((t=>e(t,i))).then((e=>i.hooks?i.hooks.processAllTokens(e):e)).then((e=>i.walkTokens?Promise.all(this.walkTokens(e,i.walkTokens)).then((()=>e)):e)).then((e=>t(e,i))).then((e=>i.hooks?i.hooks.postprocess(e):e)).catch(l);try{i.hooks&&(n=i.hooks.preprocess(n));let s=e(n,i);i.hooks&&(s=i.hooks.processAllTokens(s)),i.walkTokens&&this.walkTokens(s,i.walkTokens);let r=t(s,i);return i.hooks&&(r=i.hooks.postprocess(r)),r}catch(e){return l(e)}}}#t(e,t){return n=>{if(n.message+="\nPlease report this to https://github.com/markedjs/marked.",e){const e="

    An error occurred:

    "+c(n.message+"",!0)+"
    ";return t?Promise.resolve(e):e}if(t)return Promise.reject(n);throw n}}}const ae=new oe;function ce(e,t){return ae.parse(e,t)}ce.options=ce.setOptions=function(e){return ae.setOptions(e),ce.defaults=ae.defaults,n(ce.defaults),ce},ce.getDefaults=t,ce.defaults=e.defaults,ce.use=function(...e){return ae.use(...e),ce.defaults=ae.defaults,n(ce.defaults),ce},ce.walkTokens=function(e,t){return ae.walkTokens(e,t)},ce.parseInline=ae.parseInline,ce.Parser=ie,ce.parser=ie.parse,ce.Renderer=se,ce.TextRenderer=re,ce.Lexer=ne,ce.lexer=ne.lex,ce.Tokenizer=w,ce.Hooks=le,ce.parse=ce;const he=ce.options,pe=ce.setOptions,ue=ce.use,ke=ce.walkTokens,ge=ce.parseInline,fe=ce,de=ie.parse,xe=ne.lex;e.Hooks=le,e.Lexer=ne,e.Marked=oe,e.Parser=ie,e.Renderer=se,e.TextRenderer=re,e.Tokenizer=w,e.getDefaults=t,e.lexer=xe,e.marked=ce,e.options=he,e.parse=fe,e.parseInline=ge,e.parser=de,e.setOptions=pe,e.use=ue,e.walkTokens=ke})); diff --git a/custom_nodes/ComfyUI-KJNodes/kjweb_async/purify.min.js b/custom_nodes/ComfyUI-KJNodes/kjweb_async/purify.min.js new file mode 100644 index 0000000000000000000000000000000000000000..c2f5164618eebcc44b0186f594ccb8092639c670 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/kjweb_async/purify.min.js @@ -0,0 +1,3 @@ +/*! @license DOMPurify 3.0.11 | (c) Cure53 and other contributors | Released under the Apache license 2.0 and Mozilla Public License 2.0 | github.com/cure53/DOMPurify/blob/3.0.11/LICENSE */ +!function(e,t){"object"==typeof exports&&"undefined"!=typeof module?module.exports=t():"function"==typeof define&&define.amd?define(t):(e="undefined"!=typeof globalThis?globalThis:e||self).DOMPurify=t()}(this,(function(){"use strict";const{entries:e,setPrototypeOf:t,isFrozen:n,getPrototypeOf:o,getOwnPropertyDescriptor:r}=Object;let{freeze:i,seal:a,create:l}=Object,{apply:c,construct:s}="undefined"!=typeof Reflect&&Reflect;i||(i=function(e){return e}),a||(a=function(e){return e}),c||(c=function(e,t,n){return e.apply(t,n)}),s||(s=function(e,t){return new e(...t)});const u=b(Array.prototype.forEach),m=b(Array.prototype.pop),p=b(Array.prototype.push),f=b(String.prototype.toLowerCase),d=b(String.prototype.toString),h=b(String.prototype.match),g=b(String.prototype.replace),T=b(String.prototype.indexOf),y=b(String.prototype.trim),E=b(Object.prototype.hasOwnProperty),A=b(RegExp.prototype.test),_=(N=TypeError,function(){for(var e=arguments.length,t=new Array(e),n=0;n1?n-1:0),r=1;r2&&void 0!==arguments[2]?arguments[2]:f;t&&t(e,null);let i=o.length;for(;i--;){let t=o[i];if("string"==typeof t){const e=r(t);e!==t&&(n(o)||(o[i]=e),t=e)}e[t]=!0}return e}function R(e){for(let t=0;t/gm),B=a(/\${[\w\W]*}/gm),W=a(/^data-[\-\w.\u00B7-\uFFFF]/),G=a(/^aria-[\-\w]+$/),Y=a(/^(?:(?:(?:f|ht)tps?|mailto|tel|callto|sms|cid|xmpp):|[^a-z]|[a-z+.\-]+(?:[^a-z+.\-:]|$))/i),j=a(/^(?:\w+script|data):/i),X=a(/[\u0000-\u0020\u00A0\u1680\u180E\u2000-\u2029\u205F\u3000]/g),q=a(/^html$/i),$=a(/^[a-z][.\w]*(-[.\w]+)+$/i);var K=Object.freeze({__proto__:null,MUSTACHE_EXPR:H,ERB_EXPR:z,TMPLIT_EXPR:B,DATA_ATTR:W,ARIA_ATTR:G,IS_ALLOWED_URI:Y,IS_SCRIPT_OR_DATA:j,ATTR_WHITESPACE:X,DOCTYPE_NAME:q,CUSTOM_ELEMENT:$});const V=function(){return"undefined"==typeof window?null:window},Z=function(e,t){if("object"!=typeof e||"function"!=typeof e.createPolicy)return null;let n=null;const o="data-tt-policy-suffix";t&&t.hasAttribute(o)&&(n=t.getAttribute(o));const r="dompurify"+(n?"#"+n:"");try{return e.createPolicy(r,{createHTML:e=>e,createScriptURL:e=>e})}catch(e){return console.warn("TrustedTypes policy "+r+" could not be created."),null}};var J=function t(){let n=arguments.length>0&&void 0!==arguments[0]?arguments[0]:V();const o=e=>t(e);if(o.version="3.0.11",o.removed=[],!n||!n.document||9!==n.document.nodeType)return o.isSupported=!1,o;let{document:r}=n;const a=r,c=a.currentScript,{DocumentFragment:s,HTMLTemplateElement:N,Node:b,Element:R,NodeFilter:H,NamedNodeMap:z=n.NamedNodeMap||n.MozNamedAttrMap,HTMLFormElement:B,DOMParser:W,trustedTypes:G}=n,j=R.prototype,X=L(j,"cloneNode"),$=L(j,"nextSibling"),J=L(j,"childNodes"),Q=L(j,"parentNode");if("function"==typeof N){const e=r.createElement("template");e.content&&e.content.ownerDocument&&(r=e.content.ownerDocument)}let ee,te="";const{implementation:ne,createNodeIterator:oe,createDocumentFragment:re,getElementsByTagName:ie}=r,{importNode:ae}=a;let le={};o.isSupported="function"==typeof e&&"function"==typeof Q&&ne&&void 0!==ne.createHTMLDocument;const{MUSTACHE_EXPR:ce,ERB_EXPR:se,TMPLIT_EXPR:ue,DATA_ATTR:me,ARIA_ATTR:pe,IS_SCRIPT_OR_DATA:fe,ATTR_WHITESPACE:de,CUSTOM_ELEMENT:he}=K;let{IS_ALLOWED_URI:ge}=K,Te=null;const ye=S({},[...D,...C,...O,...v,...M]);let Ee=null;const Ae=S({},[...I,...U,...P,...F]);let _e=Object.seal(l(null,{tagNameCheck:{writable:!0,configurable:!1,enumerable:!0,value:null},attributeNameCheck:{writable:!0,configurable:!1,enumerable:!0,value:null},allowCustomizedBuiltInElements:{writable:!0,configurable:!1,enumerable:!0,value:!1}})),Ne=null,be=null,Se=!0,Re=!0,we=!1,Le=!0,De=!1,Ce=!0,Oe=!1,xe=!1,ve=!1,ke=!1,Me=!1,Ie=!1,Ue=!0,Pe=!1;const Fe="user-content-";let He=!0,ze=!1,Be={},We=null;const Ge=S({},["annotation-xml","audio","colgroup","desc","foreignobject","head","iframe","math","mi","mn","mo","ms","mtext","noembed","noframes","noscript","plaintext","script","style","svg","template","thead","title","video","xmp"]);let Ye=null;const je=S({},["audio","video","img","source","image","track"]);let Xe=null;const qe=S({},["alt","class","for","id","label","name","pattern","placeholder","role","summary","title","value","style","xmlns"]),$e="http://www.w3.org/1998/Math/MathML",Ke="http://www.w3.org/2000/svg",Ve="http://www.w3.org/1999/xhtml";let Ze=Ve,Je=!1,Qe=null;const et=S({},[$e,Ke,Ve],d);let tt=null;const nt=["application/xhtml+xml","text/html"],ot="text/html";let rt=null,it=null;const at=r.createElement("form"),lt=function(e){return e instanceof RegExp||e instanceof Function},ct=function(){let e=arguments.length>0&&void 0!==arguments[0]?arguments[0]:{};if(!it||it!==e){if(e&&"object"==typeof e||(e={}),e=w(e),tt=-1===nt.indexOf(e.PARSER_MEDIA_TYPE)?ot:e.PARSER_MEDIA_TYPE,rt="application/xhtml+xml"===tt?d:f,Te=E(e,"ALLOWED_TAGS")?S({},e.ALLOWED_TAGS,rt):ye,Ee=E(e,"ALLOWED_ATTR")?S({},e.ALLOWED_ATTR,rt):Ae,Qe=E(e,"ALLOWED_NAMESPACES")?S({},e.ALLOWED_NAMESPACES,d):et,Xe=E(e,"ADD_URI_SAFE_ATTR")?S(w(qe),e.ADD_URI_SAFE_ATTR,rt):qe,Ye=E(e,"ADD_DATA_URI_TAGS")?S(w(je),e.ADD_DATA_URI_TAGS,rt):je,We=E(e,"FORBID_CONTENTS")?S({},e.FORBID_CONTENTS,rt):Ge,Ne=E(e,"FORBID_TAGS")?S({},e.FORBID_TAGS,rt):{},be=E(e,"FORBID_ATTR")?S({},e.FORBID_ATTR,rt):{},Be=!!E(e,"USE_PROFILES")&&e.USE_PROFILES,Se=!1!==e.ALLOW_ARIA_ATTR,Re=!1!==e.ALLOW_DATA_ATTR,we=e.ALLOW_UNKNOWN_PROTOCOLS||!1,Le=!1!==e.ALLOW_SELF_CLOSE_IN_ATTR,De=e.SAFE_FOR_TEMPLATES||!1,Ce=!1!==e.SAFE_FOR_XML,Oe=e.WHOLE_DOCUMENT||!1,ke=e.RETURN_DOM||!1,Me=e.RETURN_DOM_FRAGMENT||!1,Ie=e.RETURN_TRUSTED_TYPE||!1,ve=e.FORCE_BODY||!1,Ue=!1!==e.SANITIZE_DOM,Pe=e.SANITIZE_NAMED_PROPS||!1,He=!1!==e.KEEP_CONTENT,ze=e.IN_PLACE||!1,ge=e.ALLOWED_URI_REGEXP||Y,Ze=e.NAMESPACE||Ve,_e=e.CUSTOM_ELEMENT_HANDLING||{},e.CUSTOM_ELEMENT_HANDLING&<(e.CUSTOM_ELEMENT_HANDLING.tagNameCheck)&&(_e.tagNameCheck=e.CUSTOM_ELEMENT_HANDLING.tagNameCheck),e.CUSTOM_ELEMENT_HANDLING&<(e.CUSTOM_ELEMENT_HANDLING.attributeNameCheck)&&(_e.attributeNameCheck=e.CUSTOM_ELEMENT_HANDLING.attributeNameCheck),e.CUSTOM_ELEMENT_HANDLING&&"boolean"==typeof e.CUSTOM_ELEMENT_HANDLING.allowCustomizedBuiltInElements&&(_e.allowCustomizedBuiltInElements=e.CUSTOM_ELEMENT_HANDLING.allowCustomizedBuiltInElements),De&&(Re=!1),Me&&(ke=!0),Be&&(Te=S({},M),Ee=[],!0===Be.html&&(S(Te,D),S(Ee,I)),!0===Be.svg&&(S(Te,C),S(Ee,U),S(Ee,F)),!0===Be.svgFilters&&(S(Te,O),S(Ee,U),S(Ee,F)),!0===Be.mathMl&&(S(Te,v),S(Ee,P),S(Ee,F))),e.ADD_TAGS&&(Te===ye&&(Te=w(Te)),S(Te,e.ADD_TAGS,rt)),e.ADD_ATTR&&(Ee===Ae&&(Ee=w(Ee)),S(Ee,e.ADD_ATTR,rt)),e.ADD_URI_SAFE_ATTR&&S(Xe,e.ADD_URI_SAFE_ATTR,rt),e.FORBID_CONTENTS&&(We===Ge&&(We=w(We)),S(We,e.FORBID_CONTENTS,rt)),He&&(Te["#text"]=!0),Oe&&S(Te,["html","head","body"]),Te.table&&(S(Te,["tbody"]),delete Ne.tbody),e.TRUSTED_TYPES_POLICY){if("function"!=typeof e.TRUSTED_TYPES_POLICY.createHTML)throw _('TRUSTED_TYPES_POLICY configuration option must provide a "createHTML" hook.');if("function"!=typeof e.TRUSTED_TYPES_POLICY.createScriptURL)throw _('TRUSTED_TYPES_POLICY configuration option must provide a "createScriptURL" hook.');ee=e.TRUSTED_TYPES_POLICY,te=ee.createHTML("")}else void 0===ee&&(ee=Z(G,c)),null!==ee&&"string"==typeof te&&(te=ee.createHTML(""));i&&i(e),it=e}},st=S({},["mi","mo","mn","ms","mtext"]),ut=S({},["foreignobject","desc","title","annotation-xml"]),mt=S({},["title","style","font","a","script"]),pt=S({},[...C,...O,...x]),ft=S({},[...v,...k]),dt=function(e){let t=Q(e);t&&t.tagName||(t={namespaceURI:Ze,tagName:"template"});const n=f(e.tagName),o=f(t.tagName);return!!Qe[e.namespaceURI]&&(e.namespaceURI===Ke?t.namespaceURI===Ve?"svg"===n:t.namespaceURI===$e?"svg"===n&&("annotation-xml"===o||st[o]):Boolean(pt[n]):e.namespaceURI===$e?t.namespaceURI===Ve?"math"===n:t.namespaceURI===Ke?"math"===n&&ut[o]:Boolean(ft[n]):e.namespaceURI===Ve?!(t.namespaceURI===Ke&&!ut[o])&&(!(t.namespaceURI===$e&&!st[o])&&(!ft[n]&&(mt[n]||!pt[n]))):!("application/xhtml+xml"!==tt||!Qe[e.namespaceURI]))},ht=function(e){p(o.removed,{element:e});try{e.parentNode.removeChild(e)}catch(t){e.remove()}},gt=function(e,t){try{p(o.removed,{attribute:t.getAttributeNode(e),from:t})}catch(e){p(o.removed,{attribute:null,from:t})}if(t.removeAttribute(e),"is"===e&&!Ee[e])if(ke||Me)try{ht(t)}catch(e){}else try{t.setAttribute(e,"")}catch(e){}},Tt=function(e){let t=null,n=null;if(ve)e=""+e;else{const t=h(e,/^[\r\n\t ]+/);n=t&&t[0]}"application/xhtml+xml"===tt&&Ze===Ve&&(e=''+e+"");const o=ee?ee.createHTML(e):e;if(Ze===Ve)try{t=(new W).parseFromString(o,tt)}catch(e){}if(!t||!t.documentElement){t=ne.createDocument(Ze,"template",null);try{t.documentElement.innerHTML=Je?te:o}catch(e){}}const i=t.body||t.documentElement;return e&&n&&i.insertBefore(r.createTextNode(n),i.childNodes[0]||null),Ze===Ve?ie.call(t,Oe?"html":"body")[0]:Oe?t.documentElement:i},yt=function(e){return oe.call(e.ownerDocument||e,e,H.SHOW_ELEMENT|H.SHOW_COMMENT|H.SHOW_TEXT|H.SHOW_PROCESSING_INSTRUCTION|H.SHOW_CDATA_SECTION,null)},Et=function(e){return e instanceof B&&("string"!=typeof e.nodeName||"string"!=typeof e.textContent||"function"!=typeof e.removeChild||!(e.attributes instanceof z)||"function"!=typeof e.removeAttribute||"function"!=typeof e.setAttribute||"string"!=typeof e.namespaceURI||"function"!=typeof e.insertBefore||"function"!=typeof e.hasChildNodes)},At=function(e){return"function"==typeof b&&e instanceof b},_t=function(e,t,n){le[e]&&u(le[e],(e=>{e.call(o,t,n,it)}))},Nt=function(e){let t=null;if(_t("beforeSanitizeElements",e,null),Et(e))return ht(e),!0;const n=rt(e.nodeName);if(_t("uponSanitizeElement",e,{tagName:n,allowedTags:Te}),e.hasChildNodes()&&!At(e.firstElementChild)&&A(/<[/\w]/g,e.innerHTML)&&A(/<[/\w]/g,e.textContent))return ht(e),!0;if(7===e.nodeType)return ht(e),!0;if(Ce&&8===e.nodeType&&A(/<[/\w]/g,e.data))return ht(e),!0;if(!Te[n]||Ne[n]){if(!Ne[n]&&St(n)){if(_e.tagNameCheck instanceof RegExp&&A(_e.tagNameCheck,n))return!1;if(_e.tagNameCheck instanceof Function&&_e.tagNameCheck(n))return!1}if(He&&!We[n]){const t=Q(e)||e.parentNode,n=J(e)||e.childNodes;if(n&&t){for(let o=n.length-1;o>=0;--o)t.insertBefore(X(n[o],!0),$(e))}}return ht(e),!0}return e instanceof R&&!dt(e)?(ht(e),!0):"noscript"!==n&&"noembed"!==n&&"noframes"!==n||!A(/<\/no(script|embed|frames)/i,e.innerHTML)?(De&&3===e.nodeType&&(t=e.textContent,u([ce,se,ue],(e=>{t=g(t,e," ")})),e.textContent!==t&&(p(o.removed,{element:e.cloneNode()}),e.textContent=t)),_t("afterSanitizeElements",e,null),!1):(ht(e),!0)},bt=function(e,t,n){if(Ue&&("id"===t||"name"===t)&&(n in r||n in at))return!1;if(Re&&!be[t]&&A(me,t));else if(Se&&A(pe,t));else if(!Ee[t]||be[t]){if(!(St(e)&&(_e.tagNameCheck instanceof RegExp&&A(_e.tagNameCheck,e)||_e.tagNameCheck instanceof Function&&_e.tagNameCheck(e))&&(_e.attributeNameCheck instanceof RegExp&&A(_e.attributeNameCheck,t)||_e.attributeNameCheck instanceof Function&&_e.attributeNameCheck(t))||"is"===t&&_e.allowCustomizedBuiltInElements&&(_e.tagNameCheck instanceof RegExp&&A(_e.tagNameCheck,n)||_e.tagNameCheck instanceof Function&&_e.tagNameCheck(n))))return!1}else if(Xe[t]);else if(A(ge,g(n,de,"")));else if("src"!==t&&"xlink:href"!==t&&"href"!==t||"script"===e||0!==T(n,"data:")||!Ye[e]){if(we&&!A(fe,g(n,de,"")));else if(n)return!1}else;return!0},St=function(e){return"annotation-xml"!==e&&h(e,he)},Rt=function(e){_t("beforeSanitizeAttributes",e,null);const{attributes:t}=e;if(!t)return;const n={attrName:"",attrValue:"",keepAttr:!0,allowedAttributes:Ee};let r=t.length;for(;r--;){const i=t[r],{name:a,namespaceURI:l,value:c}=i,s=rt(a);let p="value"===a?c:y(c);if(n.attrName=s,n.attrValue=p,n.keepAttr=!0,n.forceKeepAttr=void 0,_t("uponSanitizeAttribute",e,n),p=n.attrValue,n.forceKeepAttr)continue;if(gt(a,e),!n.keepAttr)continue;if(!Le&&A(/\/>/i,p)){gt(a,e);continue}De&&u([ce,se,ue],(e=>{p=g(p,e," ")}));const f=rt(e.nodeName);if(bt(f,s,p)){if(!Pe||"id"!==s&&"name"!==s||(gt(a,e),p=Fe+p),ee&&"object"==typeof G&&"function"==typeof G.getAttributeType)if(l);else switch(G.getAttributeType(f,s)){case"TrustedHTML":p=ee.createHTML(p);break;case"TrustedScriptURL":p=ee.createScriptURL(p)}try{l?e.setAttributeNS(l,a,p):e.setAttribute(a,p),m(o.removed)}catch(e){}}}_t("afterSanitizeAttributes",e,null)},wt=function e(t){let n=null;const o=yt(t);for(_t("beforeSanitizeShadowDOM",t,null);n=o.nextNode();)_t("uponSanitizeShadowNode",n,null),Nt(n)||(n.content instanceof s&&e(n.content),Rt(n));_t("afterSanitizeShadowDOM",t,null)};return o.sanitize=function(e){let t=arguments.length>1&&void 0!==arguments[1]?arguments[1]:{},n=null,r=null,i=null,l=null;if(Je=!e,Je&&(e="\x3c!--\x3e"),"string"!=typeof e&&!At(e)){if("function"!=typeof e.toString)throw _("toString is not a function");if("string"!=typeof(e=e.toString()))throw _("dirty is not a string, aborting")}if(!o.isSupported)return e;if(xe||ct(t),o.removed=[],"string"==typeof e&&(ze=!1),ze){if(e.nodeName){const t=rt(e.nodeName);if(!Te[t]||Ne[t])throw _("root node is forbidden and cannot be sanitized in-place")}}else if(e instanceof b)n=Tt("\x3c!----\x3e"),r=n.ownerDocument.importNode(e,!0),1===r.nodeType&&"BODY"===r.nodeName||"HTML"===r.nodeName?n=r:n.appendChild(r);else{if(!ke&&!De&&!Oe&&-1===e.indexOf("<"))return ee&&Ie?ee.createHTML(e):e;if(n=Tt(e),!n)return ke?null:Ie?te:""}n&&ve&&ht(n.firstChild);const c=yt(ze?e:n);for(;i=c.nextNode();)Nt(i)||(i.content instanceof s&&wt(i.content),Rt(i));if(ze)return e;if(ke){if(Me)for(l=re.call(n.ownerDocument);n.firstChild;)l.appendChild(n.firstChild);else l=n;return(Ee.shadowroot||Ee.shadowrootmode)&&(l=ae.call(a,l,!0)),l}let m=Oe?n.outerHTML:n.innerHTML;return Oe&&Te["!doctype"]&&n.ownerDocument&&n.ownerDocument.doctype&&n.ownerDocument.doctype.name&&A(q,n.ownerDocument.doctype.name)&&(m="\n"+m),De&&u([ce,se,ue],(e=>{m=g(m,e," ")})),ee&&Ie?ee.createHTML(m):m},o.setConfig=function(){let e=arguments.length>0&&void 0!==arguments[0]?arguments[0]:{};ct(e),xe=!0},o.clearConfig=function(){it=null,xe=!1},o.isValidAttribute=function(e,t,n){it||ct({});const o=rt(e),r=rt(t);return bt(o,r,n)},o.addHook=function(e,t){"function"==typeof t&&(le[e]=le[e]||[],p(le[e],t))},o.removeHook=function(e){if(le[e])return m(le[e])},o.removeHooks=function(e){le[e]&&(le[e]=[])},o.removeAllHooks=function(){le={}},o}();return J})); +//# sourceMappingURL=purify.min.js.map diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/audioscheduler_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/audioscheduler_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..69d0422e7da875298f87fe60a7f6d1494530dca2 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/audioscheduler_nodes.py @@ -0,0 +1,251 @@ +# to be used with https://github.com/a1lazydog/ComfyUI-AudioScheduler +import torch +from torchvision.transforms import functional as TF +from PIL import Image, ImageDraw +import numpy as np +from ..utility.utility import pil2tensor +from nodes import MAX_RESOLUTION + +class NormalizedAmplitudeToMask: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "normalized_amp": ("NORMALIZED_AMPLITUDE",), + "width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_offset": ("INT", {"default": 0,"min": -255, "max": 255, "step": 1}), + "location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), + "location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), + "size": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + "shape": ( + [ + 'none', + 'circle', + 'square', + 'triangle', + ], + { + "default": 'none' + }), + "color": ( + [ + 'white', + 'amplitude', + ], + { + "default": 'amplitude' + }), + },} + + CATEGORY = "KJNodes/audio" + RETURN_TYPES = ("MASK",) + FUNCTION = "convert" + DESCRIPTION = """ +Works as a bridge to the AudioScheduler -nodes: +https://github.com/a1lazydog/ComfyUI-AudioScheduler +Creates masks based on the normalized amplitude. +""" + + def convert(self, normalized_amp, width, height, frame_offset, shape, location_x, location_y, size, color): + # Ensure normalized_amp is an array and within the range [0, 1] + normalized_amp = np.clip(normalized_amp, 0.0, 1.0) + + # Offset the amplitude values by rolling the array + normalized_amp = np.roll(normalized_amp, frame_offset) + + # Initialize an empty list to hold the image tensors + out = [] + # Iterate over each amplitude value to create an image + for amp in normalized_amp: + # Scale the amplitude value to cover the full range of grayscale values + if color == 'amplitude': + grayscale_value = int(amp * 255) + elif color == 'white': + grayscale_value = 255 + # Convert the grayscale value to an RGB format + gray_color = (grayscale_value, grayscale_value, grayscale_value) + finalsize = size * amp + + if shape == 'none': + shapeimage = Image.new("RGB", (width, height), gray_color) + else: + shapeimage = Image.new("RGB", (width, height), "black") + + draw = ImageDraw.Draw(shapeimage) + if shape == 'circle' or shape == 'square': + # Define the bounding box for the shape + left_up_point = (location_x - finalsize, location_y - finalsize) + right_down_point = (location_x + finalsize,location_y + finalsize) + two_points = [left_up_point, right_down_point] + + if shape == 'circle': + draw.ellipse(two_points, fill=gray_color) + elif shape == 'square': + draw.rectangle(two_points, fill=gray_color) + + elif shape == 'triangle': + # Define the points for the triangle + left_up_point = (location_x - finalsize, location_y + finalsize) # bottom left + right_down_point = (location_x + finalsize, location_y + finalsize) # bottom right + top_point = (location_x, location_y) # top point + draw.polygon([top_point, left_up_point, right_down_point], fill=gray_color) + + shapeimage = pil2tensor(shapeimage) + mask = shapeimage[:, :, :, 0] + out.append(mask) + + return (torch.cat(out, dim=0),) + +class NormalizedAmplitudeToFloatList: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "normalized_amp": ("NORMALIZED_AMPLITUDE",), + },} + + CATEGORY = "KJNodes/audio" + RETURN_TYPES = ("FLOAT",) + FUNCTION = "convert" + DESCRIPTION = """ +Works as a bridge to the AudioScheduler -nodes: +https://github.com/a1lazydog/ComfyUI-AudioScheduler +Creates a list of floats from the normalized amplitude. +""" + + def convert(self, normalized_amp): + # Ensure normalized_amp is an array and within the range [0, 1] + normalized_amp = np.clip(normalized_amp, 0.0, 1.0) + return (normalized_amp.tolist(),) + +class OffsetMaskByNormalizedAmplitude: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "normalized_amp": ("NORMALIZED_AMPLITUDE",), + "mask": ("MASK",), + "x": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "y": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "rotate": ("BOOLEAN", { "default": False }), + "angle_multiplier": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001, "display": "number" }), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "offset" + CATEGORY = "KJNodes/audio" + DESCRIPTION = """ +Works as a bridge to the AudioScheduler -nodes: +https://github.com/a1lazydog/ComfyUI-AudioScheduler +Offsets masks based on the normalized amplitude. +""" + + def offset(self, mask, x, y, angle_multiplier, rotate, normalized_amp): + + # Ensure normalized_amp is an array and within the range [0, 1] + offsetmask = mask.clone() + normalized_amp = np.clip(normalized_amp, 0.0, 1.0) + + batch_size, height, width = mask.shape + + if rotate: + for i in range(batch_size): + rotation_amp = int(normalized_amp[i] * (360 * angle_multiplier)) + rotation_angle = rotation_amp + offsetmask[i] = TF.rotate(offsetmask[i].unsqueeze(0), rotation_angle).squeeze(0) + if x != 0 or y != 0: + for i in range(batch_size): + offset_amp = normalized_amp[i] * 10 + shift_x = min(x*offset_amp, width-1) + shift_y = min(y*offset_amp, height-1) + if shift_x != 0: + offsetmask[i] = torch.roll(offsetmask[i], shifts=int(shift_x), dims=1) + if shift_y != 0: + offsetmask[i] = torch.roll(offsetmask[i], shifts=int(shift_y), dims=0) + + return offsetmask, + +class ImageTransformByNormalizedAmplitude: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "normalized_amp": ("NORMALIZED_AMPLITUDE",), + "zoom_scale": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001, "display": "number" }), + "x_offset": ("INT", { "default": 0, "min": (1 -MAX_RESOLUTION), "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "y_offset": ("INT", { "default": 0, "min": (1 -MAX_RESOLUTION), "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "cumulative": ("BOOLEAN", { "default": False }), + "image": ("IMAGE",), + }} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "amptransform" + CATEGORY = "KJNodes/audio" + DESCRIPTION = """ +Works as a bridge to the AudioScheduler -nodes: +https://github.com/a1lazydog/ComfyUI-AudioScheduler +Transforms image based on the normalized amplitude. +""" + + def amptransform(self, image, normalized_amp, zoom_scale, cumulative, x_offset, y_offset): + # Ensure normalized_amp is an array and within the range [0, 1] + normalized_amp = np.clip(normalized_amp, 0.0, 1.0) + transformed_images = [] + + # Initialize the cumulative zoom factor + prev_amp = 0.0 + + for i in range(image.shape[0]): + img = image[i] # Get the i-th image in the batch + amp = normalized_amp[i] # Get the corresponding amplitude value + + # Incrementally increase the cumulative zoom factor + if cumulative: + prev_amp += amp + amp += prev_amp + + # Convert the image tensor from BxHxWxC to CxHxW format expected by torchvision + img = img.permute(2, 0, 1) + + # Convert PyTorch tensor to PIL Image for processing + pil_img = TF.to_pil_image(img) + + # Calculate the crop size based on the amplitude + width, height = pil_img.size + crop_size = int(min(width, height) * (1 - amp * zoom_scale)) + crop_size = max(crop_size, 1) + + # Calculate the crop box coordinates (centered crop) + left = (width - crop_size) // 2 + top = (height - crop_size) // 2 + right = (width + crop_size) // 2 + bottom = (height + crop_size) // 2 + + # Crop and resize back to original size + cropped_img = TF.crop(pil_img, top, left, crop_size, crop_size) + resized_img = TF.resize(cropped_img, (height, width)) + + # Convert back to tensor in CxHxW format + tensor_img = TF.to_tensor(resized_img) + + # Convert the tensor back to BxHxWxC format + tensor_img = tensor_img.permute(1, 2, 0) + + # Offset the image based on the amplitude + offset_amp = amp * 10 # Calculate the offset magnitude based on the amplitude + shift_x = min(x_offset * offset_amp, img.shape[1] - 1) # Calculate the shift in x direction + shift_y = min(y_offset * offset_amp, img.shape[0] - 1) # Calculate the shift in y direction + + # Apply the offset to the image tensor + if shift_x != 0: + tensor_img = torch.roll(tensor_img, shifts=int(shift_x), dims=1) + if shift_y != 0: + tensor_img = torch.roll(tensor_img, shifts=int(shift_y), dims=0) + + # Add to the list + transformed_images.append(tensor_img) + + # Stack all transformed images into a batch + transformed_batch = torch.stack(transformed_images) + + return (transformed_batch,) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/batchcrop_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/batchcrop_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..304ff1265770265fe4786e424a05f372daaf1f03 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/batchcrop_nodes.py @@ -0,0 +1,768 @@ +from ..utility.utility import tensor2pil, pil2tensor +from PIL import Image, ImageDraw, ImageFilter +import numpy as np +import torch +from torchvision.transforms import Resize, CenterCrop, InterpolationMode +import math + +#based on nodes from mtb https://github.com/melMass/comfy_mtb + +def bbox_to_region(bbox, target_size=None): + bbox = bbox_check(bbox, target_size) + return (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]) + +def bbox_check(bbox, target_size=None): + if not target_size: + return bbox + + new_bbox = ( + bbox[0], + bbox[1], + min(target_size[0] - bbox[0], bbox[2]), + min(target_size[1] - bbox[1], bbox[3]), + ) + return new_bbox + +class BatchCropFromMask: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "original_images": ("IMAGE",), + "masks": ("MASK",), + "crop_size_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}), + "bbox_smooth_alpha": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "BBOX", + "INT", + "INT", + ) + RETURN_NAMES = ( + "original_images", + "cropped_images", + "bboxes", + "width", + "height", + ) + FUNCTION = "crop" + CATEGORY = "KJNodes/masking" + + def smooth_bbox_size(self, prev_bbox_size, curr_bbox_size, alpha): + if alpha == 0: + return prev_bbox_size + return round(alpha * curr_bbox_size + (1 - alpha) * prev_bbox_size) + + def smooth_center(self, prev_center, curr_center, alpha=0.5): + if alpha == 0: + return prev_center + return ( + round(alpha * curr_center[0] + (1 - alpha) * prev_center[0]), + round(alpha * curr_center[1] + (1 - alpha) * prev_center[1]) + ) + + def crop(self, masks, original_images, crop_size_mult, bbox_smooth_alpha): + + bounding_boxes = [] + cropped_images = [] + + self.max_bbox_width = 0 + self.max_bbox_height = 0 + + # First, calculate the maximum bounding box size across all masks + curr_max_bbox_width = 0 + curr_max_bbox_height = 0 + for mask in masks: + _mask = tensor2pil(mask)[0] + non_zero_indices = np.nonzero(np.array(_mask)) + min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1]) + min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0]) + width = max_x - min_x + height = max_y - min_y + curr_max_bbox_width = max(curr_max_bbox_width, width) + curr_max_bbox_height = max(curr_max_bbox_height, height) + + # Smooth the changes in the bounding box size + self.max_bbox_width = self.smooth_bbox_size(self.max_bbox_width, curr_max_bbox_width, bbox_smooth_alpha) + self.max_bbox_height = self.smooth_bbox_size(self.max_bbox_height, curr_max_bbox_height, bbox_smooth_alpha) + + # Apply the crop size multiplier + self.max_bbox_width = round(self.max_bbox_width * crop_size_mult) + self.max_bbox_height = round(self.max_bbox_height * crop_size_mult) + bbox_aspect_ratio = self.max_bbox_width / self.max_bbox_height + + # Then, for each mask and corresponding image... + for i, (mask, img) in enumerate(zip(masks, original_images)): + _mask = tensor2pil(mask)[0] + non_zero_indices = np.nonzero(np.array(_mask)) + min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1]) + min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0]) + + # Calculate center of bounding box + center_x = np.mean(non_zero_indices[1]) + center_y = np.mean(non_zero_indices[0]) + curr_center = (round(center_x), round(center_y)) + + # If this is the first frame, initialize prev_center with curr_center + if not hasattr(self, 'prev_center'): + self.prev_center = curr_center + + # Smooth the changes in the center coordinates from the second frame onwards + if i > 0: + center = self.smooth_center(self.prev_center, curr_center, bbox_smooth_alpha) + else: + center = curr_center + + # Update prev_center for the next frame + self.prev_center = center + + # Create bounding box using max_bbox_width and max_bbox_height + half_box_width = round(self.max_bbox_width / 2) + half_box_height = round(self.max_bbox_height / 2) + min_x = max(0, center[0] - half_box_width) + max_x = min(img.shape[1], center[0] + half_box_width) + min_y = max(0, center[1] - half_box_height) + max_y = min(img.shape[0], center[1] + half_box_height) + + # Append bounding box coordinates + bounding_boxes.append((min_x, min_y, max_x - min_x, max_y - min_y)) + + # Crop the image from the bounding box + cropped_img = img[min_y:max_y, min_x:max_x, :] + + # Calculate the new dimensions while maintaining the aspect ratio + new_height = min(cropped_img.shape[0], self.max_bbox_height) + new_width = round(new_height * bbox_aspect_ratio) + + # Resize the image + resize_transform = Resize((new_height, new_width)) + resized_img = resize_transform(cropped_img.permute(2, 0, 1)) + + # Perform the center crop to the desired size + crop_transform = CenterCrop((self.max_bbox_height, self.max_bbox_width)) # swap the order here if necessary + cropped_resized_img = crop_transform(resized_img) + + cropped_images.append(cropped_resized_img.permute(1, 2, 0)) + + cropped_out = torch.stack(cropped_images, dim=0) + + return (original_images, cropped_out, bounding_boxes, self.max_bbox_width, self.max_bbox_height, ) + +class BatchUncrop: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "original_images": ("IMAGE",), + "cropped_images": ("IMAGE",), + "bboxes": ("BBOX",), + "border_blending": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "crop_rescale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "border_top": ("BOOLEAN", {"default": True}), + "border_bottom": ("BOOLEAN", {"default": True}), + "border_left": ("BOOLEAN", {"default": True}), + "border_right": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "uncrop" + + CATEGORY = "KJNodes/masking" + + def uncrop(self, original_images, cropped_images, bboxes, border_blending, crop_rescale, border_top, border_bottom, border_left, border_right): + def inset_border(image, border_width, border_color, border_top, border_bottom, border_left, border_right): + draw = ImageDraw.Draw(image) + width, height = image.size + if border_top: + draw.rectangle((0, 0, width, border_width), fill=border_color) + if border_bottom: + draw.rectangle((0, height - border_width, width, height), fill=border_color) + if border_left: + draw.rectangle((0, 0, border_width, height), fill=border_color) + if border_right: + draw.rectangle((width - border_width, 0, width, height), fill=border_color) + return image + + if len(original_images) != len(cropped_images): + raise ValueError(f"The number of original_images ({len(original_images)}) and cropped_images ({len(cropped_images)}) should be the same") + + # Ensure there are enough bboxes, but drop the excess if there are more bboxes than images + if len(bboxes) > len(original_images): + print(f"Warning: Dropping excess bounding boxes. Expected {len(original_images)}, but got {len(bboxes)}") + bboxes = bboxes[:len(original_images)] + elif len(bboxes) < len(original_images): + raise ValueError("There should be at least as many bboxes as there are original and cropped images") + + input_images = tensor2pil(original_images) + crop_imgs = tensor2pil(cropped_images) + + out_images = [] + for i in range(len(input_images)): + img = input_images[i] + crop = crop_imgs[i] + bbox = bboxes[i] + + # uncrop the image based on the bounding box + bb_x, bb_y, bb_width, bb_height = bbox + + paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size) + + # scale factors + scale_x = crop_rescale + scale_y = crop_rescale + + # scaled paste_region + paste_region = (round(paste_region[0]*scale_x), round(paste_region[1]*scale_y), round(paste_region[2]*scale_x), round(paste_region[3]*scale_y)) + + # rescale the crop image to fit the paste_region + crop = crop.resize((round(paste_region[2]-paste_region[0]), round(paste_region[3]-paste_region[1]))) + crop_img = crop.convert("RGB") + + if border_blending > 1.0: + border_blending = 1.0 + elif border_blending < 0.0: + border_blending = 0.0 + + blend_ratio = (max(crop_img.size) / 2) * float(border_blending) + + blend = img.convert("RGBA") + mask = Image.new("L", img.size, 0) + + mask_block = Image.new("L", (paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]), 255) + mask_block = inset_border(mask_block, round(blend_ratio / 2), (0), border_top, border_bottom, border_left, border_right) + + mask.paste(mask_block, paste_region) + blend.paste(crop_img, paste_region) + + mask = mask.filter(ImageFilter.BoxBlur(radius=blend_ratio / 4)) + mask = mask.filter(ImageFilter.GaussianBlur(radius=blend_ratio / 4)) + + blend.putalpha(mask) + img = Image.alpha_composite(img.convert("RGBA"), blend) + out_images.append(img.convert("RGB")) + + return (pil2tensor(out_images),) + +class BatchCropFromMaskAdvanced: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "original_images": ("IMAGE",), + "masks": ("MASK",), + "crop_size_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "bbox_smooth_alpha": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "MASK", + "IMAGE", + "MASK", + "BBOX", + "BBOX", + "INT", + "INT", + ) + RETURN_NAMES = ( + "original_images", + "cropped_images", + "cropped_masks", + "combined_crop_image", + "combined_crop_masks", + "bboxes", + "combined_bounding_box", + "bbox_width", + "bbox_height", + ) + FUNCTION = "crop" + CATEGORY = "KJNodes/masking" + + def smooth_bbox_size(self, prev_bbox_size, curr_bbox_size, alpha): + return round(alpha * curr_bbox_size + (1 - alpha) * prev_bbox_size) + + def smooth_center(self, prev_center, curr_center, alpha=0.5): + return (round(alpha * curr_center[0] + (1 - alpha) * prev_center[0]), + round(alpha * curr_center[1] + (1 - alpha) * prev_center[1])) + + def crop(self, masks, original_images, crop_size_mult, bbox_smooth_alpha): + bounding_boxes = [] + combined_bounding_box = [] + cropped_images = [] + cropped_masks = [] + cropped_masks_out = [] + combined_crop_out = [] + combined_cropped_images = [] + combined_cropped_masks = [] + + def calculate_bbox(mask): + non_zero_indices = np.nonzero(np.array(mask)) + + # handle empty masks + min_x, max_x, min_y, max_y = 0, 0, 0, 0 + if len(non_zero_indices[1]) > 0 and len(non_zero_indices[0]) > 0: + min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1]) + min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0]) + + width = max_x - min_x + height = max_y - min_y + bbox_size = max(width, height) + return min_x, max_x, min_y, max_y, bbox_size + + combined_mask = torch.max(masks, dim=0)[0] + _mask = tensor2pil(combined_mask)[0] + new_min_x, new_max_x, new_min_y, new_max_y, combined_bbox_size = calculate_bbox(_mask) + center_x = (new_min_x + new_max_x) / 2 + center_y = (new_min_y + new_max_y) / 2 + half_box_size = round(combined_bbox_size // 2) + new_min_x = max(0, round(center_x - half_box_size)) + new_max_x = min(original_images[0].shape[1], round(center_x + half_box_size)) + new_min_y = max(0, round(center_y - half_box_size)) + new_max_y = min(original_images[0].shape[0], round(center_y + half_box_size)) + + combined_bounding_box.append((new_min_x, new_min_y, new_max_x - new_min_x, new_max_y - new_min_y)) + + self.max_bbox_size = 0 + + # First, calculate the maximum bounding box size across all masks + curr_max_bbox_size = max(calculate_bbox(tensor2pil(mask)[0])[-1] for mask in masks) + # Smooth the changes in the bounding box size + self.max_bbox_size = self.smooth_bbox_size(self.max_bbox_size, curr_max_bbox_size, bbox_smooth_alpha) + # Apply the crop size multiplier + self.max_bbox_size = round(self.max_bbox_size * crop_size_mult) + # Make sure max_bbox_size is divisible by 16, if not, round it upwards so it is + self.max_bbox_size = math.ceil(self.max_bbox_size / 16) * 16 + + if self.max_bbox_size > original_images[0].shape[0] or self.max_bbox_size > original_images[0].shape[1]: + # max_bbox_size can only be as big as our input's width or height, and it has to be even + self.max_bbox_size = math.floor(min(original_images[0].shape[0], original_images[0].shape[1]) / 2) * 2 + + # Then, for each mask and corresponding image... + for i, (mask, img) in enumerate(zip(masks, original_images)): + _mask = tensor2pil(mask)[0] + non_zero_indices = np.nonzero(np.array(_mask)) + + # check for empty masks + if len(non_zero_indices[0]) > 0 and len(non_zero_indices[1]) > 0: + min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1]) + min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0]) + + # Calculate center of bounding box + center_x = np.mean(non_zero_indices[1]) + center_y = np.mean(non_zero_indices[0]) + curr_center = (round(center_x), round(center_y)) + + # If this is the first frame, initialize prev_center with curr_center + if not hasattr(self, 'prev_center'): + self.prev_center = curr_center + + # Smooth the changes in the center coordinates from the second frame onwards + if i > 0: + center = self.smooth_center(self.prev_center, curr_center, bbox_smooth_alpha) + else: + center = curr_center + + # Update prev_center for the next frame + self.prev_center = center + + # Create bounding box using max_bbox_size + half_box_size = self.max_bbox_size // 2 + min_x = max(0, center[0] - half_box_size) + max_x = min(img.shape[1], center[0] + half_box_size) + min_y = max(0, center[1] - half_box_size) + max_y = min(img.shape[0], center[1] + half_box_size) + + # Append bounding box coordinates + bounding_boxes.append((min_x, min_y, max_x - min_x, max_y - min_y)) + + # Crop the image from the bounding box + cropped_img = img[min_y:max_y, min_x:max_x, :] + cropped_mask = mask[min_y:max_y, min_x:max_x] + + # Resize the cropped image to a fixed size + new_size = max(cropped_img.shape[0], cropped_img.shape[1]) + resize_transform = Resize(new_size, interpolation=InterpolationMode.NEAREST, max_size=max(img.shape[0], img.shape[1])) + resized_mask = resize_transform(cropped_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0) + resized_img = resize_transform(cropped_img.permute(2, 0, 1)) + # Perform the center crop to the desired size + # Constrain the crop to the smaller of our bbox or our image so we don't expand past the image dimensions. + crop_transform = CenterCrop((min(self.max_bbox_size, resized_img.shape[1]), min(self.max_bbox_size, resized_img.shape[2]))) + + cropped_resized_img = crop_transform(resized_img) + cropped_images.append(cropped_resized_img.permute(1, 2, 0)) + + cropped_resized_mask = crop_transform(resized_mask) + cropped_masks.append(cropped_resized_mask) + + combined_cropped_img = original_images[i][new_min_y:new_max_y, new_min_x:new_max_x, :] + combined_cropped_images.append(combined_cropped_img) + + combined_cropped_mask = masks[i][new_min_y:new_max_y, new_min_x:new_max_x] + combined_cropped_masks.append(combined_cropped_mask) + else: + bounding_boxes.append((0, 0, img.shape[1], img.shape[0])) + cropped_images.append(img) + cropped_masks.append(mask) + combined_cropped_images.append(img) + combined_cropped_masks.append(mask) + + cropped_out = torch.stack(cropped_images, dim=0) + combined_crop_out = torch.stack(combined_cropped_images, dim=0) + cropped_masks_out = torch.stack(cropped_masks, dim=0) + combined_crop_mask_out = torch.stack(combined_cropped_masks, dim=0) + + return (original_images, cropped_out, cropped_masks_out, combined_crop_out, combined_crop_mask_out, bounding_boxes, combined_bounding_box, self.max_bbox_size, self.max_bbox_size) + +class FilterZeroMasksAndCorrespondingImages: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "masks": ("MASK",), + }, + "optional": { + "original_images": ("IMAGE",), + }, + } + + RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "INDEXES",) + RETURN_NAMES = ("non_zero_masks_out", "non_zero_mask_images_out", "zero_mask_images_out", "zero_mask_images_out_indexes",) + FUNCTION = "filter" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Filter out all the empty (i.e. all zero) mask in masks +Also filter out all the corresponding images in original_images by indexes if provide + +original_images (optional): If provided, need have same length as masks. +""" + + def filter(self, masks, original_images=None): + non_zero_masks = [] + non_zero_mask_images = [] + zero_mask_images = [] + zero_mask_images_indexes = [] + + masks_num = len(masks) + also_process_images = False + if original_images is not None: + imgs_num = len(original_images) + if len(original_images) == masks_num: + also_process_images = True + else: + print(f"[WARNING] ignore input: original_images, due to number of original_images ({imgs_num}) is not equal to number of masks ({masks_num})") + + for i in range(masks_num): + non_zero_num = np.count_nonzero(np.array(masks[i])) + if non_zero_num > 0: + non_zero_masks.append(masks[i]) + if also_process_images: + non_zero_mask_images.append(original_images[i]) + else: + zero_mask_images.append(original_images[i]) + zero_mask_images_indexes.append(i) + + non_zero_masks_out = torch.stack(non_zero_masks, dim=0) + non_zero_mask_images_out = zero_mask_images_out = zero_mask_images_out_indexes = None + + if also_process_images: + non_zero_mask_images_out = torch.stack(non_zero_mask_images, dim=0) + if len(zero_mask_images) > 0: + zero_mask_images_out = torch.stack(zero_mask_images, dim=0) + zero_mask_images_out_indexes = zero_mask_images_indexes + + return (non_zero_masks_out, non_zero_mask_images_out, zero_mask_images_out, zero_mask_images_out_indexes) + +class InsertImageBatchByIndexes: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "images_to_insert": ("IMAGE",), + "insert_indexes": ("INDEXES",), + }, + } + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("images_after_insert", ) + FUNCTION = "insert" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +This node is designed to be use with node FilterZeroMasksAndCorrespondingImages +It inserts the images_to_insert into images according to insert_indexes + +Returns: + images_after_insert: updated original images with origonal sequence order +""" + + def insert(self, images, images_to_insert, insert_indexes): + images_after_insert = images + + if images_to_insert is not None and insert_indexes is not None: + images_to_insert_num = len(images_to_insert) + insert_indexes_num = len(insert_indexes) + if images_to_insert_num == insert_indexes_num: + images_after_insert = [] + + i_images = 0 + for i in range(len(images) + images_to_insert_num): + if i in insert_indexes: + images_after_insert.append(images_to_insert[insert_indexes.index(i)]) + else: + images_after_insert.append(images[i_images]) + i_images += 1 + + images_after_insert = torch.stack(images_after_insert, dim=0) + + else: + print(f"[WARNING] skip this node, due to number of images_to_insert ({images_to_insert_num}) is not equal to number of insert_indexes ({insert_indexes_num})") + + + return (images_after_insert, ) + +class BatchUncropAdvanced: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "original_images": ("IMAGE",), + "cropped_images": ("IMAGE",), + "cropped_masks": ("MASK",), + "combined_crop_mask": ("MASK",), + "bboxes": ("BBOX",), + "border_blending": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "crop_rescale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "use_combined_mask": ("BOOLEAN", {"default": False}), + "use_square_mask": ("BOOLEAN", {"default": True}), + }, + "optional": { + "combined_bounding_box": ("BBOX", {"default": None}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "uncrop" + CATEGORY = "KJNodes/masking" + + + def uncrop(self, original_images, cropped_images, cropped_masks, combined_crop_mask, bboxes, border_blending, crop_rescale, use_combined_mask, use_square_mask, combined_bounding_box = None): + + def inset_border(image, border_width=20, border_color=(0)): + width, height = image.size + bordered_image = Image.new(image.mode, (width, height), border_color) + bordered_image.paste(image, (0, 0)) + draw = ImageDraw.Draw(bordered_image) + draw.rectangle((0, 0, width - 1, height - 1), outline=border_color, width=border_width) + return bordered_image + + if len(original_images) != len(cropped_images): + raise ValueError(f"The number of original_images ({len(original_images)}) and cropped_images ({len(cropped_images)}) should be the same") + + # Ensure there are enough bboxes, but drop the excess if there are more bboxes than images + if len(bboxes) > len(original_images): + print(f"Warning: Dropping excess bounding boxes. Expected {len(original_images)}, but got {len(bboxes)}") + bboxes = bboxes[:len(original_images)] + elif len(bboxes) < len(original_images): + raise ValueError("There should be at least as many bboxes as there are original and cropped images") + + crop_imgs = tensor2pil(cropped_images) + input_images = tensor2pil(original_images) + out_images = [] + + for i in range(len(input_images)): + img = input_images[i] + crop = crop_imgs[i] + bbox = bboxes[i] + + if use_combined_mask: + bb_x, bb_y, bb_width, bb_height = combined_bounding_box[0] + paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size) + mask = combined_crop_mask[i] + else: + bb_x, bb_y, bb_width, bb_height = bbox + paste_region = bbox_to_region((bb_x, bb_y, bb_width, bb_height), img.size) + mask = cropped_masks[i] + + # scale paste_region + scale_x = scale_y = crop_rescale + paste_region = (round(paste_region[0]*scale_x), round(paste_region[1]*scale_y), round(paste_region[2]*scale_x), round(paste_region[3]*scale_y)) + + # rescale the crop image to fit the paste_region + crop = crop.resize((round(paste_region[2]-paste_region[0]), round(paste_region[3]-paste_region[1]))) + crop_img = crop.convert("RGB") + + #border blending + if border_blending > 1.0: + border_blending = 1.0 + elif border_blending < 0.0: + border_blending = 0.0 + + blend_ratio = (max(crop_img.size) / 2) * float(border_blending) + blend = img.convert("RGBA") + + if use_square_mask: + mask = Image.new("L", img.size, 0) + mask_block = Image.new("L", (paste_region[2]-paste_region[0], paste_region[3]-paste_region[1]), 255) + mask_block = inset_border(mask_block, round(blend_ratio / 2), (0)) + mask.paste(mask_block, paste_region) + else: + original_mask = tensor2pil(mask)[0] + original_mask = original_mask.resize((paste_region[2]-paste_region[0], paste_region[3]-paste_region[1])) + mask = Image.new("L", img.size, 0) + mask.paste(original_mask, paste_region) + + mask = mask.filter(ImageFilter.BoxBlur(radius=blend_ratio / 4)) + mask = mask.filter(ImageFilter.GaussianBlur(radius=blend_ratio / 4)) + + blend.paste(crop_img, paste_region) + blend.putalpha(mask) + + img = Image.alpha_composite(img.convert("RGBA"), blend) + out_images.append(img.convert("RGB")) + + return (pil2tensor(out_images),) + +class SplitBboxes: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "bboxes": ("BBOX",), + "index": ("INT", {"default": 0,"min": 0, "max": 99999999, "step": 1}), + }, + } + + RETURN_TYPES = ("BBOX","BBOX",) + RETURN_NAMES = ("bboxes_a","bboxes_b",) + FUNCTION = "splitbbox" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Splits the specified bbox list at the given index into two lists. +""" + + def splitbbox(self, bboxes, index): + bboxes_a = bboxes[:index] # Sub-list from the start of bboxes up to (but not including) the index + bboxes_b = bboxes[index:] # Sub-list from the index to the end of bboxes + + return (bboxes_a, bboxes_b,) + +class BboxToInt: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "bboxes": ("BBOX",), + "index": ("INT", {"default": 0,"min": 0, "max": 99999999, "step": 1}), + }, + } + + RETURN_TYPES = ("INT","INT","INT","INT","INT","INT",) + RETURN_NAMES = ("x_min","y_min","width","height", "center_x","center_y",) + FUNCTION = "bboxtoint" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Returns selected index from bounding box list as integers. +""" + def bboxtoint(self, bboxes, index): + x_min, y_min, width, height = bboxes[index] + center_x = int(x_min + width / 2) + center_y = int(y_min + height / 2) + + return (x_min, y_min, width, height, center_x, center_y,) + +class BboxVisualize: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "bboxes": ("BBOX",), + "line_width": ("INT", {"default": 1,"min": 1, "max": 10, "step": 1}), + "bbox_format": (["xywh", "xyxy"], {"default": "xywh"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "visualizebbox" + DESCRIPTION = """ +Visualizes the specified bbox on the image. +""" + + CATEGORY = "KJNodes/masking" + + def visualizebbox(self, bboxes, images, line_width, bbox_format): + image_list = [] + for image, bbox in zip(images, bboxes): + # Ensure bbox is a sequence of 4 values + if isinstance(bbox, (list, tuple, np.ndarray)) and len(bbox) == 4: + if bbox_format == "xywh": + x_min, y_min, width, height = bbox + elif bbox_format == "xyxy": + x_min, y_min, x_max, y_max = bbox + width = x_max - x_min + height = y_max - y_min + else: + raise ValueError(f"Unknown bbox_format: {bbox_format}") + else: + print("Invalid bbox:", bbox) + continue + + # Ensure bbox coordinates are integers + x_min = int(x_min) + y_min = int(y_min) + width = int(width) + height = int(height) + + # Permute the image dimensions + image = image.permute(2, 0, 1) + + # Clone the image to draw bounding boxes + img_with_bbox = image.clone() + + # Define the color for the bbox, e.g., red + color = torch.tensor([1, 0, 0], dtype=torch.float32) + + # Ensure color tensor matches the image channels + if color.shape[0] != img_with_bbox.shape[0]: + color = color.unsqueeze(1).expand(-1, line_width) + + # Draw lines for each side of the bbox with the specified line width + for lw in range(line_width): + # Top horizontal line + if y_min + lw < img_with_bbox.shape[1]: + img_with_bbox[:, y_min + lw, x_min:x_min + width] = color[:, None] + + # Bottom horizontal line + if y_min + height - lw < img_with_bbox.shape[1]: + img_with_bbox[:, y_min + height - lw, x_min:x_min + width] = color[:, None] + + # Left vertical line + if x_min + lw < img_with_bbox.shape[2]: + img_with_bbox[:, y_min:y_min + height, x_min + lw] = color[:, None] + + # Right vertical line + if x_min + width - lw < img_with_bbox.shape[2]: + img_with_bbox[:, y_min:y_min + height, x_min + width - lw] = color[:, None] + + # Permute the image dimensions back + img_with_bbox = img_with_bbox.permute(1, 2, 0).unsqueeze(0) + image_list.append(img_with_bbox) + + return (torch.cat(image_list, dim=0),) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/curve_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/curve_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..35dc3558790ca2f5e7860c04827e47a012bdebc7 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/curve_nodes.py @@ -0,0 +1,1642 @@ +import torch +from torchvision import transforms +import json +from PIL import Image, ImageDraw, ImageFont, ImageFilter +import numpy as np +from ..utility.utility import pil2tensor, tensor2pil +import folder_paths +import base64 +from io import BytesIO + + +def parse_color(color): + if isinstance(color, str) and ',' in color: + return tuple(int(c.strip()) for c in color.split(',')) + return color + +def parse_json_tracks(tracks): + tracks_data = [] + try: + # If tracks is a string, try to parse it as JSON + if isinstance(tracks, str): + parsed = json.loads(tracks.replace("'", '"')) + tracks_data.extend(parsed) + else: + # If tracks is a list of strings, parse each one + for track_str in tracks: + parsed = json.loads(track_str.replace("'", '"')) + tracks_data.append(parsed) + + # Check if we have a single track (dict with x,y) or a list of tracks + if tracks_data and isinstance(tracks_data[0], dict) and 'x' in tracks_data[0]: + # Single track detected, wrap it in a list + tracks_data = [tracks_data] + elif tracks_data and isinstance(tracks_data[0], list) and tracks_data[0] and isinstance(tracks_data[0][0], dict) and 'x' in tracks_data[0][0]: + # Already a list of tracks, nothing to do + pass + else: + # Unexpected format + print(f"Warning: Unexpected track format: {type(tracks_data[0])}") + + except json.JSONDecodeError as e: + print(f"Error parsing tracks JSON: {e}") + tracks_data = [] + + return tracks_data + +def plot_coordinates_to_tensor(coordinates, height, width, bbox_height, bbox_width, size_multiplier, prompt): + import matplotlib + matplotlib.use('Agg') + from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas + text_color = '#999999' + bg_color = '#353535' + matplotlib.pyplot.rcParams['text.color'] = text_color + fig, ax = matplotlib.pyplot.subplots(figsize=(width/100, height/100), dpi=100) + fig.patch.set_facecolor(bg_color) + ax.set_facecolor(bg_color) + ax.grid(color=text_color, linestyle='-', linewidth=0.5) + ax.set_xlabel('x', color=text_color) + ax.set_ylabel('y', color=text_color) + for text in ax.get_xticklabels() + ax.get_yticklabels(): + text.set_color(text_color) + ax.set_title('position for: ' + prompt) + ax.set_xlabel('X Coordinate') + ax.set_ylabel('Y Coordinate') + #ax.legend().remove() + ax.set_xlim(0, width) # Set the x-axis to match the input latent width + ax.set_ylim(height, 0) # Set the y-axis to match the input latent height, with (0,0) at top-left + # Adjust the margins of the subplot + matplotlib.pyplot.subplots_adjust(left=0.08, right=0.95, bottom=0.05, top=0.95, wspace=0.2, hspace=0.2) + + cmap = matplotlib.pyplot.get_cmap('rainbow') + image_batch = [] + canvas = FigureCanvas(fig) + width, height = fig.get_size_inches() * fig.get_dpi() + # Draw a box at each coordinate + for i, ((x, y), size) in enumerate(zip(coordinates, size_multiplier)): + color_index = i / (len(coordinates) - 1) + color = cmap(color_index) + draw_height = bbox_height * size + draw_width = bbox_width * size + rect = matplotlib.patches.Rectangle((x - draw_width/2, y - draw_height/2), draw_width, draw_height, + linewidth=1, edgecolor=color, facecolor='none', alpha=0.5) + ax.add_patch(rect) + + # Check if there is a next coordinate to draw an arrow to + if i < len(coordinates) - 1: + x1, y1 = coordinates[i] + x2, y2 = coordinates[i + 1] + ax.annotate("", xy=(x2, y2), xytext=(x1, y1), + arrowprops=dict(arrowstyle="->", + linestyle="-", + lw=1, + color=color, + mutation_scale=20)) + canvas.draw() + try: + image_np = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3).copy() + except AttributeError: + image_np = np.frombuffer(canvas.tostring_argb(), dtype='uint8').reshape(int(height), int(width), 4) + image_np = image_np[:, :, [1, 2, 3]] # Convert ARGB to RGB + image_tensor = torch.from_numpy(image_np).float() / 255.0 + image_tensor = image_tensor.unsqueeze(0) + image_batch.append(image_tensor) + + matplotlib.pyplot.close(fig) + image_batch_tensor = torch.cat(image_batch, dim=0) + + return image_batch_tensor + +class PlotCoordinates: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "coordinates": ("STRING", {"forceInput": True}), + "text": ("STRING", {"default": 'title', "multiline": False}), + "width": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}), + "height": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}), + "bbox_width": ("INT", {"default": 128, "min": 8, "max": 4096, "step": 8}), + "bbox_height": ("INT", {"default": 128, "min": 8, "max": 4096, "step": 8}), + }, + "optional": {"size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True})}, + } + RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "INT",) + RETURN_NAMES = ("images", "width", "height", "bbox_width", "bbox_height",) + FUNCTION = "append" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = """ +Plots coordinates to sequence of images using Matplotlib. + +""" + + def append(self, coordinates, text, width, height, bbox_width, bbox_height, size_multiplier=[1.0]): + coordinates = json.loads(coordinates.replace("'", '"')) + coordinates = [(coord['x'], coord['y']) for coord in coordinates] + batch_size = len(coordinates) + if not size_multiplier or len(size_multiplier) != batch_size: + size_multiplier = [0] * batch_size + else: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + + plot_image_tensor = plot_coordinates_to_tensor(coordinates, height, width, bbox_height, bbox_width, size_multiplier, text) + + return (plot_image_tensor, width, height, bbox_width, bbox_height) + +class SplineEditor: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "points_store": ("STRING", {"multiline": False, "advanced": True}), + "coordinates": ("STRING", {"multiline": False, "advanced": True}), + "mask_width": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}), + "mask_height": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}), + "points_to_sample": ("INT", {"default": 16, "min": 2, "max": 1000, "step": 1}), + "sampling_method": ( + [ + 'path', + 'time', + 'controlpoints', + 'speed' + ], + { + "default": 'time' + }), + "interpolation": ( + [ + 'cardinal', + 'monotone', + 'basis', + 'linear', + 'step-before', + 'step-after', + 'polar', + 'polar-reverse', + 'bezier', + ], + { + "default": 'cardinal' + }), + "tension": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "repeat_output": ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1}), + "float_output_type": ( + [ + 'list', + 'pandas series', + 'tensor', + ], + { + "default": 'list' + }), + }, + "optional": { + "min_value": ("FLOAT", {"default": 0.0, "min": -10000.0, "max": 10000.0, "step": 0.01}), + "max_value": ("FLOAT", {"default": 1.0, "min": -10000.0, "max": 10000.0, "step": 0.01}), + "bg_image": ("IMAGE", ), + } + } + + RETURN_TYPES = ("MASK", "STRING", "FLOAT", "INT", "STRING",) + RETURN_NAMES = ("mask", "coord_str", "float", "count", "normalized_str",) + FUNCTION = "splinedata" + CATEGORY = "KJNodes/weights" + DESCRIPTION = """ +# WORK IN PROGRESS +Do not count on this as part of your workflow yet, +probably contains lots of bugs and stability is not +guaranteed!! + +## Graphical editor to create values for various +## schedules and/or mask batches. + +**Shift + click** to add control point at end. +**Ctrl + click** to add control point (subdivide) between two points. +**Right click on a point** to delete it. +Note that you can't delete from start/end. + +Right click on canvas for context menu: +NEW!: +- Add new spline + - Creates a new spline on same canvas, currently these paths are only outputed + as coordinates. +- Add single point + - Creates a single point that only returns it's current position coords +- Delete spline + - Deletes the currently selected spline, you can select a spline by clicking on + it's path, or cycle through them with the 'Next spline' -option. + +These are purely visual options, doesn't affect the output: + - Toggle handles visibility + - Display sample points: display the points to be returned. + +**points_to_sample** value sets the number of samples +returned from the **drawn spline itself**, this is independent from the +actual control points, so the interpolation type matters. +sampling_method: + - time: samples along the time axis, used for schedules + - path: samples along the path itself, useful for coordinates + - controlpoints: samples only the control points themselves + +output types: + - mask batch + example compatible nodes: anything that takes masks + - list of floats + example compatible nodes: IPAdapter weights + - pandas series + example compatible nodes: anything that takes Fizz' + nodes Batch Value Schedule + - torch tensor + example compatible nodes: unknown +""" + + def splinedata(self, mask_width, mask_height, coordinates, float_output_type, interpolation, + points_to_sample, sampling_method, points_store, tension, repeat_output, + min_value=0.0, max_value=1.0, bg_image=None): + + coordinates = json.loads(coordinates) + + # Handle nested list structure if present + all_normalized = [] + all_normalized_y_values = [] + + # Check if we have a nested list structure + if isinstance(coordinates, list) and len(coordinates) > 0 and isinstance(coordinates[0], list): + # Process each list of coordinates in the nested structure + coordinate_sets = coordinates + else: + # If not nested, treat as a single list of coordinates + coordinate_sets = [coordinates] + + first_spline = coordinate_sets[0] if coordinate_sets else [] + + # Process each set of coordinates + for coord_set in coordinate_sets: + normalized = [] + normalized_y_values = [] + + for coord in coord_set: + coord['x'] = int(round(coord['x'])) + coord['y'] = int(round(coord['y'])) + norm_x = (1.0 - (coord['x'] / mask_height) - 0.0) * (max_value - min_value) + min_value + norm_y = (1.0 - (coord['y'] / mask_height) - 0.0) * (max_value - min_value) + min_value + normalized_y_values.append(norm_y) + normalized.append({'x':norm_x, 'y':norm_y}) + + all_normalized.extend(normalized) + all_normalized_y_values.extend(normalized_y_values) + + # Use the combined normalized values for output + if float_output_type == 'list': + out_floats = all_normalized_y_values * repeat_output + elif float_output_type == 'pandas series': + try: + import pandas as pd + except ImportError as e: + raise ImportError("MaskOrImageToWeight: pandas is not installed. Please install pandas to use this output_type") from e + out_floats = pd.Series(all_normalized_y_values * repeat_output), + elif float_output_type == 'tensor': + out_floats = torch.tensor(all_normalized_y_values * repeat_output, dtype=torch.float32) + + color_map = lambda y: torch.full((mask_height, mask_width, 3), y, dtype=torch.float32) + y_values_for_masks = normalized_y_values if normalized_y_values else [min_value] + mask_tensors = [color_map(y) for y in y_values_for_masks] + masks_out = torch.stack(mask_tensors) + masks_out = masks_out.repeat(repeat_output, 1, 1, 1) + masks_out = masks_out.mean(dim=-1) + + single_spline_count = len(first_spline) + + if bg_image is None: + return (masks_out, json.dumps(coordinates if len(coordinates) > 1 else coordinates[0]), out_floats, single_spline_count, json.dumps(all_normalized)) + else: + transform = transforms.ToPILImage() + image = transform(bg_image[0].permute(2, 0, 1)) + buffered = BytesIO() + image.save(buffered, format="JPEG", quality=75) + + # Encode the image bytes to a Base64 string + img_bytes = buffered.getvalue() + img_base64 = base64.b64encode(img_bytes).decode('utf-8') + + return { + "ui": {"bg_image": [img_base64]}, + "result": (masks_out, json.dumps(coordinates if len(coordinates) > 1 else coordinates[0]), out_floats, single_spline_count, json.dumps(all_normalized)) + } + + +class CreateShapeMaskOnPath: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createshapemask" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Creates a mask or batch of masks with the specified shape. +Locations are center locations. +""" + DEPRECATED = True + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "shape": ( + [ 'circle', + 'square', + 'triangle', + ], + { + "default": 'circle' + }), + "coordinates": ("STRING", {"forceInput": True}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "shape_width": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + "shape_height": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + }, + "optional": { + "size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True}), + } + } + + def createshapemask(self, coordinates, frame_width, frame_height, shape_width, shape_height, shape, size_multiplier=[1.0]): + # Define the number of images in the batch + coordinates = coordinates.replace("'", '"') + coordinates = json.loads(coordinates) + + batch_size = len(coordinates) + out = [] + color = "white" + if not size_multiplier or len(size_multiplier) != batch_size: + size_multiplier = [0] * batch_size + else: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + for i, coord in enumerate(coordinates): + image = Image.new("RGB", (frame_width, frame_height), "black") + draw = ImageDraw.Draw(image) + + # Calculate the size for this frame and ensure it's not less than 0 + current_width = max(0, shape_width + i * size_multiplier[i]) + current_height = max(0, shape_height + i * size_multiplier[i]) + + location_x = coord['x'] + location_y = coord['y'] + + if shape == 'circle' or shape == 'square': + # Define the bounding box for the shape + left_up_point = (location_x - current_width // 2, location_y - current_height // 2) + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) + two_points = [left_up_point, right_down_point] + + if shape == 'circle': + draw.ellipse(two_points, fill=color) + elif shape == 'square': + draw.rectangle(two_points, fill=color) + + elif shape == 'triangle': + # Define the points for the triangle + left_up_point = (location_x - current_width // 2, location_y + current_height // 2) # bottom left + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) # bottom right + top_point = (location_x, location_y - current_height // 2) # top point + draw.polygon([top_point, left_up_point, right_down_point], fill=color) + + image = pil2tensor(image) + mask = image[:, :, :, 0] + out.append(mask) + outstack = torch.cat(out, dim=0) + return (outstack, 1.0 - outstack,) + + + +class CreateShapeImageOnPath: + + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image","mask", ) + FUNCTION = "createshapemask" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates an image or batch of images with the specified shape. +Locations are center locations. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "shape": ( + [ 'circle', + 'square', + 'triangle', + ], + { + "default": 'circle' + }), + "coordinates": ("STRING", {"forceInput": True}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "shape_width": ("INT", {"default": 128,"min": 2, "max": 4096, "step": 1}), + "shape_height": ("INT", {"default": 128,"min": 2, "max": 4096, "step": 1}), + "shape_color": ("STRING", {"default": 'white'}), + "bg_color": ("STRING", {"default": 'black'}), + "blur_radius": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100, "step": 0.1}), + "intensity": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + }, + "optional": { + "size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True}), + "trailing": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "border_width": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), + "border_color": ("STRING", {"default": 'black'}), + } + } + + def createshapemask(self, coordinates, frame_width, frame_height, shape_width, shape_height, shape_color, + bg_color, blur_radius, shape, intensity, size_multiplier=[1.0], trailing=1.0, border_width=0, border_color='black'): + + shape_color = parse_color(shape_color) + border_color = parse_color(border_color) + bg_color = parse_color(bg_color) + coords_list = parse_json_tracks(coordinates) + + batch_size = len(coords_list[0]) + images_list = [] + masks_list = [] + + if not size_multiplier or len(size_multiplier) != batch_size: + size_multiplier = [1] * batch_size + else: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + + previous_output = None + + for i in range(batch_size): + image = Image.new("RGB", (frame_width, frame_height), bg_color) + draw = ImageDraw.Draw(image) + + # Calculate the size for this frame and ensure it's not less than 0 + current_width = shape_width * size_multiplier[i] + current_height = shape_height * size_multiplier[i] + + for coords in coords_list: + location_x = coords[i]['x'] + location_y = coords[i]['y'] + + if shape == 'circle' or shape == 'square': + # Define the bounding box for the shape + left_up_point = (location_x - current_width // 2, location_y - current_height // 2) + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) + two_points = [left_up_point, right_down_point] + + if shape == 'circle': + if border_width > 0: + draw.ellipse(two_points, fill=shape_color, outline=border_color, width=border_width) + else: + draw.ellipse(two_points, fill=shape_color) + elif shape == 'square': + if border_width > 0: + draw.rectangle(two_points, fill=shape_color, outline=border_color, width=border_width) + else: + draw.rectangle(two_points, fill=shape_color) + + elif shape == 'triangle': + # Define the points for the triangle + left_up_point = (location_x - current_width // 2, location_y + current_height // 2) # bottom left + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) # bottom right + top_point = (location_x, location_y - current_height // 2) # top point + + if border_width > 0: + draw.polygon([top_point, left_up_point, right_down_point], fill=shape_color, outline=border_color, width=border_width) + else: + draw.polygon([top_point, left_up_point, right_down_point], fill=shape_color) + + if blur_radius != 0: + image = image.filter(ImageFilter.GaussianBlur(blur_radius)) + # Blend the current image with the accumulated image + + image = pil2tensor(image) + if trailing != 1.0 and previous_output is not None: + # Add the decayed previous output to the current frame + image += trailing * previous_output + image = image / image.max() + previous_output = image + image = image * intensity + mask = image[:, :, :, 0] + masks_list.append(mask) + images_list.append(image) + out_images = torch.cat(images_list, dim=0).cpu().float() + out_masks = torch.cat(masks_list, dim=0) + return (out_images, out_masks) + +class CreateTextOnPath: + + RETURN_TYPES = ("IMAGE", "MASK", "MASK",) + RETURN_NAMES = ("image", "mask", "mask_inverted",) + FUNCTION = "createtextmask" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Creates a mask or batch of masks with the specified text. +Locations are center locations. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "coordinates": ("STRING", {"forceInput": True}), + "text": ("STRING", {"default": 'text', "multiline": True}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "font": (folder_paths.get_filename_list("kjnodes_fonts"), ), + "font_size": ("INT", {"default": 42}), + "alignment": ( + [ 'left', + 'center', + 'right' + ], + {"default": 'center'} + ), + "text_color": ("STRING", {"default": 'white'}), + }, + "optional": { + "size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True}), + } + } + + def createtextmask(self, coordinates, frame_width, frame_height, font, font_size, text, text_color, alignment, size_multiplier=[1.0]): + coordinates = coordinates.replace("'", '"') + coordinates = json.loads(coordinates) + + batch_size = len(coordinates) + mask_list = [] + image_list = [] + color = parse_color(text_color) + font_path = folder_paths.get_full_path("kjnodes_fonts", font) + + if len(size_multiplier) != batch_size: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + + for i, coord in enumerate(coordinates): + image = Image.new("RGB", (frame_width, frame_height), "black") + draw = ImageDraw.Draw(image) + lines = text.split('\n') # Split the text into lines + # Apply the size multiplier to the font size for this iteration + current_font_size = int(font_size * size_multiplier[i]) + current_font = ImageFont.truetype(font_path, current_font_size) + line_heights = [current_font.getbbox(line)[3] for line in lines] # List of line heights + total_text_height = sum(line_heights) # Total height of text block + + # Calculate the starting Y position to center the block of text + start_y = coord['y'] - total_text_height // 2 + for j, line in enumerate(lines): + text_width, text_height = current_font.getbbox(line)[2], line_heights[j] + if alignment == 'left': + location_x = coord['x'] + elif alignment == 'center': + location_x = int(coord['x'] - text_width // 2) + elif alignment == 'right': + location_x = int(coord['x'] - text_width) + + location_y = int(start_y + sum(line_heights[:j])) + text_position = (location_x, location_y) + # Draw the text + try: + draw.text(text_position, line, fill=color, font=current_font, features=['-liga']) + except Exception: + draw.text(text_position, line, fill=color, font=current_font) + + image = pil2tensor(image) + non_black_pixels = (image > 0).any(dim=-1) + mask = non_black_pixels.to(image.dtype) + mask_list.append(mask) + image_list.append(image) + + out_images = torch.cat(image_list, dim=0).cpu().float() + out_masks = torch.cat(mask_list, dim=0) + return (out_images, out_masks, 1.0 - out_masks,) + +class CreateGradientFromCoords: + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("image", ) + FUNCTION = "generate" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates a gradient image from coordinates. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "coordinates": ("STRING", {"forceInput": True}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "start_color": ("STRING", {"default": 'white'}), + "end_color": ("STRING", {"default": 'black'}), + "multiplier": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01}), + }, + } + + def generate(self, coordinates, frame_width, frame_height, start_color, end_color, multiplier): + # Parse the coordinates + coordinates = json.loads(coordinates.replace("'", '"')) + + # Create an image + image = Image.new("RGB", (frame_width, frame_height)) + draw = ImageDraw.Draw(image) + + # Extract start and end points for the gradient + start_coord = coordinates[0] + end_coord = coordinates[1] + + start_color = parse_color(start_color) + end_color = parse_color(end_color) + + # Calculate the gradient direction (vector) + gradient_direction = (end_coord['x'] - start_coord['x'], end_coord['y'] - start_coord['y']) + gradient_length = (gradient_direction[0] ** 2 + gradient_direction[1] ** 2) ** 0.5 + + # Iterate over each pixel in the image + for y in range(frame_height): + for x in range(frame_width): + # Calculate the projection of the point on the gradient line + point_vector = (x - start_coord['x'], y - start_coord['y']) + projection = (point_vector[0] * gradient_direction[0] + point_vector[1] * gradient_direction[1]) / gradient_length + projection = max(min(projection, gradient_length), 0) # Clamp the projection value + + # Calculate the blend factor for the current pixel + blend = projection * multiplier / gradient_length + + # Determine the color of the current pixel + color = ( + int(start_color[0] + (end_color[0] - start_color[0]) * blend), + int(start_color[1] + (end_color[1] - start_color[1]) * blend), + int(start_color[2] + (end_color[2] - start_color[2]) * blend) + ) + + # Set the pixel color + draw.point((x, y), fill=color) + + # Convert the PIL image to a tensor (assuming such a function exists in your context) + image_tensor = pil2tensor(image) + + return (image_tensor,) + +class GradientToFloat: + + RETURN_TYPES = ("FLOAT", "FLOAT",) + RETURN_NAMES = ("float_x", "float_y", ) + FUNCTION = "sample" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Calculates list of floats from image. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "steps": ("INT", {"default": 10, "min": 2, "max": 10000, "step": 1}), + }, + } + + def sample(self, image, steps): + # Assuming image is a tensor with shape [B, H, W, C] + B, H, W, C = image.shape + + # Sample along the width axis (W) + w_intervals = torch.linspace(0, W - 1, steps=steps, dtype=torch.int64) + # Assuming we're sampling from the first batch and the first channel + w_sampled = image[0, :, w_intervals, 0] + + # Sample along the height axis (H) + h_intervals = torch.linspace(0, H - 1, steps=steps, dtype=torch.int64) + # Assuming we're sampling from the first batch and the first channel + h_sampled = image[0, h_intervals, :, 0] + + # Taking the mean across the height for width sampling, and across the width for height sampling + w_values = w_sampled.mean(dim=0).tolist() + h_values = h_sampled.mean(dim=1).tolist() + + return (w_values, h_values) + +class MaskOrImageToWeight: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "output_type": ( + [ + 'list', + 'pandas series', + 'tensor', + 'string' + ], + { + "default": 'list' + }), + }, + "optional": { + "images": ("IMAGE",), + "masks": ("MASK",), + }, + + } + RETURN_TYPES = ("FLOAT", "STRING",) + FUNCTION = "execute" + CATEGORY = "KJNodes/weights" + DESCRIPTION = """ +Gets the mean values from mask or image batch +and returns that as the selected output type. +""" + + def execute(self, output_type, images=None, masks=None): + mean_values = [] + if masks is not None and images is None: + for mask in masks: + mean_values.append(mask.mean().item()) + elif masks is None and images is not None: + for image in images: + mean_values.append(image.mean().item()) + elif masks is not None and images is not None: + raise ValueError("MaskOrImageToWeight: Use either mask or image input only.") + + # Convert mean_values to the specified output_type + if output_type == 'list': + out = mean_values + elif output_type == 'pandas series': + try: + import pandas as pd + except ImportError as e: + raise ImportError("MaskOrImageToWeight: pandas is not installed. Please install pandas to use this output_type") from e + out = pd.Series(mean_values), + elif output_type == 'tensor': + out = torch.tensor(mean_values, dtype=torch.float32), + return (out, [str(value) for value in mean_values],) + +class WeightScheduleConvert: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_values": ("FLOAT", {"default": 0.0, "forceInput": True}), + "output_type": ( + [ + 'match_input', + 'list', + 'pandas series', + 'tensor', + ], + { + "default": 'list' + }), + "invert": ("BOOLEAN", {"default": False}), + "repeat": ("INT", {"default": 1,"min": 1, "max": 255, "step": 1}), + }, + "optional": { + "remap_to_frames": ("INT", {"default": 0}), + "interpolation_curve": ("FLOAT", {"forceInput": True}), + "remap_values": ("BOOLEAN", {"default": False}), + "remap_min": ("FLOAT", {"default": 0.0, "min": -100000, "max": 100000.0, "step": 0.01}), + "remap_max": ("FLOAT", {"default": 1.0, "min": -100000, "max": 100000.0, "step": 0.01}), + }, + + } + RETURN_TYPES = ("FLOAT", "STRING", "INT",) + FUNCTION = "execute" + CATEGORY = "KJNodes/weights" + DESCRIPTION = """ +Converts different value lists/series to another type. +""" + + def detect_input_type(self, input_values): + import pandas as pd + if isinstance(input_values, list): + return 'list' + elif isinstance(input_values, pd.Series): + return 'pandas series' + elif isinstance(input_values, torch.Tensor): + return 'tensor' + else: + raise ValueError("Unsupported input type") + + def execute(self, input_values, output_type, invert, repeat, remap_to_frames=0, interpolation_curve=None, remap_min=0.0, remap_max=1.0, remap_values=False): + import pandas as pd + input_type = self.detect_input_type(input_values) + + if input_type == 'pandas series': + float_values = input_values.tolist() + elif input_type == 'tensor': + float_values = input_values + else: + float_values = input_values + + if invert: + float_values = [1 - value for value in float_values] + + if interpolation_curve is not None: + interpolated_pattern = [] + orig_float_values = float_values + for value in interpolation_curve: + min_val = min(orig_float_values) + max_val = max(orig_float_values) + # Normalize the values to [0, 1] + normalized_values = [(value - min_val) / (max_val - min_val) for value in orig_float_values] + # Interpolate the normalized values to the new frame count + remapped_float_values = np.interp(np.linspace(0, 1, int(remap_to_frames * value)), np.linspace(0, 1, len(normalized_values)), normalized_values).tolist() + interpolated_pattern.extend(remapped_float_values) + float_values = interpolated_pattern + else: + # Remap float_values to match target_frame_amount + if remap_to_frames > 0 and remap_to_frames != len(float_values): + min_val = min(float_values) + max_val = max(float_values) + # Normalize the values to [0, 1] + normalized_values = [(value - min_val) / (max_val - min_val) for value in float_values] + # Interpolate the normalized values to the new frame count + float_values = np.interp(np.linspace(0, 1, remap_to_frames), np.linspace(0, 1, len(normalized_values)), normalized_values).tolist() + + float_values = float_values * repeat + if remap_values: + float_values = self.remap_values(float_values, remap_min, remap_max) + + if output_type == 'list': + out = float_values, + elif output_type == 'pandas series': + out = pd.Series(float_values), + elif output_type == 'tensor': + if input_type == 'pandas series': + out = torch.tensor(float_values.values, dtype=torch.float32), + else: + out = torch.tensor(float_values, dtype=torch.float32), + elif output_type == 'match_input': + out = float_values, + return (out, [str(value) for value in float_values], [int(value) for value in float_values]) + + def remap_values(self, values, target_min, target_max): + # Determine the current range + current_min = min(values) + current_max = max(values) + current_range = current_max - current_min + + # Determine the target range + target_range = target_max - target_min + + # Perform the linear interpolation for each value + remapped_values = [(value - current_min) / current_range * target_range + target_min for value in values] + + return remapped_values + + +class FloatToMask: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_values": ("FLOAT", {"forceInput": True, "default": 0}), + "width": ("INT", {"default": 100, "min": 1}), + "height": ("INT", {"default": 100, "min": 1}), + }, + } + RETURN_TYPES = ("MASK",) + FUNCTION = "execute" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Generates a batch of masks based on the input float values. +The batch size is determined by the length of the input float values. +Each mask is generated with the specified width and height. +""" + + def execute(self, input_values, width, height): + import pandas as pd + # Ensure input_values is a list + if isinstance(input_values, (float, int)): + input_values = [input_values] + elif isinstance(input_values, pd.Series): + input_values = input_values.tolist() + elif isinstance(input_values, list) and all(isinstance(item, list) for item in input_values): + input_values = [item for sublist in input_values for item in sublist] + + # Generate a batch of masks based on the input_values + masks = [] + for value in input_values: + # Assuming value is a float between 0 and 1 representing the mask's intensity + mask = torch.ones((height, width), dtype=torch.float32) * value + masks.append(mask) + masks_out = torch.stack(masks, dim=0) + + return(masks_out,) +class WeightScheduleExtend: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_values_1": ("FLOAT", {"default": 0.0, "forceInput": True}), + "input_values_2": ("FLOAT", {"default": 0.0, "forceInput": True}), + "output_type": ( + [ + 'match_input', + 'list', + 'pandas series', + 'tensor', + ], + { + "default": 'match_input' + }), + }, + + } + RETURN_TYPES = ("FLOAT",) + FUNCTION = "execute" + CATEGORY = "KJNodes/weights" + DESCRIPTION = """ +Extends, and converts if needed, different value lists/series +""" + + def detect_input_type(self, input_values): + import pandas as pd + if isinstance(input_values, list): + return 'list' + elif isinstance(input_values, pd.Series): + return 'pandas series' + elif isinstance(input_values, torch.Tensor): + return 'tensor' + else: + raise ValueError("Unsupported input type") + + def execute(self, input_values_1, input_values_2, output_type): + import pandas as pd + input_type_1 = self.detect_input_type(input_values_1) + input_type_2 = self.detect_input_type(input_values_2) + # Convert input_values_2 to the same format as input_values_1 if they do not match + if not input_type_1 == input_type_2: + print("Converting input_values_2 to the same format as input_values_1") + if input_type_1 == 'pandas series': + # Convert input_values_2 to a pandas Series + float_values_2 = pd.Series(input_values_2) + elif input_type_1 == 'tensor': + # Convert input_values_2 to a tensor + float_values_2 = torch.tensor(input_values_2, dtype=torch.float32) + else: + print("Input types match, no conversion needed") + # If the types match, no conversion is needed + float_values_2 = input_values_2 + + float_values = input_values_1 + float_values_2 + + if output_type == 'list': + return float_values, + elif output_type == 'pandas series': + return pd.Series(float_values), + elif output_type == 'tensor': + if input_type_1 == 'pandas series': + return torch.tensor(float_values.values, dtype=torch.float32), + else: + return torch.tensor(float_values, dtype=torch.float32), + elif output_type == 'match_input': + return float_values, + else: + raise ValueError(f"Unsupported output_type: {output_type}") + +class FloatToSigmas: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "float_list": ("FLOAT", {"default": 0.0, "forceInput": True}), + } + } + RETURN_TYPES = ("SIGMAS",) + RETURN_NAMES = ("SIGMAS",) + CATEGORY = "KJNodes/noise" + FUNCTION = "customsigmas" + DESCRIPTION = """ +Creates a sigmas tensor from list of float values. + +""" + def customsigmas(self, float_list): + return torch.tensor(float_list, dtype=torch.float32), + +class SigmasToFloat: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "sigmas": ("SIGMAS",), + } + } + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("float",) + CATEGORY = "KJNodes/noise" + FUNCTION = "customsigmas" + DESCRIPTION = """ +Creates a float list from sigmas tensors. + +""" + def customsigmas(self, sigmas): + return sigmas.tolist(), + +class GLIGENTextBoxApplyBatchCoords: + @classmethod + def INPUT_TYPES(s): + return {"required": {"conditioning_to": ("CONDITIONING", ), + "latents": ("LATENT", ), + "clip": ("CLIP", ), + "gligen_textbox_model": ("GLIGEN", ), + "coordinates": ("STRING", {"forceInput": True}), + "text": ("STRING", {"multiline": True}), + "width": ("INT", {"default": 128, "min": 8, "max": 4096, "step": 8}), + "height": ("INT", {"default": 128, "min": 8, "max": 4096, "step": 8}), + }, + "optional": {"size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True})}, + } + RETURN_TYPES = ("CONDITIONING", "IMAGE", ) + RETURN_NAMES = ("conditioning", "coord_preview", ) + FUNCTION = "append" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = """ +This node allows scheduling GLIGEN text box positions in a batch, +to be used with AnimateDiff-Evolved. Intended to pair with the +Spline Editor -node. + +GLIGEN model can be downloaded through the Manage's "Install Models" menu. +Or directly from here: +https://huggingface.co/comfyanonymous/GLIGEN_pruned_safetensors/tree/main + +Inputs: +- **latents** input is used to calculate batch size +- **clip** is your standard text encoder, use same as for the main prompt +- **gligen_textbox_model** connects to GLIGEN Loader +- **coordinates** takes a json string of points, directly compatible +with the spline editor node. +- **text** is the part of the prompt to set position for +- **width** and **height** are the size of the GLIGEN bounding box + +Outputs: +- **conditioning** goes between to clip text encode and the sampler +- **coord_preview** is an optional preview of the coordinates and +bounding boxes. + +""" + + def append(self, latents, coordinates, conditioning_to, clip, gligen_textbox_model, text, width, height, size_multiplier=[1.0]): + coordinates = json.loads(coordinates.replace("'", '"')) + coordinates = [(coord['x'], coord['y']) for coord in coordinates] + + batch_size = sum(tensor.size(0) for tensor in latents.values()) + if len(coordinates) != batch_size: + print("GLIGENTextBoxApplyBatchCoords WARNING: The number of coordinates does not match the number of latents") + + c = [] + _, cond_pooled = clip.encode_from_tokens(clip.tokenize(text), return_pooled=True) + + for t in conditioning_to: + n = [t[0], t[1].copy()] + + position_params_batch = [[] for _ in range(batch_size)] # Initialize a list of empty lists for each batch item + if len(size_multiplier) != batch_size: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + + for i in range(batch_size): + x_position, y_position = coordinates[i] + position_param = (cond_pooled, int((height // 8) * size_multiplier[i]), int((width // 8) * size_multiplier[i]), (y_position - height // 2) // 8, (x_position - width // 2) // 8) + position_params_batch[i].append(position_param) # Append position_param to the correct sublist + + prev = [] + if "gligen" in n[1]: + prev = n[1]['gligen'][2] + else: + prev = [[] for _ in range(batch_size)] + # Concatenate prev and position_params_batch, ensuring both are lists of lists + # and each sublist corresponds to a batch item + combined_position_params = [prev_item + batch_item for prev_item, batch_item in zip(prev, position_params_batch)] + n[1]['gligen'] = ("position_batched", gligen_textbox_model, combined_position_params) + c.append(n) + + image_height = latents['samples'].shape[-2] * 8 + image_width = latents['samples'].shape[-1] * 8 + plot_image_tensor = plot_coordinates_to_tensor(coordinates, image_height, image_width, height, width, size_multiplier, text) + + return (c, plot_image_tensor,) + +class CreateInstanceDiffusionTracking: + + RETURN_TYPES = ("TRACKING", "STRING", "INT", "INT", "INT", "INT",) + RETURN_NAMES = ("tracking", "prompt", "width", "height", "bbox_width", "bbox_height",) + FUNCTION = "tracking" + CATEGORY = "KJNodes/InstanceDiffusion" + DESCRIPTION = """ +Creates tracking data to be used with InstanceDiffusion: +https://github.com/logtd/ComfyUI-InstanceDiffusion + +InstanceDiffusion prompt format: +"class_id.class_name": "prompt", +for example: +"1.head": "((head))", +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "coordinates": ("STRING", {"forceInput": True}), + "width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "bbox_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "bbox_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "class_name": ("STRING", {"default": "class_name"}), + "class_id": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "prompt": ("STRING", {"default": "prompt", "multiline": True}), + }, + "optional": { + "size_multiplier": ("FLOAT", {"default": [1.0], "forceInput": True}), + "fit_in_frame": ("BOOLEAN", {"default": True}), + } + } + + def tracking(self, coordinates, class_name, class_id, width, height, bbox_width, bbox_height, prompt, size_multiplier=[1.0], fit_in_frame=True): + # Define the number of images in the batch + coordinates = coordinates.replace("'", '"') + coordinates = json.loads(coordinates) + + tracked = {} + tracked[class_name] = {} + batch_size = len(coordinates) + # Initialize a list to hold the coordinates for the current ID + id_coordinates = [] + if not size_multiplier or len(size_multiplier) != batch_size: + size_multiplier = [0] * batch_size + else: + size_multiplier = size_multiplier * (batch_size // len(size_multiplier)) + size_multiplier[:batch_size % len(size_multiplier)] + for i, coord in enumerate(coordinates): + x = coord['x'] + y = coord['y'] + adjusted_bbox_width = bbox_width * size_multiplier[i] + adjusted_bbox_height = bbox_height * size_multiplier[i] + # Calculate the top left and bottom right coordinates + top_left_x = x - adjusted_bbox_width // 2 + top_left_y = y - adjusted_bbox_height // 2 + bottom_right_x = x + adjusted_bbox_width // 2 + bottom_right_y = y + adjusted_bbox_height // 2 + + if fit_in_frame: + # Clip the coordinates to the frame boundaries + top_left_x = max(0, top_left_x) + top_left_y = max(0, top_left_y) + bottom_right_x = min(width, bottom_right_x) + bottom_right_y = min(height, bottom_right_y) + # Ensure width and height are positive + adjusted_bbox_width = max(1, bottom_right_x - top_left_x) + adjusted_bbox_height = max(1, bottom_right_y - top_left_y) + + # Update the coordinates with the new width and height + bottom_right_x = top_left_x + adjusted_bbox_width + bottom_right_y = top_left_y + adjusted_bbox_height + + # Append the top left and bottom right coordinates to the list for the current ID + id_coordinates.append([top_left_x, top_left_y, bottom_right_x, bottom_right_y, width, height]) + + class_id = int(class_id) + # Assign the list of coordinates to the specified ID within the class_id dictionary + tracked[class_name][class_id] = id_coordinates + + prompt_string = "" + for class_name, class_data in tracked.items(): + for class_id in class_data.keys(): + class_id_str = str(class_id) + # Use the incoming prompt for each class name and ID + prompt_string += f'"{class_id_str}.{class_name}": "({prompt})",\n' + + # Remove the last comma and newline + prompt_string = prompt_string.rstrip(",\n") + + return (tracked, prompt_string, width, height, bbox_width, bbox_height) + +class AppendInstanceDiffusionTracking: + + RETURN_TYPES = ("TRACKING", "STRING",) + RETURN_NAMES = ("tracking", "prompt",) + FUNCTION = "append" + CATEGORY = "KJNodes/InstanceDiffusion" + DESCRIPTION = """ +Appends tracking data to be used with InstanceDiffusion: +https://github.com/logtd/ComfyUI-InstanceDiffusion + +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "tracking_1": ("TRACKING", {"forceInput": True}), + "tracking_2": ("TRACKING", {"forceInput": True}), + }, + "optional": { + "prompt_1": ("STRING", {"default": "", "forceInput": True}), + "prompt_2": ("STRING", {"default": "", "forceInput": True}), + } + } + + def append(self, tracking_1, tracking_2, prompt_1="", prompt_2=""): + tracking_copy = tracking_1.copy() + # Check for existing class names and class IDs, and raise an error if they exist + for class_name, class_data in tracking_2.items(): + if class_name not in tracking_copy: + tracking_copy[class_name] = class_data + else: + # If the class name exists, merge the class data from tracking_2 into tracking_copy + # This will add new class IDs under the same class name without raising an error + tracking_copy[class_name].update(class_data) + prompt_string = prompt_1 + "," + prompt_2 + return (tracking_copy, prompt_string) + +class InterpolateCoords: + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("coordinates",) + FUNCTION = "interpolate" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = """ +Interpolates coordinates based on a curve. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "coordinates": ("STRING", {"forceInput": True}), + "interpolation_curve": ("FLOAT", {"forceInput": True}), + + }, + } + + def interpolate(self, coordinates, interpolation_curve): + # Parse the JSON string to get the list of coordinates + coordinates = json.loads(coordinates.replace("'", '"')) + + # Convert the list of dictionaries to a list of (x, y) tuples for easier processing + coordinates = [(coord['x'], coord['y']) for coord in coordinates] + + # Calculate the total length of the original path + path_length = sum(np.linalg.norm(np.array(coordinates[i]) - np.array(coordinates[i-1])) + for i in range(1, len(coordinates))) + + # Initialize variables for interpolation + interpolated_coords = [] + current_length = 0 + current_index = 0 + + # Iterate over the normalized curve + for normalized_length in interpolation_curve: + target_length = normalized_length * path_length # Convert to the original scale + while current_index < len(coordinates) - 1: + segment_start, segment_end = np.array(coordinates[current_index]), np.array(coordinates[current_index + 1]) + segment_length = np.linalg.norm(segment_end - segment_start) + if current_length + segment_length >= target_length: + break + current_length += segment_length + current_index += 1 + + # Interpolate between the last two points + if current_index < len(coordinates) - 1: + p1, p2 = np.array(coordinates[current_index]), np.array(coordinates[current_index + 1]) + segment_length = np.linalg.norm(p2 - p1) + if segment_length > 0: + t = (target_length - current_length) / segment_length + interpolated_point = p1 + t * (p2 - p1) + interpolated_coords.append(interpolated_point.tolist()) + else: + interpolated_coords.append(p1.tolist()) + else: + # If the target_length is at or beyond the end of the path, add the last coordinate + interpolated_coords.append(coordinates[-1]) + + # Convert back to string format if necessary + interpolated_coords_str = "[" + ", ".join([f"{{'x': {round(coord[0])}, 'y': {round(coord[1])}}}" for coord in interpolated_coords]) + "]" + print(interpolated_coords_str) + + return (interpolated_coords_str,) + +class DrawInstanceDiffusionTracking: + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image", ) + FUNCTION = "draw" + CATEGORY = "KJNodes/InstanceDiffusion" + DESCRIPTION = """ +Draws the tracking data from +CreateInstanceDiffusionTracking -node. + +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "tracking": ("TRACKING", {"forceInput": True}), + "box_line_width": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}), + "draw_text": ("BOOLEAN", {"default": True}), + "font": (folder_paths.get_filename_list("kjnodes_fonts"), ), + "font_size": ("INT", {"default": 20}), + }, + } + + def draw(self, image, tracking, box_line_width, draw_text, font, font_size): + import matplotlib.cm as cm + + modified_images = [] + + colormap = cm.get_cmap('rainbow', len(tracking)) + if draw_text: + font_path = folder_paths.get_full_path("kjnodes_fonts", font) + font = ImageFont.truetype(font_path, font_size) + + # Iterate over each image in the batch + for i in range(image.shape[0]): + # Extract the current image and convert it to a PIL image + current_image = image[i, :, :, :].permute(2, 0, 1) + pil_image = transforms.ToPILImage()(current_image) + + draw = ImageDraw.Draw(pil_image) + + # Iterate over the bounding boxes for the current image + for j, (class_name, class_data) in enumerate(tracking.items()): + for class_id, bbox_list in class_data.items(): + # Check if the current index is within the bounds of the bbox_list + if i < len(bbox_list): + bbox = bbox_list[i] + # Ensure bbox is a list or tuple before unpacking + if isinstance(bbox, (list, tuple)): + x1, y1, x2, y2, _, _ = bbox + # Convert coordinates to integers + x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) + # Generate a color from the rainbow colormap + color = tuple(int(255 * x) for x in colormap(j / len(tracking)))[:3] + # Draw the bounding box on the image with the generated color + draw.rectangle([x1, y1, x2, y2], outline=color, width=box_line_width) + if draw_text: + # Draw the class name and ID as text above the box with the generated color + text = f"{class_id}.{class_name}" + # Calculate the width and height of the text + _, _, text_width, text_height = draw.textbbox((0, 0), text=text, font=font) + # Position the text above the top-left corner of the box + text_position = (x1, y1 - text_height) + draw.text(text_position, text, fill=color, font=font) + else: + print(f"Unexpected data type for bbox: {type(bbox)}") + + # Convert the drawn image back to a torch tensor and adjust back to (H, W, C) + modified_image_tensor = transforms.ToTensor()(pil_image).permute(1, 2, 0) + modified_images.append(modified_image_tensor) + + # Stack the modified images back into a batch + image_tensor_batch = torch.stack(modified_images).cpu().float() + + return image_tensor_batch, + +class PointsEditor: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "points_store": ("STRING", {"multiline": False, "advanced": True}), + "coordinates": ("STRING", {"multiline": False, "socketless": True, "advanced": True}), + "neg_coordinates": ("STRING", {"multiline": False, "socketless": True, "advanced": True}), + "bbox_store": ("STRING", {"multiline": False, "advanced": True}), + "bboxes": ("STRING", {"multiline": False, "socketless": True, "advanced": True}), + "bbox_format": ( + [ + 'xyxy', + 'xywh', + ], + ), + "width": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}), + "height": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}), + "normalize": ("BOOLEAN", {"default": False}), + }, + "optional": { + "bg_image": ("IMAGE", ), + }, + } + + RETURN_TYPES = ("STRING", "STRING", "BBOX", "MASK", "IMAGE") + RETURN_NAMES = ("positive_coords", "negative_coords", "bbox", "bbox_mask", "cropped_image") + FUNCTION = "pointdata" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = """ +# WORK IN PROGRESS +Do not count on this as part of your workflow yet, +probably contains lots of bugs and stability is not +guaranteed!! + +## Graphical editor to create coordinates + +**Shift + click** to add a positive (green) point. +**Shift + right click** to add a negative (red) point. +**Right click on a point** to delete it. +**Ctrl + click** to draw a bounding box. +**Drag bbox corners** to resize, **drag inside** to move. +**Right click on bbox** to delete it. + +To add an image select the node and copy/paste or drag in the image. +Or from the bg_image input on queue (first frame of the batch). + +**THE IMAGE IS SAVED TO THE NODE AND WORKFLOW METADATA** +you can clear the image from the context menu by right clicking on the canvas + +""" + + def pointdata(self, points_store, bbox_store, width, height, coordinates, neg_coordinates, normalize, bboxes, bbox_format="xyxy", bg_image=None): + coordinates = json.loads(coordinates) + if not coordinates: + raise ValueError("No points on the canvas. Use Shift+click to add positive points or Shift+right-click to add negative points before executing.") + pos_coordinates = [] + for coord in coordinates: + coord['x'] = int(round(coord['x'])) + coord['y'] = int(round(coord['y'])) + if normalize: + norm_x = coord['x'] / width + norm_y = coord['y'] / height + pos_coordinates.append({'x': norm_x, 'y': norm_y}) + else: + pos_coordinates.append({'x': coord['x'], 'y': coord['y']}) + + if neg_coordinates: + coordinates = json.loads(neg_coordinates) + neg_coordinates = [] + for coord in coordinates: + coord['x'] = int(round(coord['x'])) + coord['y'] = int(round(coord['y'])) + if normalize: + norm_x = coord['x'] / width + norm_y = coord['y'] / height + neg_coordinates.append({'x': norm_x, 'y': norm_y}) + else: + neg_coordinates.append({'x': coord['x'], 'y': coord['y']}) + + # Create a blank mask + mask = np.zeros((height, width), dtype=np.uint8) + bboxes = json.loads(bboxes) + valid_bboxes = [] + for bbox in bboxes: + if (bbox.get("startX") is None or + bbox.get("startY") is None or + bbox.get("endX") is None or + bbox.get("endY") is None): + continue # Skip this bounding box if any value is None + else: + # Ensure that endX and endY are greater than startX and startY + x_min = min(int(bbox["startX"]), int(bbox["endX"])) + y_min = min(int(bbox["startY"]), int(bbox["endY"])) + x_max = max(int(bbox["startX"]), int(bbox["endX"])) + y_max = max(int(bbox["startY"]), int(bbox["endY"])) + + valid_bboxes.append((x_min, y_min, x_max, y_max)) + + bboxes_xyxy = [] + for bbox in valid_bboxes: + x_min, y_min, x_max, y_max = bbox + bboxes_xyxy.append((x_min, y_min, x_max, y_max)) + mask[y_min:y_max, x_min:x_max] = 1 # Fill the bounding box area with 1s + + if bbox_format == "xywh": + bboxes_xywh = [] + for bbox in valid_bboxes: + x_min, y_min, x_max, y_max = bbox + width = x_max - x_min + height = y_max - y_min + bboxes_xywh.append((x_min, y_min, width, height)) + bboxes = bboxes_xywh + else: + bboxes = bboxes_xyxy + + mask_tensor = torch.from_numpy(mask) + mask_tensor = mask_tensor.unsqueeze(0).float().cpu() + + if bg_image is not None and len(valid_bboxes) > 0: + x_min, y_min, x_max, y_max = bboxes[0] + cropped_image = bg_image[:, y_min:y_max, x_min:x_max, :] + elif bg_image is not None: + cropped_image = bg_image + else: + cropped_image = torch.zeros(1, height, width, 3, dtype=torch.float32) + + if bg_image is None: + return (json.dumps(pos_coordinates), json.dumps(neg_coordinates), bboxes, mask_tensor, cropped_image) + else: + transform = transforms.ToPILImage() + image = transform(bg_image[0].permute(2, 0, 1)) + buffered = BytesIO() + image.save(buffered, format="JPEG", quality=75) + img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8') + + return { + "ui": {"bg_image": [img_base64]}, + "result": (json.dumps(pos_coordinates), json.dumps(neg_coordinates), bboxes, mask_tensor, cropped_image) + } + +class CutAndDragOnPath: + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image","mask", ) + FUNCTION = "cutanddrag" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Cuts the masked area from the image, and drags it along the path. If inpaint is enabled, and no bg_image is provided, the cut area is filled using cv2 TELEA algorithm. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "coordinates": ("STRING", {"forceInput": True}), + "mask": ("MASK",), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "inpaint": ("BOOLEAN", {"default": True}), + }, + "optional": { + "bg_image": ("IMAGE",), + } + } + + def cutanddrag(self, image, coordinates, mask, frame_width, frame_height, inpaint, bg_image=None): + # Parse coordinates + coords_list = parse_json_tracks(coordinates) + + batch_size = len(coords_list[0]) + images_list = [] + masks_list = [] + + # Convert input image and mask to PIL + input_image = tensor2pil(image)[0] + input_mask = tensor2pil(mask)[0] + + # Find masked region bounds + mask_array = np.array(input_mask) + y_indices, x_indices = np.where(mask_array > 0) + if len(x_indices) == 0 or len(y_indices) == 0: + return (image, mask) + + x_min, x_max = x_indices.min(), x_indices.max() + y_min, y_max = y_indices.min(), y_indices.max() + + # Cut out the masked region + cut_width = x_max - x_min + cut_height = y_max - y_min + cut_image = input_image.crop((x_min, y_min, x_max, y_max)) + cut_mask = input_mask.crop((x_min, y_min, x_max, y_max)) + + # Create inpainted background + if bg_image is None: + background = input_image.copy() + # Inpaint the cut area + if inpaint: + import cv2 + border = 5 # Create small border around cut area for better inpainting + fill_mask = Image.new("L", background.size, 0) + draw = ImageDraw.Draw(fill_mask) + draw.rectangle([x_min-border, y_min-border, x_max+border, y_max+border], fill=255) + background = cv2.inpaint( + np.array(background), + np.array(fill_mask), + inpaintRadius=3, + flags=cv2.INPAINT_TELEA + ) + background = Image.fromarray(background) + else: + background = tensor2pil(bg_image)[0] + + # Create batch of images with cut region at different positions + for i in range(batch_size): + # Create new image + new_image = background.copy() + new_mask = Image.new("L", (frame_width, frame_height), 0) + + # Get target position from coordinates + for coords in coords_list: + target_x = int(coords[i]['x'] - cut_width/2) + target_y = int(coords[i]['y'] - cut_height/2) + + # Paste cut region at new position + new_image.paste(cut_image, (target_x, target_y), cut_mask) + new_mask.paste(cut_mask, (target_x, target_y)) + + # Convert to tensor and append + image_tensor = pil2tensor(new_image) + mask_tensor = pil2tensor(new_mask) + + images_list.append(image_tensor) + masks_list.append(mask_tensor) + + # Stack tensors into batches + out_images = torch.cat(images_list, dim=0).cpu().float() + out_masks = torch.cat(masks_list, dim=0) + + return (out_images, out_masks) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/hdr_preview_node.py b/custom_nodes/ComfyUI-KJNodes/nodes/hdr_preview_node.py new file mode 100644 index 0000000000000000000000000000000000000000..82d2f585a963de873815c4dacf1a4e1419bc13ed --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/hdr_preview_node.py @@ -0,0 +1,166 @@ +import os +import random + +import torch +from PIL import Image +from comfy_api.latest import io +import comfy.model_management as mm +import folder_paths + + +# LogC3 constants (ARRI LogC3, EI 800). Kept in sync with the WebGL shader in web/js/hdr_preview.js. +LC_A = 5.555556 +LC_B = 0.052272 +LC_C = 0.247190 +LC_D = 0.385537 +LC_E = 5.367655 +LC_F = 0.092809 +LC_CUT = 0.010591 +LC_CUT_LOG = LC_E * LC_CUT + LC_F # ~0.14966 + +def _logc3_decompress(logc: torch.Tensor) -> torch.Tensor: + logc = logc.clamp(0.0, 1.0) + lin_from_log = (10.0 ** ((logc - LC_D) / LC_C) - LC_B) / LC_A + lin_from_lin = (logc - LC_F) / LC_E + return torch.where(logc >= LC_CUT_LOG, lin_from_log, lin_from_lin) + + +def _linear_to_srgb(x: torch.Tensor) -> torch.Tensor: + cutoff = 0.0031308 + return torch.where( + x <= cutoff, + 12.92 * x, + 1.055 * torch.pow(x.clamp(min=cutoff), 1.0 / 2.4) - 0.055, + ).clamp_(0.0, 1.0) + + +def _srgb_to_linear(x: torch.Tensor) -> torch.Tensor: + cutoff = 0.04045 + return torch.where( + x <= cutoff, + x / 12.92, + ((x.clamp(min=0.0) + 0.055) / 1.055) ** 2.4, + ) + + +class HDRPreviewKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="HDRPreviewKJ", + display_name="HDR Preview KJ", + category="KJNodes/image", + is_output_node=True, + is_experimental=True, + description=( + "Realtime-exposure preview for HDR-compressed images.\n\n" + "Input: LogC3-compressed [0,1] image/video batch (e.g. the VAE-decoded output " + "of an HDR IC-LoRA workflow, prior to HDR decompression).\n\n" + "Decompression + exposure + saturation + Reinhard tonemap + sRGB runs in a WebGL " + "fragment shader in the browser for realtime slider feedback, and the same math " + "runs server-side to produce the baked sRGB IMAGE output. Slider changes update " + "the preview immediately; the IMAGE output only updates when the workflow is re-queued." + ), + inputs=[ + io.Image.Input("image", + tooltip="LogC3-compressed HDR image/video in [0,1], or linear HDR if 'input_space' is 'linear'."), + io.Float.Input("exposure", default=0.0, min=-10.0, max=10.0, step=0.01, + tooltip="Exposure in EV stops. 0 = no change, +1 = 2x brighter."), + io.Float.Input("saturation", default=1.0, min=0.0, max=2.0, step=0.01, + tooltip="Saturation multiplier. 0 = grayscale, 1 = unchanged, 2 = 2x."), + io.Float.Input("fps", default=24.0, min=1.0, max=120.0, step=0.1, optional=True, + tooltip="Playback frame rate for video (batch) inputs."), + io.Combo.Input("input_space", options=["logc3", "linear", "srgb"], default="logc3", optional=True, + tooltip="Color space of input. 'logc3' = ARRI LogC3 compressed HDR; 'linear' = linear HDR directly; 'srgb' = already-graded sRGB image (skips Reinhard tonemap)."), + ], + outputs=[ + io.Image.Output(display_name="image", + tooltip="Tonemapped sRGB image, ready for preview/save."), + ], + ) + + @classmethod + def execute(cls, image: torch.Tensor, exposure: float = 0.0, saturation: float = 1.0, fps: float = 24.0, input_space: str = "logc3") -> io.NodeOutput: + temp_dir = folder_paths.get_temp_directory() + os.makedirs(temp_dir, exist_ok=True) + prefix = f"hdrprv_{random.randint(0, 0xFFFFFF):06x}" + + B, H, W, _ = image.shape + device = mm.get_torch_device() + exposure_mul = 2.0 ** exposure + luma_weights = torch.tensor([0.2126, 0.7152, 0.0722], device=device) + + bytes_per_frame = H * W * 3 * 4 + chunk_size = max(1, min(B, int(1_000_000_000 // max(bytes_per_frame * 10, 1)))) + + # For linear input we need the global max across all frames to normalize previews. + norm_scale = 1.0 + if input_space == "linear": + max_val = float(image[..., :3].max().item()) + norm_scale = max_val if max_val > 1.0 else 1.0 + + filenames = [] + srgb_chunks = [] + + for start in range(0, B, chunk_size): + end = min(start + chunk_size, B) + image_rgb = image[start:end, ..., :3].float().to(device, non_blocking=True) + + # --- Preview frames (8-bit PNG, always in the "raw" pre-exposure space) --- + if input_space == "linear": + preview = (image_rgb / norm_scale).clamp_(0.0, 1.0) + else: + preview = image_rgb.clamp(0.0, 1.0) + + preview_np = preview.mul_(255.0).add_(0.5).clamp_(0.0, 255.0).to(torch.uint8).cpu().numpy() + del preview + + for i in range(end - start): + fname = f"{prefix}_{start + i:05d}.png" + Image.fromarray(preview_np[i], mode="RGB").save( + os.path.join(temp_dir, fname), + format="PNG", + compress_level=1, + ) + filenames.append(fname) + del preview_np + + # --- Baked sRGB output (same math as the shader in hdr_preview.js) --- + if input_space == "logc3": + hdr = _logc3_decompress(image_rgb).clamp_(min=0.0) + elif input_space == "srgb": + hdr = _srgb_to_linear(image_rgb).clamp_(min=0.0) + else: + hdr = image_rgb.clamp(min=0.0) + del image_rgb + + exposed = hdr.mul_(exposure_mul) + luma = (exposed * luma_weights.to(exposed.dtype)).sum(dim=-1, keepdim=True) + saturated = (luma + (exposed - luma) * saturation).clamp_(min=0.0) + del luma, exposed + + if input_space == "srgb": + # Already display-ready linear; skip Reinhard, just clip over-exposed highlights. + tonemapped = saturated.clamp_(0.0, 1.0) + else: + tonemapped = saturated / (1.0 + saturated) + del saturated + + srgb_chunks.append(_linear_to_srgb(tonemapped).cpu()) + del tonemapped + + srgb = torch.cat(srgb_chunks, dim=0) + del srgb_chunks + + data = { + "frames": [{"filename": f, "type": "temp"} for f in filenames], + "width": int(W), + "height": int(H), + "fps": float(fps), + "input_space": input_space, + "linear_scale": float(norm_scale), + "frame_count": int(B), + "exposure": float(exposure), + "saturation": float(saturation), + } + return io.NodeOutput(srgb, ui={"hdr_preview_data": [data]}) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/ideogram4_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/ideogram4_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..4b033b53da7496456ed849e526c30e8b32331acb --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/ideogram4_nodes.py @@ -0,0 +1,411 @@ +"""Ideogram 4 prompt builder. + +A single self-contained node with a visual bbox editor: draw regions on a blank +canvas, set each region's type/desc/text/color palette, and assemble the Ideogram 4 JSON caption prompt. +""" + +import json +import os + +import numpy as np +import torch +from PIL import Image, ImageDraw, ImageFont, ImageEnhance + +from comfy_api.latest import io + + +_FONT_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "fonts", "FreeMono.ttf") + + +def _hex_rgb(h): + h = h.lstrip("#") + return (int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)) if len(h) == 6 else (255, 255, 255) + + +def _readable(rgb): + # Lighten toward white if too dark, so box-colored text stays legible on the dark canvas. + r, g, b = rgb + lum = 0.299 * r + 0.587 * g + 0.114 * b + if lum < 130: + t = (130 - lum) / (255 - lum) + r, g, b = round(r + (255 - r) * t), round(g + (255 - g) * t), round(b + (255 - b) * t) + return (r, g, b) + + +def _font(size): + try: + return ImageFont.truetype(_FONT_PATH, size) + except Exception: + try: + return ImageFont.load_default(size) + except Exception: + return ImageFont.load_default() + + +def _wrap(draw, text, font, max_w): + lines = [] + for para in text.split("\n"): + line = "" + for word in para.split(): + test = word if not line else line + " " + word + if line and draw.textlength(test, font=font) > max_w: + lines.append(line) + line = word + else: + line = test + lines.append(line) + return lines + + +def _render_preview(boxes, width, height, bg=None, brightness=50): + # Render the regions + prompts over the reference image (or a black canvas). + if bg is not None: + iw, ih = bg.size + long_edge = max(iw, ih) + scale = min(1.0, 1024 / long_edge) if long_edge > 0 else 1.0 + rw, rh = max(1, round(iw * scale)), max(1, round(ih * scale)) + base = bg.convert("RGB").resize((rw, rh), Image.LANCZOS) + if brightness < 100: # dim to match the editor's brightness slider + base = ImageEnhance.Brightness(base).enhance(max(0.0, brightness / 100.0)) + img = base.convert("RGBA") + else: + long_edge = max(width, height) + scale = min(1.0, 1024 / long_edge) if long_edge > 0 else 1.0 + rw = max(1, round(width * scale)) + rh = max(1, round(height * scale)) + img = Image.new("RGBA", (rw, rh), (0, 0, 0, 255)) # black so the overlay composites cleanly + overlay = Image.new("RGBA", (rw, rh), (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay) + fs = max(10, round(rh / 64)) + font = _font(fs) + tag_font = _font(max(9, fs - 2)) + lh = fs + 2 + + for i, box in enumerate(boxes): + if not isinstance(box, dict) or box.get("nobbox"): + continue # skip unplaced elements (no real location) + palette = [c for c in (box.get("palette") or []) if c] + r, g, b = _hex_rgb(palette[0]) if palette else (140, 140, 140) # box = first palette color, else grey + x1 = max(0, min(rw, round(box.get("x", 0) * rw))) + y1 = max(0, min(rh, round(box.get("y", 0) * rh))) + x2 = max(0, min(rw, round((box.get("x", 0) + box.get("w", 0)) * rw))) + y2 = max(0, min(rh, round((box.get("y", 0) + box.get("h", 0)) * rh))) + if x2 < x1: + x1, x2 = x2, x1 + if y2 < y1: + y1, y2 = y2, y1 + + draw.rectangle([x1, y1, x2, y2], outline=(r, g, b, 255), width=2) + + pal5 = palette[:5] # palette shown as a strip along the top edge + if pal5 and (x2 - x1) > 2: + sh = max(5, fs // 2) + seg = (x2 - x1) / len(pal5) + for p, hexc in enumerate(pal5): + sx = x1 + round(p * seg) + draw.rectangle([sx, y1, x1 + round((p + 1) * seg), y1 + sh], fill=_hex_rgb(hexc)) + + etype = "text" if box.get("type") == "text" else "obj" + tag = str(i + 1).zfill(2) + tw = draw.textlength(tag, font=tag_font) + draw.rectangle([x1, y1, x1 + tw + 6, y1 + fs + 2], fill=(r, g, b, 255)) # tag chip = box color + tagfill = (0, 0, 0, 255) if (0.299 * r + 0.587 * g + 0.114 * b) > 140 else (255, 255, 255, 255) + draw.text((x1 + 3, y1 + 1), tag, fill=tagfill, font=tag_font) + + body = box.get("desc", "") or "" + if etype == "text" and box.get("text"): + body = '"%s"%s' % (box["text"], " — " + body if body else "") + if body and (x2 - x1) > 8: + ty = y1 + fs + 5 + for line in _wrap(draw, body, font, x2 - x1 - 8): + if ty > y2: + break + draw.text((x1 + 4, ty), line, fill=_readable((r, g, b)) + (255,), font=font) + ty += lh + + img = Image.alpha_composite(img, overlay).convert("RGB") + arr = np.asarray(img, dtype=np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) + + +def _norm_bbox(box): + # Normalized {x, y, w, h} fractions (0-1) -> [ymin, xmin, ymax, xmax] on a 0-1000 grid. + def c(v): + return max(0, min(1000, round(v * 1000))) + x, y, w, h = box.get("x", 0.0), box.get("y", 0.0), box.get("w", 0.0), box.get("h", 0.0) + ymin, xmin, ymax, xmax = c(y), c(x), c(y + h), c(x + w) + if ymin > ymax: + ymin, ymax = ymax, ymin + if xmin > xmax: + xmin, xmax = xmax, xmin + return [ymin, xmin, ymax, xmax] + + +def _palette(colors): + # ["#rrggbb", ...] (or autogrow dict) -> ["#RRGGBB", ...] in order, dropping empties. + if isinstance(colors, dict): + colors = colors.values() + return [c.upper() for c in colors if c] + + +def _dumps(v, lvl=0): + # Like json.dumps(ensure_ascii=False, indent=4), but scalar arrays stay on one line. + pad, end = " " * (lvl + 1), " " * lvl + if isinstance(v, str): + return json.dumps(v, ensure_ascii=False) + if isinstance(v, list): + if not v: + return "[]" + if all(not isinstance(x, (dict, list)) for x in v): + return "[" + ", ".join(_dumps(x, lvl) for x in v) + "]" + return "[\n" + ",\n".join(pad + _dumps(x, lvl + 1) for x in v) + "\n" + end + "]" + if isinstance(v, dict): + if not v: + return "{}" + items = [pad + json.dumps(k, ensure_ascii=False) + ": " + _dumps(val, lvl + 1) for k, val in v.items()] + return "{\n" + ",\n".join(items) + "\n" + end + "}" + return json.dumps(v, ensure_ascii=False) + + +def _parse_json_list(s): + if s: + try: + v = json.loads(s) + if isinstance(v, list): + return v + except json.JSONDecodeError: + pass + return [] + + +def _caption_to_boxes(cap): + # Caption dict -> editor box list ({x,y,w,h, type, text, desc, palette}) for preview/bboxes. + cd = cap.get("compositional_deconstruction") or {} + boxes = [] + for el in (cd.get("elements") or []): + if not isinstance(el, dict): + continue + box = {"type": "text" if el.get("type") == "text" else "obj", + "text": el.get("text", "") or "", "desc": el.get("desc", "") or "", + "palette": list(el.get("color_palette") or [])} + bb = el.get("bbox") + if isinstance(bb, (list, tuple)) and len(bb) == 4: + ymin, xmin, ymax, xmax = bb + box.update(x=xmin / 1000.0, y=ymin / 1000.0, + w=(xmax - xmin) / 1000.0, h=(ymax - ymin) / 1000.0) + else: # no bbox: unplaced placeholder + box.update(x=0.03, y=0.03, w=0.22, h=0.14, nobbox=True) + boxes.append(box) + return boxes + + +class Ideogram4PromptBuilderKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="Ideogram4PromptBuilderKJ", + display_name="Ideogram 4 Prompt Builder KJ", + category="KJNodes/text", + search_aliases=["ideogram", "caption", "bbox", "prompt builder", "json prompt"], + is_experimental=True, + description=""" +Visual prompt builder for Ideogram 4's structured JSON caption format. + +Drag on the canvas to draw regions; select a region to set its type (obj/text), +description, text, and color palette. Set the background and optional style fields +as widgets. Outputs the assembled caption JSON string. + +bbox is normalized to a 0-1000 grid as [ymin, xmin, ymax, xmax]; width/height set +the canvas aspect ratio. + +Canvas controls: +- Drag: draw a new region +- Ctrl/Cmd-drag: force-draw a new region even on top of an existing one +- Click: select a region · Alt-click: cycle overlapping regions +- Double-click: edit the description inline +- Right-click: region list (select / delete / duplicate / reorder, top = front) +- Del / Backspace: remove the selected region +- Ctrl/Cmd + C / V / D: copy / paste / duplicate the selected region +- bbox fields (px / out) next to obj/text are editable + +Color swatches: +- Click: edit · Drag: reorder · Right-click: remove +- Hover + Ctrl/Cmd + C / V: copy / paste the hex +- "+": add a color (uses the clipboard color if it is one) + +Toolbar: +- Live: use the live sampling preview as the background (and grab the final result) +- Grab BG / Clear BG: use the last generated image as the background +- brightness slider, token estimate, and Copy / Paste / Clear all""", + inputs=[ + io.Int.Input("width", default=1024, min=64, max=16384, step=16, + tooltip="Canvas aspect width (also the pixel grid the bbox is measured in). Ideogram 4 needs multiples of 16."), + io.Int.Input("height", default=1024, min=64, max=16384, step=16, + tooltip="Canvas aspect height (also the pixel grid the bbox is measured in). Ideogram 4 needs multiples of 16."), + io.String.Input("high_level_description", multiline=True, default="", + tooltip="Optional one-line overview of the whole image (blank = omitted)."), + io.String.Input("background", multiline=True, default="", + tooltip="Required scene background description."), + io.DynamicCombo.Input("style", options=[ + io.DynamicCombo.Option("none", []), + io.DynamicCombo.Option("photo", [ + io.String.Input("photo", default=""), + ]), + io.DynamicCombo.Option("art_style", [ + io.String.Input("art_style", default=""), + ]), + ]), + io.String.Input("aesthetics", default="", tooltip="Style descriptor (blank = omitted)."), + io.String.Input("lighting", default="", tooltip="Style descriptor (blank = omitted)."), + io.String.Input("medium", default="", tooltip="Style descriptor (blank = omitted)."), + io.Image.Input("image", optional=True, + tooltip="Optional reference image shown as the editor background (and behind the preview)."), + io.String.Input("import_json", default="", optional=True, force_input=True, + tooltip="Optional: a full caption JSON. When connected, it loads into the editor " + "and drives the output per 'import_mode'."), + io.String.Input("style_palette_data", default="", socketless=True, advanced=True, + tooltip="Serialized style color palette from the editor (managed by the node UI)."), + io.String.Input("elements_data", default="", socketless=True, advanced=True, + tooltip="Serialized regions from the editor (managed by the node UI)."), + io.Int.Input("bg_brightness", default=25, min=0, max=100, socketless=True, advanced=True, + tooltip="Background image brightness % (managed by the node UI slider)."), + io.Combo.Input("import_mode", options=["when empty", "always"], default="when empty", + tooltip="How a wired import_json is used: 'when empty' only seeds the editor while " + "it has no regions (then the editor wins, so you can edit); 'always' makes " + "the wired JSON authoritative so its changes always propagate to the output."), + io.String.Input("output_format", default="compact", socketless=True, advanced=True, + tooltip="Output JSON formatting (set via the editor toolbar): 'compact' (default, what " + "Ideogram 4 expects) or 'pretty' (indented, for readability)."), + io.BoundingBox.Input("bboxes", optional=True, force_input=True, + tooltip="Optional pixel-space boxes ({x, y, width, height}) used to seed the " + "editor's regions when it has none. Ignored once regions exist."), + ], + outputs=[ + io.String.Output(display_name="prompt"), + io.Image.Output(display_name="preview"), + io.BoundingBox.Output(display_name="bboxes"), + io.Int.Output(display_name="width"), + io.Int.Output(display_name="height"), + ], + ) + + @classmethod + def execute(cls, width, height, background, style, + high_level_description="", aesthetics="", lighting="", medium="", + style_palette_data="", elements_data="", import_json="", import_mode="when empty", + output_format="pretty", bboxes=None, image=None, bg_brightness=25) -> io.NodeOutput: + if import_mode not in ("when empty", "always"): # old workflows saved before this widget existed + import_mode = "when empty" + dump = _dumps if output_format == "pretty" else (lambda v: json.dumps(v, ensure_ascii=False, separators=(",", ":"))) + boxes = _parse_json_list(elements_data) + boxes_seeded = False + if not boxes and bboxes: + if isinstance(bboxes, dict): # a single BoundingBox is a bare {x,y,width,height} dict + frame = [bboxes] + elif bboxes and isinstance(bboxes[0], (list, tuple)): + frame = bboxes[0] # per-frame nesting: [[box, ...], ...] + else: + frame = bboxes # flat list of boxes + for bb in frame: + if not isinstance(bb, dict): + continue + boxes.append({"x": bb.get("x", 0) / width, "y": bb.get("y", 0) / height, + "w": bb.get("width", 0) / width, "h": bb.get("height", 0) / height, + "type": "obj", "text": "", "desc": "", "palette": []}) + boxes_seeded = bool(boxes) + + imported = None + if import_json and import_json.strip(): + try: + c = json.loads(import_json) + if isinstance(c, dict): + imported = c + except json.JSONDecodeError: + pass + + kind = style["style"] # "none" | "photo" | "art_style" + + # Use the wired import_json directly per import_mode: "always" -> authoritative (its changes + # always propagate); "when empty" -> only seed the editor while it has no regions, then the + # editor wins so manual edits stick. The editor mirrors it via ui when used. + used_import = imported is not None and (import_mode == "always" or not boxes) + + if used_import: + caption = imported + boxes = _caption_to_boxes(imported) + else: + caption = {} + if high_level_description.strip(): + caption["high_level_description"] = high_level_description + + if kind != "none": + # The verifier requires every style key present (in order) once a style is + # chosen; only color_palette is conditional. Emit blanks rather than omit. + sd = {"aesthetics": aesthetics, "lighting": lighting} + # photo: ...photo, medium... | art_style: ...medium, art_style... (key order) + if kind == "photo": + sd["photo"] = style.get("photo", "") + sd["medium"] = medium + else: + sd["medium"] = medium + sd["art_style"] = style.get("art_style", "") + palette = _palette(_parse_json_list(style_palette_data)) + if palette: + sd["color_palette"] = palette + caption["style_description"] = sd + + elements = [] + for box in boxes: + if not isinstance(box, dict): + continue + etype = "text" if box.get("type") == "text" else "obj" + elem = {"type": etype} # key order matters + if not box.get("nobbox"): # unplaced elements omit bbox + elem["bbox"] = _norm_bbox(box) + if etype == "text": + elem["text"] = box.get("text", "") + elem["desc"] = box.get("desc", "") + palette = _palette(box.get("palette", [])) + if palette: + elem["color_palette"] = palette[:5] + elements.append(elem) + + caption["compositional_deconstruction"] = { + "background": background, + "elements": elements, + } + bg = None + if image is not None: # composite over the input image, else black + try: + bg = Image.fromarray((image[0].detach().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)) + except Exception: + bg = None + preview = _render_preview(boxes, width, height, bg, bg_brightness) + + # Pixel-space bboxes ({x, y, width, height}) for SAM3 / BoundingBox consumers. + bbox_dicts = [] + for box in boxes: + if not isinstance(box, dict) or box.get("nobbox"): + continue + x, y = box.get("x", 0.0), box.get("y", 0.0) + bw, bh = box.get("w", 0.0), box.get("h", 0.0) + if bw < 0: + x += bw + bw = -bw + if bh < 0: + y += bh + bh = -bh + bbox_dicts.append({"x": round(x * width), "y": round(y * height), + "width": round(bw * width), "height": round(bh * height)}) + # Per-frame nesting (list[list[dict]]) — the canonical BoundingBox shape that + # SAM3 / crop nodes expect (bboxes[frame] -> list of boxes). + bboxes_out = [bbox_dicts] if bbox_dicts else [] + + # ui: send the resolved width/height so the editor canvas can follow connected + # inputs; import_json (if wired) loads into the editor (output reflects editor only). + ui = {"dims": [width, height]} + if boxes_seeded: + ui["boxes"] = [json.dumps(boxes)] + if used_import: # mirror the import in the editor (only when used) + ui["caption"] = [_dumps(imported)] + return io.NodeOutput(dump(caption), preview, bboxes_out, width, height, ui=ui) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/image_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/image_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..cfe460af9e92337edf1feea0732d7fd68032a1ce --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/image_nodes.py @@ -0,0 +1,5156 @@ +import numpy as np +import time +import torch +import torch.nn.functional as F +import torchvision.transforms as T +import base64 +import random +import math +import os +import re +import json +import importlib +import hashlib +import pathlib +import logging +from io import BytesIO + +try: + import cv2 + HAS_CV2 = True +except ImportError: + logging.warning("OpenCV not installed") + HAS_CV2 = False + +from PIL import ImageGrab, ImageDraw, ImageFont, Image, ImageOps, ImageSequence, ImageStat +from PIL.PngImagePlugin import PngInfo + +from nodes import MAX_RESOLUTION, SaveImage +from comfy_extras.nodes_mask import composite +from comfy.cli_args import args +from comfy.utils import ProgressBar, common_upscale, tiled_scale_multidim +from comfy import model_management +from comfy_api.latest import io, InputImpl, Types, ui +from fractions import Fraction +import node_helpers +import folder_paths + +from ..utility.utility import string_to_color + +try: + from server import PromptServer, BinaryEventTypes +except ImportError: + PromptServer = None + BinaryEventTypes = None +from concurrent.futures import ThreadPoolExecutor + +script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +class ImagePass: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "image": ("IMAGE",), + }, + } + RETURN_TYPES = ("IMAGE",) + FUNCTION = "passthrough" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Passes the image through without modifying it. +""" + + def passthrough(self, image=None): + return image, + +class ColorMatch: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_ref": ("IMAGE",), + "image_target": ("IMAGE",), + "method": (['mkl','hm', 'reinhard', 'mvgd', 'hm-mvgd-hm', 'hm-mkl-hm'], { + "default": 'mkl' + }), + }, + "optional": { + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "multithread": ("BOOLEAN", {"default": True}), + } + } + + CATEGORY = "KJNodes/image" + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "colormatch" + DEPRECATED = True + DESCRIPTION = """ +color-matcher enables color transfer across images which comes in handy for automatic +color-grading of photographs, paintings and film sequences as well as light-field +and stopmotion corrections. + +The methods behind the mappings are based on the approach from Reinhard et al., +the Monge-Kantorovich Linearization (MKL) as proposed by Pitie et al. and our analytical solution +to a Multi-Variate Gaussian Distribution (MVGD) transfer in conjunction with classical histogram +matching. As shown below our HM-MVGD-HM compound outperforms existing methods. +https://github.com/hahnec/color-matcher/ + +""" + + def colormatch(self, image_ref, image_target, method, strength=1.0, multithread=True): + # Skip unnecessary processing + if strength == 0: + return (image_target,) + + try: + from color_matcher import ColorMatcher + except ImportError as e: + raise ImportError("Can't import color-matcher, did you install requirements.txt? Manual install: pip install color-matcher") from e + + image_ref = image_ref.cpu() + image_target = image_target.cpu() + batch_size = image_target.size(0) + + images_target = image_target.squeeze() + images_ref = image_ref.squeeze() + + image_ref_np = images_ref.numpy() + images_target_np = images_target.numpy() + + def process(i): + cm = ColorMatcher() + image_target_np_i = images_target_np if batch_size == 1 else images_target[i].numpy() + image_ref_np_i = image_ref_np if image_ref.size(0) == 1 else images_ref[i].numpy() + try: + image_result = cm.transfer(src=image_target_np_i, ref=image_ref_np_i, method=method) # Avoid potential blur when only the fully color-matched image is used + if strength != 1: + image_result = image_target_np_i + strength * (image_result - image_target_np_i) + + return torch.from_numpy(image_result) + + except Exception as e: + logging.warning(f"Thread {i} error: {e}") + return torch.from_numpy(image_target_np_i) # fallback + + if multithread and batch_size > 1: + max_threads = min(os.cpu_count() or 1, batch_size) + with ThreadPoolExecutor(max_workers=max_threads) as executor: + out = list(executor.map(process, range(batch_size))) + else: + out = [process(i) for i in range(batch_size)] + + out = torch.stack(out, dim=0).to(torch.float32) + out.clamp_(0, 1) + return (out,) + +class ColorMatchV2(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ColorMatchV2", + category="KJNodes/image", + description=""" +color-matcher enables color transfer across images which comes in handy for automatic +color-grading of photographs, paintings and film sequences as well as light-field +and stopmotion corrections. + +The methods behind the mappings are based on the approach from Reinhard et al., +the Monge-Kantorovich Linearization (MKL) as proposed by Pitie et al. and our analytical solution +to a Multi-Variate Gaussian Distribution (MVGD) transfer in conjunction with classical histogram +matching. As shown below our HM-MVGD-HM compound outperforms existing methods. +https://github.com/hahnec/color-matcher/ + +'reinhard_lab_gpu' method uses Kornia for GPU-accelerated color transfer in Lab color space. +""", + inputs=[ + io.Image.Input("image_target"), + io.Image.Input("image_ref"), + io.Combo.Input("method", + options=['mkl', 'hm', 'reinhard', 'mvgd', 'hm-mvgd-hm', 'hm-mkl-hm', 'reinhard_lab_gpu'], + default='mkl'), + io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.01), + io.Boolean.Input("multithread", default=True), + ], + outputs=[ + io.Image.Output(display_name="image"), + ], + ) + + @classmethod + def execute(cls, image_target, image_ref, method, strength=1.0, multithread=True) -> io.NodeOutput: + # Skip unnecessary processing + if strength == 0: + return io.NodeOutput(image_target) + + if method == "reinhard_lab_gpu": + import kornia + device = model_management.get_torch_device() + + B, H, W, C = image_target.shape + + src_bchw = image_target.to(device).permute(0, 3, 1, 2).contiguous() # (B, H, W, C) -> (B, C, H, W) + ref_bchw = image_ref.to(device).permute(0, 3, 1, 2).contiguous() + # RGB->Lab + src_lab = kornia.color.rgb_to_lab(src_bchw) + ref_lab = kornia.color.rgb_to_lab(ref_bchw) + + src_lab_flat = src_lab.view(B, C, -1) # (B, C, HW) + ref_lab_flat = ref_lab.view(ref_lab.shape[0], C, -1) # (B or 1, C, HW) + + src_std, src_mean = torch.std_mean(src_lab_flat, dim=-1, keepdim=True, unbiased=False) + ref_std, ref_mean = torch.std_mean(ref_lab_flat, dim=-1, keepdim=True, unbiased=False) + src_std = src_std.clamp_min_(1e-6) + + if ref_lab.shape[0] == 1 and B > 1: + ref_mean = ref_mean.expand(B, -1, -1) + ref_std = ref_std.expand(B, -1, -1) + + corrected_lab_flat = (src_lab_flat - src_mean) * (ref_std / src_std) + ref_mean + corrected_lab = corrected_lab_flat.view(B, C, H, W) + + # Lab->RGB + corrected_rgb_01 = kornia.color.lab_to_rgb(corrected_lab) + out = (1.0 - strength) * src_bchw + strength * corrected_rgb_01 + out = out.permute(0, 2, 3, 1).contiguous() # (B, C, H, W) -> (B, H, W, C) + + return io.NodeOutput(out.cpu().float().clamp_(0, 1)) + + try: + from color_matcher import ColorMatcher + except ImportError as e: + raise ImportError("Can't import color-matcher, did you install requirements.txt? Manual install: pip install color-matcher") from e + + batch_size = image_target.size(0) + ref_batch_size = image_ref.size(0) + + def process(i): + cm = ColorMatcher() + image_target_np = image_target[i].cpu().numpy() + image_ref_np = image_ref[min(i, ref_batch_size - 1)].cpu().numpy() + try: + image_result = cm.transfer(src=image_target_np, ref=image_ref_np, method=method) + if strength != 1: + image_result = image_target_np + strength * (image_result - image_target_np) + + return torch.from_numpy(image_result) + + except Exception as e: + logging.error(f"Thread {i} error: {e}") + return torch.from_numpy(image_target_np) # fallback + if multithread and batch_size > 1: + max_threads = min(os.cpu_count() or 1, batch_size) + with ThreadPoolExecutor(max_workers=max_threads) as executor: + out = list(executor.map(process, range(batch_size))) + else: + out = [process(i) for i in range(batch_size)] + + out = torch.stack(out, dim=0).to(torch.float32).clamp_(0, 1) + + return io.NodeOutput(out) + +class SaveImageWithAlpha: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + + @classmethod + def INPUT_TYPES(s): + return {"required": + {"images": ("IMAGE", ), + "mask": ("MASK", ), + "filename_prefix": ("STRING", {"default": "ComfyUI"})}, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = () + FUNCTION = "save_images_alpha" + OUTPUT_NODE = True + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Saves an image and mask as .PNG with the mask as the alpha channel. +""" + + def save_images_alpha(self, images, mask, filename_prefix="ComfyUI_image_with_alpha", prompt=None, extra_pnginfo=None): + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + results = list() + def file_counter(): + max_counter = 0 + # Loop through the existing files + for existing_file in sorted(os.listdir(full_output_folder)): + # Check if the file matches the expected format + match = re.fullmatch(fr"{filename}_(\d+)_?\.[a-zA-Z0-9]+", existing_file) + if match: + # Extract the numeric portion of the filename + file_counter = int(match.group(1)) + # Update the maximum counter value if necessary + if file_counter > max_counter: + max_counter = file_counter + return max_counter + + for image, alpha in zip(images, mask): + i = 255. * image.cpu().numpy() + a = 255. * (1.0 - alpha.cpu().float()).numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + # Resize the mask to match the image size + a_resized = Image.fromarray(a).resize(img.size, Image.LANCZOS) + a_resized = np.clip(a_resized, 0, 255).astype(np.uint8) + img.putalpha(Image.fromarray(a_resized, mode='L')) + metadata = None + if not args.disable_metadata: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + # Increment the counter by 1 to get the next available value + counter = file_counter() + 1 + file = f"{filename}_{counter:05}.png" + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + + return { "ui": { "images": results } } + + +class ImageConcanate(io.ComfyNode): + @classmethod + def define_schema(cls): + # image1 drives the output type; image2 can independently be IMAGE or MASK and gets + # converted to image1's type inside concatenate(). + type_template = io.MatchType.Template("image_or_mask", allowed_types=[io.Image, io.Mask]) + return io.Schema( + node_id="ImageConcanate", + category="KJNodes/image", + description=( + "Concatenates image2 to image1 in the specified direction.\n" + "Both inputs accept IMAGE or MASK; the output type follows image1.\n" + "If image2 is a different type than image1 it's converted (RGB mean for image→mask,\n" + "channel-replicate for mask→image).\n" + "When match_image_size is False and dimensions don't match along the shared axis,\n" + "the smaller image is centered and zero-padded instead of erroring." + ), + inputs=[ + io.MatchType.Input("image1", template=type_template), + io.MultiType.Input("image2", types=[io.Image, io.Mask]), + io.Combo.Input("direction", options=['right', 'down', 'left', 'up'], default='right'), + io.Boolean.Input("match_image_size", default=True), + ], + outputs=[ + io.MatchType.Output(template=type_template, display_name="output"), + ], + ) + + @classmethod + def execute(cls, image1, image2, direction, match_image_size) -> io.NodeOutput: + return io.NodeOutput(cls.concatenate(image1, image2, direction, match_image_size)) + + @staticmethod + def concatenate(image1, image2, direction, match_image_size, first_image_shape=None): + # IMAGE is BHWC, MASK is BHW. Output type follows image1; convert image2 to match, + # then unsqueeze any masks to BHW1 so the rest of the function can stay BHWC-only. + output_is_mask = image1.dim() == 3 + if output_is_mask and image2.dim() == 4: + ch = min(3, image2.shape[-1]) + image2 = image2[..., :ch].mean(dim=-1) + elif not output_is_mask and image2.dim() == 3: + image2 = image2.unsqueeze(-1).expand(-1, -1, -1, image1.shape[-1]) + if output_is_mask: + image1 = image1.unsqueeze(-1) + image2 = image2.unsqueeze(-1) + + bs1 = image1.shape[0] + bs2 = image2.shape[0] + B = max(bs1, bs2) + + H1, W1 = image1.shape[1], image1.shape[2] + C1, C2 = image1.shape[-1], image2.shape[-1] + out_C = max(C1, C2) + + if match_image_size: + target_shape = first_image_shape if first_image_shape is not None else image1.shape + orig_aspect = image2.shape[2] / image2.shape[1] + if direction in ('left', 'right'): + H2 = target_shape[1] + W2 = int(H2 * orig_aspect) + else: + W2 = target_shape[2] + H2 = int(W2 / orig_aspect) + else: + H2, W2 = image2.shape[1], image2.shape[2] + + if direction in ('right', 'left'): + out_H, out_W = max(H1, H2), W1 + W2 + else: + out_H, out_W = H1 + H2, max(W1, W2) + + if direction == 'right': + i1_y, i1_x, i2_y, i2_x = (out_H - H1) // 2, 0, (out_H - H2) // 2, W1 + elif direction == 'left': + i1_y, i1_x, i2_y, i2_x = (out_H - H1) // 2, W2, (out_H - H2) // 2, 0 + elif direction == 'down': + i1_y, i1_x, i2_y, i2_x = 0, (out_W - W1) // 2, H1, (out_W - W2) // 2 + else: # 'up' + i1_y, i1_x, i2_y, i2_x = H2, (out_W - W1) // 2, 0, (out_W - W2) // 2 + + output = torch.zeros( + (B, out_H, out_W, out_C), + dtype=model_management.intermediate_dtype(), + device=model_management.intermediate_device(), + ) + + def write(dst, src, src_C): + if dst.shape[-1] == src_C: + dst.copy_(src) + else: + dst[..., :src_C].copy_(src) + dst[..., src_C:].fill_(1.0) + + slot1 = output[:, i1_y:i1_y + H1, i1_x:i1_x + W1, :] + if bs1 == B: + write(slot1, image1, C1) + else: + write(slot1[:bs1], image1, C1) + write(slot1[bs1:], image1[-1:].expand(B - bs1, -1, -1, -1), C1) + del slot1 + + slot2 = output[:, i2_y:i2_y + H2, i2_x:i2_x + W2, :] + if match_image_size: + pbar = ProgressBar(B) + device = model_management.get_torch_device() + for i in range(B): + src_i = min(i, bs2 - 1) + frame = image2[src_i:src_i + 1].to(device, non_blocking=True).permute(0, 3, 1, 2) + resized = F.interpolate(frame, size=(H2, W2), mode='bicubic', antialias=True).permute(0, 2, 3, 1) + write(slot2[i:i + 1], resized, C2) + del frame, resized + pbar.update(1) + else: + if bs2 == B: + write(slot2, image2, C2) + else: + write(slot2[:bs2], image2, C2) + write(slot2[bs2:], image2[-1:].expand(B - bs2, -1, -1, -1), C2) + del slot2 + + if output_is_mask: + return output.squeeze(-1) + return output + + +class ImageConcatFromBatch: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE",), + "num_columns": ("INT", {"default": 3, "min": 1, "max": 255, "step": 1}), + "match_image_size": ("BOOLEAN", {"default": False}), + "max_resolution": ("INT", {"default": 4096}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "concat" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ + Concatenates images from a batch into a grid with a specified number of columns. + """ + + def concat(self, images, num_columns, match_image_size, max_resolution): + # Assuming images is a batch of images (B, H, W, C) + batch_size, height, width, channels = images.shape + num_rows = (batch_size + num_columns - 1) // num_columns # Calculate number of rows + + logging.info(f"Initial dimensions: batch_size={batch_size}, height={height}, width={width}, channels={channels}") + logging.info(f"num_rows={num_rows}, num_columns={num_columns}") + + if match_image_size: + target_shape = images[0].shape + + resized_images = [] + for image in images: + original_height = image.shape[0] + original_width = image.shape[1] + original_aspect_ratio = original_width / original_height + + if original_aspect_ratio > 1: + target_height = target_shape[0] + target_width = int(target_height * original_aspect_ratio) + else: + target_width = target_shape[1] + target_height = int(target_width / original_aspect_ratio) + + logging.info(f"Resizing image from ({original_height}, {original_width}) to ({target_height}, {target_width})") + + # Resize the image to match the target size while preserving aspect ratio + resized_image = common_upscale(image.movedim(-1, 0), target_width, target_height, "lanczos", "disabled") + resized_image = resized_image.movedim(0, -1) # Move channels back to the last dimension + resized_images.append(resized_image) + + # Convert the list of resized images back to a tensor + images = torch.stack(resized_images) + + height, width = target_shape[:2] # Update height and width + + # Initialize an empty grid + grid_height = num_rows * height + grid_width = num_columns * width + + logging.info(f"Grid dimensions before scaling: grid_height={grid_height}, grid_width={grid_width}") + + # Original scale factor calculation remains unchanged + scale_factor = min(max_resolution / grid_height, max_resolution / grid_width, 1.0) + + # Apply scale factor to height and width + scaled_height = height * scale_factor + scaled_width = width * scale_factor + + # Round scaled dimensions to the nearest number divisible by 8 + height = max(1, int(round(scaled_height / 8) * 8)) + width = max(1, int(round(scaled_width / 8) * 8)) + + if abs(scaled_height - height) > 4: + height = max(1, int(round((scaled_height + 4) / 8) * 8)) + if abs(scaled_width - width) > 4: + width = max(1, int(round((scaled_width + 4) / 8) * 8)) + + # Recalculate grid dimensions with adjusted height and width + grid_height = num_rows * height + grid_width = num_columns * width + logging.info(f"Grid dimensions after scaling: grid_height={grid_height}, grid_width={grid_width}") + logging.info(f"Final image dimensions: height={height}, width={width}") + + grid = torch.zeros((grid_height, grid_width, channels), dtype=images.dtype) + + for idx, image in enumerate(images): + resized_image = torch.nn.functional.interpolate(image.unsqueeze(0).permute(0, 3, 1, 2), size=(height, width), mode="bilinear").squeeze().permute(1, 2, 0) + row = idx // num_columns + col = idx % num_columns + grid[row*height:(row+1)*height, col*width:(col+1)*width, :] = resized_image + + return grid.unsqueeze(0), + + +class ImageGridComposite2x2: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "image3": ("IMAGE",), + "image4": ("IMAGE",), + }} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "compositegrid" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Concatenates the 4 input images into a 2x2 grid. +""" + + def compositegrid(self, image1, image2, image3, image4): + top_row = torch.cat((image1, image2), dim=2) + bottom_row = torch.cat((image3, image4), dim=2) + grid = torch.cat((top_row, bottom_row), dim=1) + return (grid,) + +class ImageGridComposite3x3: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "image3": ("IMAGE",), + "image4": ("IMAGE",), + "image5": ("IMAGE",), + "image6": ("IMAGE",), + "image7": ("IMAGE",), + "image8": ("IMAGE",), + "image9": ("IMAGE",), + }} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "compositegrid" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Concatenates the 9 input images into a 3x3 grid. +""" + + def compositegrid(self, image1, image2, image3, image4, image5, image6, image7, image8, image9): + top_row = torch.cat((image1, image2, image3), dim=2) + mid_row = torch.cat((image4, image5, image6), dim=2) + bottom_row = torch.cat((image7, image8, image9), dim=2) + grid = torch.cat((top_row, mid_row, bottom_row), dim=1) + return (grid,) + + +class ImageBatchTestPattern(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ImageBatchTestPattern", + category="KJNodes/text", + description="Generate a batch of images with sequential numbers rendered in a chosen font.", + inputs=[ + io.Int.Input("batch_size", default=1, min=1, max=4096, step=1), + io.Int.Input("start_from", default=0, min=0, max=4096, step=1), + io.Int.Input("text_x", default=256, min=0, max=4096, step=1), + io.Int.Input("text_y", default=256, min=0, max=4096, step=1), + io.Int.Input("width", default=512, min=16, max=4096, step=1), + io.Int.Input("height", default=512, min=16, max=4096, step=1), + io.Combo.Input("font", options=folder_paths.get_filename_list("kjnodes_fonts")), + io.Int.Input("font_size", default=255, min=8, max=4096, step=1), + ], + outputs=[ + io.Image.Output(display_name="image"), + ], + ) + + @classmethod + def execute(cls, batch_size, font, font_size, start_from, width, height, text_x, text_y) -> io.NodeOutput: + font_path = folder_paths.get_full_path("kjnodes_fonts", font) + pil_font = ImageFont.truetype(font_path, font_size) + + # Probe once whether the '-liga' feature is supported by this PIL build/font + use_liga = True + try: + ImageDraw.Draw(Image.new("RGB", (1, 1))).text( + (0, 0), "0", font=pil_font, fill=(0, 0, 0), features=['-liga'] + ) + except Exception: + use_liga = False + + image = Image.new("RGB", (width, height), color='black') + draw = ImageDraw.Draw(image) + + out_buf = np.empty((batch_size, height, width, 3), dtype=np.uint8) + pbar = ProgressBar(batch_size) + + for i in range(batch_size): + # Reset canvas to black instead of allocating a new PIL image + draw.rectangle((0, 0, width, height), fill='black') + + font_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) + text = str(start_from + i) + + if use_liga: + draw.text((text_x, text_y), text, font=pil_font, fill=font_color, features=['-liga']) + else: + draw.text((text_x, text_y), text, font=pil_font, fill=font_color) + + out_buf[i] = np.asarray(image) + pbar.update(1) + + out_tensor = torch.from_numpy(out_buf).to( + device=model_management.intermediate_device(), + dtype=model_management.intermediate_dtype(), + ).div_(255.0) + return io.NodeOutput(out_tensor) + +class ImageGrabPIL: + + @classmethod + def IS_CHANGED(cls): + + return + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "screencap" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Captures an area specified by screen coordinates. +Can be used for realtime diffusion with autoqueue. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "x": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "width": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), + "num_frames": ("INT", {"default": 1,"min": 1, "max": 255, "step": 1}), + "delay": ("FLOAT", {"default": 0.1,"min": 0.0, "max": 10.0, "step": 0.01}), + }, + } + + def screencap(self, x, y, width, height, num_frames, delay): + start_time = time.time() + captures = [] + bbox = (x, y, x + width, y + height) + + for _ in range(num_frames): + # Capture screen + screen_capture = ImageGrab.grab(bbox=bbox) + screen_capture_torch = torch.from_numpy(np.array(screen_capture, dtype=np.float32) / 255.0).unsqueeze(0) + captures.append(screen_capture_torch) + + # Wait for a short delay if more than one frame is to be captured + if num_frames > 1: + time.sleep(delay) + + elapsed_time = time.time() - start_time + logging.info(f"screengrab took {elapsed_time} seconds.") + + return (torch.cat(captures, dim=0),) + +class Screencap_mss: + + @classmethod + def IS_CHANGED(s, **kwargs): + return float("NaN") + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "screencap" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Captures an area specified by screen coordinates. +Can be used for realtime diffusion with autoqueue. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "x": ("INT", {"default": 0,"min": 0, "max": 10000, "step": 1}), + "y": ("INT", {"default": 0,"min": 0, "max": 10000, "step": 1}), + "width": ("INT", {"default": 512,"min": 0, "max": 10000, "step": 1}), + "height": ("INT", {"default": 512,"min": 0, "max": 10000, "step": 1}), + "num_frames": ("INT", {"default": 1,"min": 1, "max": 255, "step": 1}), + "delay": ("FLOAT", {"default": 0.1,"min": 0.0, "max": 10.0, "step": 0.01}), + }, + } + + def screencap(self, x, y, width, height, num_frames, delay): + from mss import mss + captures = [] + with mss() as sct: + bbox = {'top': y, 'left': x, 'width': width, 'height': height} + + for _ in range(num_frames): + sct_img = sct.grab(bbox) + img_np = np.array(sct_img) + img_torch = torch.from_numpy(img_np[..., [2, 1, 0]]).float() / 255.0 + captures.append(img_torch) + + if num_frames > 1: + time.sleep(delay) + + return (torch.stack(captures, 0),) + +class ScreencapStream: + + @classmethod + def IS_CHANGED(s, **kwargs): + return float("NaN") + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "capture" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Captures a frame from a browser screen/window share stream. +Click 'Start capture' to select a screen or window to share. +Live preview is shown in the node. Works with auto-queue. + +Crop controls: +- Drag on preview to draw a crop box +- Drag inside the box to move it +- Drag edges or corners to resize +- Shift+drag to lock aspect ratio +- Right-click or double-click to clear crop +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frame_data": ("STRING", {"default": "", "multiline": False}), + "crop_width": ("INT", {"default": 1, "min": 1, "max": MAX_RESOLUTION, "step": 1}), + "crop_height": ("INT", {"default": 1, "min": 1, "max": MAX_RESOLUTION, "step": 1}), + }, + } + + MAX_FRAME_BYTES = 50 * 1024 * 1024 # 50MB base64 limit (PNG is larger than JPEG) + + def capture(self, crop_width, crop_height, frame_data): + if not frame_data: + w = crop_width if crop_width > 0 else 512 + h = crop_height if crop_height > 0 else 512 + return (torch.zeros(1, h, w, 3),) + if len(frame_data) > self.MAX_FRAME_BYTES: + raise ValueError(f"Frame data exceeds {self.MAX_FRAME_BYTES // (1024*1024)}MB limit") + try: + img_bytes = base64.b64decode(frame_data.split(",", 1)[-1]) + except Exception: + raise ValueError("Invalid frame data encoding") + img = Image.open(BytesIO(img_bytes)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_np).unsqueeze(0) + return (img_tensor,) + +class WebcamCaptureCV2: + + @classmethod + def IS_CHANGED(cls): + return + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "capture" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = """ +Captures a frame from a webcam using CV2. +Can be used for realtime diffusion with autoqueue. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "x": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "width": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), + "cam_index": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "release": ("BOOLEAN", {"default": False}), + }, + } + + def capture(self, x, y, cam_index, width, height, release): + # Check if the camera index has changed or the capture object doesn't exist + if not hasattr(self, "cap") or self.cap is None or self.current_cam_index != cam_index: + if hasattr(self, "cap") and self.cap is not None: + self.cap.release() + self.current_cam_index = cam_index + self.cap = cv2.VideoCapture(cam_index) + try: + self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width) + self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height) + except cv2.error: + pass + if not self.cap.isOpened(): + raise RuntimeError("Could not open webcam") + + ret, frame = self.cap.read() + if not ret: + raise RuntimeError("Failed to capture image from webcam") + + # Crop the frame to the specified bbox + frame = frame[y:y+height, x:x+width] + img_torch = torch.from_numpy(frame[..., [2, 1, 0]]).float() / 255.0 + + if release: + self.cap.release() + self.cap = None + + return (img_torch.unsqueeze(0),) + +class AddLabel: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image":("IMAGE",), + "text_x": ("INT", {"default": 10, "min": 0, "max": 4096, "step": 1}), + "text_y": ("INT", {"default": 2, "min": 0, "max": 4096, "step": 1}), + "height": ("INT", {"default": 48, "min": -1, "max": 4096, "step": 1}), + "font_size": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 1}), + "font_color": ("STRING", {"default": "white"}), + "label_color": ("STRING", {"default": "black"}), + "font": (folder_paths.get_filename_list("kjnodes_fonts"), ), + "text": ("STRING", {"default": "Text"}), + "direction": ( + [ 'up', + 'down', + 'left', + 'right', + 'overlay' + ], + { + "default": 'up' + }), + }, + "optional":{ + "caption": ("STRING", {"default": "", "forceInput": True}), + } + } + RETURN_TYPES = ("IMAGE",) + FUNCTION = "addlabel" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +Creates a new with the given text, and concatenates it to +either above or below the input image. +Note that this changes the input image's height! +Fonts are loaded from this folder: +ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts +""" + + def addlabel(self, image, text_x, text_y, text, height, font_size, font_color, label_color, font, direction, caption=""): + batch_size = image.shape[0] + width = image.shape[2] + + font_path = os.path.join(script_directory, "fonts", "TTNorms-Black.otf") if font == "TTNorms-Black.otf" else folder_paths.get_full_path("kjnodes_fonts", font) + + # Parse colors using helper function + font_color_rgb = string_to_color(font_color) + label_color_rgb = string_to_color(label_color) + + # Convert to tuples for PIL + font_color_tuple = tuple(font_color_rgb[:3]) # RGB only + label_color_tuple = tuple(label_color_rgb[:3]) # RGB only + + def process_image(input_image, caption_text): + font = ImageFont.truetype(font_path, font_size) + lines = [] + for text_line in caption_text.split('\n'): + if text_line.strip() == "": + # Preserve empty lines for multiple newlines + lines.append("") + continue + words = text_line.split() + current_line = [] + for word in words: + if current_line: + test_line = " ".join(current_line + [word]) + else: + test_line = word + try: + test_line_width = font.getbbox(test_line)[2] + except Exception: + test_line_width = font.getsize(test_line)[0] + if test_line_width <= width - 2 * text_x: + current_line.append(word) + else: + lines.append(" ".join(current_line)) + current_line = [word] + if current_line: + lines.append(" ".join(current_line)) + + if direction == 'overlay': + pil_image = Image.fromarray((input_image.cpu().numpy() * 255).astype(np.uint8)) + else: + if height == -1: + # Adjust the image height automatically + margin = 8 + required_height = (text_y + len(lines) * font_size) + margin # Calculate required height + pil_image = Image.new("RGB", (width, required_height), label_color_tuple) + else: + # Initialize with a minimal height + label_image = Image.new("RGB", (width, height), label_color_tuple) + pil_image = label_image + + draw = ImageDraw.Draw(pil_image) + + + y_offset = text_y + for line in lines: + try: + draw.text((text_x, y_offset), line, font=font, fill=font_color_tuple, features=['-liga']) + except Exception: + draw.text((text_x, y_offset), line, font=font, fill=font_color_tuple) + y_offset += font_size + + processed_image = torch.from_numpy(np.array(pil_image).astype(np.float32) / 255.0).unsqueeze(0) + return processed_image + + if caption == "": + processed_images = [process_image(img, text) for img in image] + else: + assert len(caption) == batch_size, f"Number of captions {(len(caption))} does not match number of images" + processed_images = [process_image(img, cap) for img, cap in zip(image, caption)] + processed_batch = torch.cat(processed_images, dim=0) + + # Combine images based on direction + if direction == 'down': + combined_images = torch.cat((image, processed_batch), dim=1) + elif direction == 'up': + combined_images = torch.cat((processed_batch, image), dim=1) + elif direction == 'left': + processed_batch = torch.rot90(processed_batch, 3, (2, 3)).permute(0, 3, 1, 2) + combined_images = torch.cat((processed_batch, image), dim=2) + elif direction == 'right': + processed_batch = torch.rot90(processed_batch, 3, (2, 3)).permute(0, 3, 1, 2) + combined_images = torch.cat((image, processed_batch), dim=2) + else: + combined_images = processed_batch + + return (combined_images,) + +class GetImageSizeAndCount: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + }} + + RETURN_TYPES = ("IMAGE","INT", "INT", "INT",) + RETURN_NAMES = ("image", "width", "height", "count",) + FUNCTION = "getsize" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Returns width, height and batch size of the image, +and passes it through unchanged. + +""" + + def getsize(self, image): + width = image.shape[2] + height = image.shape[1] + count = image.shape[0] + return {"ui": { + "text": [f"{count}x{width}x{height}"]}, + "result": (image, width, height, count) + } + +class GetLatentSizeAndCount: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "latent": ("LATENT",), + }} + + RETURN_TYPES = ("LATENT","INT", "INT", "INT", "INT", "INT") + RETURN_NAMES = ("latent", "batch_size", "channels", "frames", "height", "width") + FUNCTION = "getsize" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Returns latent tensor dimensions, +and passes the latent through unchanged. + +""" + def getsize(self, latent): + if len(latent["samples"].shape) == 5: + B, C, T, H, W = latent["samples"].shape + elif len(latent["samples"].shape) == 4: + B, C, H, W = latent["samples"].shape + T = 0 + else: + raise ValueError("Invalid latent shape") + + return {"ui": { + "text": [f"{B}x{C}x{T}x{H}x{W}"]}, + "result": (latent, B, C, T, H, W) + } + +class ImageBatchRepeatInterleaving: + RETURN_TYPES = ("IMAGE", "MASK",) + FUNCTION = "repeat" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Repeats each image in a batch by the specified number of times. +Example batch of 5 images: 0, 1 ,2, 3, 4 +with repeats 2 becomes batch of 10 images: 0, 0, 1, 1, 2, 2, 3, 3, 4, 4 +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "repeats": ("INT", {"default": 1, "min": 1, "max": 4096}), + }, + "optional": { + "mask": ("MASK",), + } + } + + def repeat(self, images, repeats, mask=None): + original_count = images.shape[0] + total_count = original_count * repeats + + repeated_images = torch.repeat_interleave(images, repeats=repeats, dim=0) + if mask is not None: + mask = torch.repeat_interleave(mask, repeats=repeats, dim=0) + else: + mask = torch.zeros((total_count, images.shape[1], images.shape[2]), + device=images.device, dtype=images.dtype) + for i in range(original_count): + mask[i * repeats] = 1.0 + + return (repeated_images, mask) + +class ImageUpscaleWithModelBatched: + @classmethod + def INPUT_TYPES(s): + return {"required": { "upscale_model": ("UPSCALE_MODEL",), + "images": ("IMAGE",), + "per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}), + }, + "optional": { + "downscale_ratio": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 1.0, "step": 0.01}), + "downscale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"], {"default": "lanczos"}), + "precision": (["float32", "float16", "bfloat16"], {"default": "float32"}), + }} + RETURN_TYPES = ("IMAGE",) + FUNCTION = "upscale" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Same as ComfyUI native model upscaling node, +but allows setting sub-batches for reduced VRAM usage. +Optionally downscale the result with a ratio. +""" + def upscale(self, upscale_model, images, per_batch, downscale_ratio=1.0, downscale_method="lanczos", precision="float32"): + dtype = torch.float16 if precision == "float16" else torch.bfloat16 if precision == "bfloat16" else torch.float32 + device = model_management.get_torch_device() + upscale_model.to(device, dtype=dtype) + in_img = images.movedim(-1,-3).to(dtype) + + steps = in_img.shape[0] + pbar = ProgressBar(steps) + t = [] + + for start_idx in range(0, in_img.shape[0], per_batch): + sub_images = upscale_model(in_img[start_idx:start_idx+per_batch].to(device)) + t.append(sub_images.cpu()) + # Calculate the number of images processed in this batch + batch_count = sub_images.shape[0] + # Update the progress bar by the number of images processed in this batch + pbar.update(batch_count) + upscale_model.cpu() + + t = torch.cat(t, dim=0).permute(0, 2, 3, 1).cpu().float() + + # Apply downscaling if ratio is less than 1.0 + if downscale_ratio < 1.0: + original_height = t.shape[1] + original_width = t.shape[2] + new_height = int(original_height * downscale_ratio) + new_width = int(original_width * downscale_ratio) + t = t.movedim(-1, 1) + t = common_upscale(t, new_width, new_height, downscale_method, "disabled") + t = t.movedim(1, -1) + + return (t,) + +class ImageNormalize_Neg1_To_1: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE",), + + }} + RETURN_TYPES = ("IMAGE",) + FUNCTION = "normalize" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Normalize the images to be in the range [-1, 1] +""" + + def normalize(self,images): + images = images * 2.0 - 1.0 + return (images,) + +class RemapImageRange: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + "min": ("FLOAT", {"default": 0.0,"min": -10.0, "max": 1.0, "step": 0.01}), + "max": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 10.0, "step": 0.01}), + "clamp": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "remap" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Remaps the image values to the specified range. +""" + + def remap(self, image, min, max, clamp): + if image.dtype == torch.float16: + image = image.to(torch.float32) + image = min + image * (max - min) + if clamp: + image = torch.clamp(image, min=0.0, max=1.0) + return (image, ) + +class SplitImageChannels: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK") + RETURN_NAMES = ("red", "green", "blue", "mask") + FUNCTION = "split" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Splits image channels into images where the selected channel +is repeated for all channels, and the alpha as a mask. +""" + + def split(self, image): + red = image[:, :, :, 0:1] # Red channel + green = image[:, :, :, 1:2] # Green channel + blue = image[:, :, :, 2:3] # Blue channel + if image.shape[3] == 4: + alpha = image[:, :, :, 4] # Alpha channel + else: + alpha = torch.zeros(image.shape[0], image.shape[1], image.shape[2], device=image.device) + + # Repeat the selected channel for all channels + red = torch.cat([red, red, red], dim=3) + green = torch.cat([green, green, green], dim=3) + blue = torch.cat([blue, blue, blue], dim=3) + return (red, green, blue, alpha) + +class MergeImageChannels: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "red": ("IMAGE",), + "green": ("IMAGE",), + "blue": ("IMAGE",), + + }, + "optional": { + "alpha": ("MASK", {"default": None}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "merge" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Merges channel data into an image. +""" + + def merge(self, red, green, blue, alpha=None): + image = torch.stack([ + red[..., 0, None], # Red channel + green[..., 1, None], # Green channel + blue[..., 2, None] # Blue channel + ], dim=-1) + image = image.squeeze(-2) + if alpha is not None: + image = torch.cat([image, alpha.unsqueeze(-1)], dim=-1) + return (image,) + +class ImagePadForOutpaintMasked: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "feathering": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + }, + "optional": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "expand_image" + + CATEGORY = "image" + + def expand_image(self, image, left, top, right, bottom, feathering, mask=None): + if mask is not None: + if torch.allclose(mask, torch.zeros_like(mask)): + logging.warning("The incoming mask is fully black. Handling it as None.") + mask = None + B, H, W, C = image.size() + + new_image = torch.ones( + (B, H + top + bottom, W + left + right, C), + dtype=torch.float32, + ) * 0.5 + + new_image[:, top:top + H, left:left + W, :] = image + + if mask is None: + new_mask = torch.ones( + (B, H + top + bottom, W + left + right), + dtype=torch.float32, + ) + + t = torch.zeros( + (B, H, W), + dtype=torch.float32 + ) + else: + # If a mask is provided, pad it to fit the new image size + mask = F.pad(mask, (left, right, top, bottom), mode='constant', value=0) + mask = 1 - mask + t = torch.zeros_like(mask) + + if feathering > 0 and feathering * 2 < H and feathering * 2 < W: + + for i in range(H): + for j in range(W): + dt = i if top != 0 else H + db = H - i if bottom != 0 else H + + dl = j if left != 0 else W + dr = W - j if right != 0 else W + + d = min(dt, db, dl, dr) + + if d >= feathering: + continue + + v = (feathering - d) / feathering + + if mask is None: + t[:, i, j] = v * v + else: + t[:, top + i, left + j] = v * v + + if mask is None: + new_mask[:, top:top + H, left:left + W] = t + return (new_image, new_mask,) + else: + return (new_image, mask,) + +class ImagePadForOutpaintTargetSize: + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "target_width": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "target_height": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "feathering": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "upscale_method": (s.upscale_methods,), + }, + "optional": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "expand_image" + + CATEGORY = "image" + + def expand_image(self, image, target_width, target_height, feathering, upscale_method, mask=None): + B, H, W, C = image.size() + new_height = H + new_width = W + # Calculate the scaling factor while maintaining aspect ratio + scaling_factor = min(target_width / W, target_height / H) + + # Check if the image needs to be downscaled + if scaling_factor < 1: + image = image.movedim(-1,1) + # Calculate the new width and height after downscaling + new_width = int(W * scaling_factor) + new_height = int(H * scaling_factor) + + # Downscale the image + image_scaled = common_upscale(image, new_width, new_height, upscale_method, "disabled").movedim(1,-1) + else: + # If downscaling is not needed, use the original image dimensions + image_scaled = image + + # Ensure mask dimensions match image dimensions + if mask is not None: + mask_scaled = mask.unsqueeze(0) # Add an extra dimension for batch size + mask_scaled = F.interpolate(mask_scaled, size=(new_height, new_width), mode="nearest") + mask_scaled = mask_scaled.squeeze(0) # Remove the extra dimension after interpolation + else: + mask_scaled = None + + # Calculate how much padding is needed to reach the target dimensions + pad_top = max(0, (target_height - new_height) // 2) + pad_bottom = max(0, target_height - new_height - pad_top) + pad_left = max(0, (target_width - new_width) // 2) + pad_right = max(0, target_width - new_width - pad_left) + + # Now call the original expand_image with the calculated padding + return ImagePadForOutpaintMasked.expand_image(self, image_scaled, pad_left, pad_top, pad_right, pad_bottom, feathering, mask_scaled) + +class ImagePrepForICLora: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "reference_image": ("IMAGE",), + "output_width": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}), + "output_height": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}), + "border_width": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + }, + "optional": { + "latent_image": ("IMAGE",), + "latent_mask": ("MASK",), + "reference_mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "expand_image" + + CATEGORY = "image" + + def expand_image(self, reference_image, output_width, output_height, border_width, latent_image=None, reference_mask=None, latent_mask=None): + + if reference_mask is not None: + if torch.allclose(reference_mask, torch.zeros_like(reference_mask)): + logging.warning("The incoming mask is fully black. Handling it as None.") + reference_mask = None + image = reference_image + if latent_image is not None: + if image.shape[0] != latent_image.shape[0]: + image = image.repeat(latent_image.shape[0], 1, 1, 1) + B, H, W, C = image.size() + + # Handle mask + if reference_mask is not None: + resized_mask = torch.nn.functional.interpolate( + reference_mask.unsqueeze(1), + size=(H, W), + mode='nearest' + ).squeeze(1) + image = image * resized_mask.unsqueeze(-1) + + # Calculate new width maintaining aspect ratio + new_width = int((W / H) * output_height) + + # Resize image to new height while maintaining aspect ratio + resized_image = common_upscale(image.movedim(-1,1), new_width, output_height, "lanczos", "disabled").movedim(1,-1) + + # Create padded image + if latent_image is None: + pad_image = torch.zeros((B, output_height, output_width, C), device=image.device) + else: + resized_latent_image = common_upscale(latent_image.movedim(-1,1), output_width, output_height, "lanczos", "disabled").movedim(1,-1) + pad_image = resized_latent_image + if latent_mask is not None: + resized_latent_mask = torch.nn.functional.interpolate( + latent_mask.unsqueeze(1), + size=(pad_image.shape[1], pad_image.shape[2]), + mode='nearest' + ).squeeze(1) + + if border_width > 0: + border = torch.zeros((B, output_height, border_width, C), device=image.device) + padded_image = torch.cat((resized_image, border, pad_image), dim=2) + if latent_mask is not None: + padded_mask = torch.zeros((B, padded_image.shape[1], padded_image.shape[2]), device=image.device) + padded_mask[:, :, (new_width + border_width):] = resized_latent_mask + else: + padded_mask = torch.ones((B, padded_image.shape[1], padded_image.shape[2]), device=image.device) + padded_mask[:, :, :new_width + border_width] = 0 + else: + padded_image = torch.cat((resized_image, pad_image), dim=2) + if latent_mask is not None: + padded_mask = torch.zeros((B, padded_image.shape[1], padded_image.shape[2]), device=image.device) + padded_mask[:, :, new_width:] = resized_latent_mask + else: + padded_mask = torch.ones((B, padded_image.shape[1], padded_image.shape[2]), device=image.device) + padded_mask[:, :, :new_width] = 0 + + return (padded_image, padded_mask) + + +class ImageAndMaskPreview(SaveImage): + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask_opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "mask_color": ("STRING", {"default": "255, 255, 255", "tooltip": "RGB (255,255,255) or RGBA (255,255,255,128) or Hex (#RRGGBB / #RRGGBBAA)"}), + "pass_through": ("BOOLEAN", {"default": False}), + }, + "optional": { + "image": ("IMAGE",), + "mask": ("MASK",), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("composite",) + FUNCTION = "execute" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Preview an image or a mask, when both inputs are used +composites the mask on top of the image. +with pass_through on the preview is disabled and the +composite is returned from the composite slot instead, +this allows for the preview to be passed for video combine +nodes for example. Supports RGBA for mask_color to adjust transparency per color. +""" + + def execute(self, mask_opacity, mask_color, pass_through, filename_prefix="ComfyUI", image=None, mask=None, prompt=None, extra_pnginfo=None): + if mask is not None and image is None: + preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + elif mask is None and image is not None: + preview = image + elif mask is not None and image is not None: + mask_adjusted = mask * mask_opacity + mask_image = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3).clone() + + # Use helper function to parse color string + color_list = string_to_color(mask_color) + + # Apply RGB channels + mask_image[:, :, :, 0] = color_list[0] / 255 # Red channel + mask_image[:, :, :, 1] = color_list[1] / 255 # Green channel + mask_image[:, :, :, 2] = color_list[2] / 255 # Blue channel + + if len(color_list) == 4: # Apply Alpha channel if present + alpha_factor = color_list[3] / 255.0 + mask_adjusted = mask_adjusted * alpha_factor + + destination, source = node_helpers.image_alpha_fix(image, mask_image) + destination = destination.clone().movedim(-1, 1) + preview = composite(destination, source.movedim(-1, 1), 0, 0, mask_adjusted, 1, True).movedim(1, -1) + if pass_through: + return (preview, ) + return(self.save_images(preview, filename_prefix, prompt, extra_pnginfo)) + +def crossfade(images_1, images_2, alpha): + crossfade = (1 - alpha) * images_1 + alpha * images_2 + return crossfade +def ease_in(t): + return t * t +def ease_out(t): + return 1 - (1 - t) * (1 - t) +def ease_in_out(t): + return 3 * t * t - 2 * t * t * t +def bounce(t): + if t < 0.5: + return ease_out(t * 2) * 0.5 + else: + return ease_in((t - 0.5) * 2) * 0.5 + 0.5 +def elastic(t): + return math.sin(13 * math.pi / 2 * t) * math.pow(2, 10 * (t - 1)) +def glitchy(t): + return t + 0.1 * math.sin(40 * t) +def exponential_ease_out(t): + return 1 - (1 - t) ** 4 + +easing_functions = { + "linear": lambda t: t, + "ease_in": ease_in, + "ease_out": ease_out, + "ease_in_out": ease_in_out, + "bounce": bounce, + "elastic": elastic, + "glitchy": glitchy, + "exponential_ease_out": exponential_ease_out, +} + +class CrossFadeImages: + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "crossfadeimages" + CATEGORY = "KJNodes/image" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images_1": ("IMAGE",), + "images_2": ("IMAGE",), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out", "bounce", "elastic", "glitchy", "exponential_ease_out"],), + "transition_start_index": ("INT", {"default": 1,"min": -4096, "max": 4096, "step": 1}), + "transitioning_frames": ("INT", {"default": 1,"min": 0, "max": 4096, "step": 1}), + "start_level": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 1.0, "step": 0.01}), + "end_level": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + def crossfadeimages(self, images_1, images_2, transition_start_index, transitioning_frames, interpolation, start_level, end_level): + + crossfade_images = [] + + if transition_start_index < 0: + transition_start_index = len(images_1) + transition_start_index + if transition_start_index < 0: + raise ValueError("Transition start index is out of range for images_1.") + + transitioning_frames = min(transitioning_frames, len(images_1) - transition_start_index, len(images_2)) + + alphas = torch.linspace(start_level, end_level, transitioning_frames) + for i in range(transitioning_frames): + alpha = alphas[i] + image1 = images_1[transition_start_index + i] + image2 = images_2[i] + easing_function = easing_functions.get(interpolation) + alpha = easing_function(alpha) # Apply the easing function to the alpha value + + crossfade_image = crossfade(image1, image2, alpha) + crossfade_images.append(crossfade_image) + + # Convert crossfade_images to tensor + crossfade_images = torch.stack(crossfade_images, dim=0) + + # Append the beginning of images_1 (before the transition) + beginning_images_1 = images_1[:transition_start_index] + crossfade_images = torch.cat([beginning_images_1, crossfade_images], dim=0) + + # Append the remaining frames of images_2 (after the transition) + remaining_images_2 = images_2[transitioning_frames:] + if len(remaining_images_2) > 0: + crossfade_images = torch.cat([crossfade_images, remaining_images_2], dim=0) + + return (crossfade_images, ) + +class CrossFadeImagesMulti: + RETURN_TYPES = ("IMAGE",) + FUNCTION = "crossfadeimages" + CATEGORY = "KJNodes/image" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "image_1": ("IMAGE",), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out", "bounce", "elastic", "glitchy", "exponential_ease_out"],), + "transitioning_frames": ("INT", {"default": 1,"min": 0, "max": 4096, "step": 1}), + }, + "optional": { + "image_2": ("IMAGE",), + } + } + + def crossfadeimages(self, inputcount, transitioning_frames, interpolation, **kwargs): + + image_1 = kwargs["image_1"] + first_image_shape = image_1.shape + first_image_device = image_1.device + height = image_1.shape[1] + width = image_1.shape[2] + + easing_function = easing_functions[interpolation] + + for c in range(1, inputcount): + frames = [] + new_image = kwargs.get(f"image_{c + 1}", torch.zeros(first_image_shape)).to(first_image_device) + new_image_height = new_image.shape[1] + new_image_width = new_image.shape[2] + + if new_image_height != height or new_image_width != width: + new_image = common_upscale(new_image.movedim(-1, 1), width, height, "lanczos", "disabled") + new_image = new_image.movedim(1, -1) # Move channels back to the last dimension + + last_frame_image_1 = image_1[-1] + first_frame_image_2 = new_image[0] + + for frame in range(transitioning_frames): + t = frame / (transitioning_frames - 1) + alpha = easing_function(t) + alpha_tensor = torch.tensor(alpha, dtype=last_frame_image_1.dtype, device=last_frame_image_1.device) + frame_image = crossfade(last_frame_image_1, first_frame_image_2, alpha_tensor) + frames.append(frame_image) + + frames = torch.stack(frames) + image_1 = torch.cat((image_1, frames, new_image), dim=0) + + return image_1, + +def transition_images(images_1, images_2, alpha, transition_type, blur_radius, reverse): + width = images_1.shape[1] + height = images_1.shape[0] + + mask = torch.zeros_like(images_1, device=images_1.device) + + alpha = alpha.item() + if reverse: + alpha = 1 - alpha + + #transitions from matteo's essential nodes + if "horizontal slide" in transition_type: + pos = round(width * alpha) + mask[:, :pos, :] = 1.0 + elif "vertical slide" in transition_type: + pos = round(height * alpha) + mask[:pos, :, :] = 1.0 + elif "box" in transition_type: + box_w = round(width * alpha) + box_h = round(height * alpha) + x1 = (width - box_w) // 2 + y1 = (height - box_h) // 2 + x2 = x1 + box_w + y2 = y1 + box_h + mask[y1:y2, x1:x2, :] = 1.0 + elif "circle" in transition_type: + radius = math.ceil(math.sqrt(pow(width, 2) + pow(height, 2)) * alpha / 2) + c_x = width // 2 + c_y = height // 2 + x = torch.arange(0, width, dtype=torch.float32, device="cpu") + y = torch.arange(0, height, dtype=torch.float32, device="cpu") + y, x = torch.meshgrid((y, x), indexing="ij") + circle = ((x - c_x) ** 2 + (y - c_y) ** 2) <= (radius ** 2) + mask[circle] = 1.0 + elif "horizontal door" in transition_type: + bar = math.ceil(height * alpha / 2) + if bar > 0: + mask[:bar, :, :] = 1.0 + mask[-bar:,:, :] = 1.0 + elif "vertical door" in transition_type: + bar = math.ceil(width * alpha / 2) + if bar > 0: + mask[:, :bar,:] = 1.0 + mask[:, -bar:,:] = 1.0 + elif "fade" in transition_type: + mask[:, :, :] = alpha + + mask = gaussian_blur(mask, blur_radius) + + return images_1 * (1 - mask) + images_2 * mask + +def gaussian_blur(mask, blur_radius): + if blur_radius > 0: + kernel_size = int(blur_radius * 2) + 1 + if kernel_size % 2 == 0: + kernel_size += 1 # Ensure kernel size is odd + sigma = blur_radius / 3 + x = torch.arange(-kernel_size // 2 + 1, kernel_size // 2 + 1, dtype=torch.float32) + x = torch.exp(-0.5 * (x / sigma) ** 2) + kernel1d = x / x.sum() + kernel2d = kernel1d[:, None] * kernel1d[None, :] + kernel2d = kernel2d.to(mask.device) + kernel2d = kernel2d.expand(mask.shape[2], 1, kernel2d.shape[0], kernel2d.shape[1]) + mask = mask.permute(2, 0, 1).unsqueeze(0) # Change to [C, H, W] and add batch dimension + mask = F.conv2d(mask, kernel2d, padding=kernel_size // 2, groups=mask.shape[1]) + mask = mask.squeeze(0).permute(1, 2, 0) # Change back to [H, W, C] + return mask + +class TransitionImagesMulti: + RETURN_TYPES = ("IMAGE",) + FUNCTION = "transition" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates transitions between images. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "image_1": ("IMAGE",), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out", "bounce", "elastic", "glitchy", "exponential_ease_out"],), + "transition_type": (["horizontal slide", "vertical slide", "box", "circle", "horizontal door", "vertical door", "fade"],), + "transitioning_frames": ("INT", {"default": 2,"min": 2, "max": 4096, "step": 1}), + "blur_radius": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 100.0, "step": 0.1}), + "reverse": ("BOOLEAN", {"default": False}), + "device": (["CPU", "GPU"], {"default": "CPU"}), + }, + "optional": { + "image_2": ("IMAGE",), + } + } + + def transition(self, inputcount, transitioning_frames, transition_type, interpolation, device, blur_radius, reverse, **kwargs): + + gpu = model_management.get_torch_device() + + image_1 = kwargs["image_1"] + height = image_1.shape[1] + width = image_1.shape[2] + first_image_shape = image_1.shape + first_image_device = image_1.device + + easing_function = easing_functions[interpolation] + + for c in range(1, inputcount): + frames = [] + new_image = kwargs.get(f"image_{c + 1}", torch.zeros(first_image_shape)).to(first_image_device) + new_image_height = new_image.shape[1] + new_image_width = new_image.shape[2] + + if new_image_height != height or new_image_width != width: + new_image = common_upscale(new_image.movedim(-1, 1), width, height, "lanczos", "disabled") + new_image = new_image.movedim(1, -1) # Move channels back to the last dimension + + last_frame_image_1 = image_1[-1] + first_frame_image_2 = new_image[0] + if device == "GPU": + last_frame_image_1 = last_frame_image_1.to(gpu) + first_frame_image_2 = first_frame_image_2.to(gpu) + + if reverse: + last_frame_image_1, first_frame_image_2 = first_frame_image_2, last_frame_image_1 + + for frame in range(transitioning_frames): + t = frame / (transitioning_frames - 1) + alpha = easing_function(t) + alpha_tensor = torch.tensor(alpha, dtype=last_frame_image_1.dtype, device=last_frame_image_1.device) + frame_image = transition_images(last_frame_image_1, first_frame_image_2, alpha_tensor, transition_type, blur_radius, reverse) + frames.append(frame_image) + + frames = torch.stack(frames).cpu() + image_1 = torch.cat((image_1, frames, new_image), dim=0) + + return image_1.cpu(), + +class TransitionImagesInBatch: + RETURN_TYPES = ("IMAGE",) + FUNCTION = "transition" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates transitions between images in a batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out", "bounce", "elastic", "glitchy", "exponential_ease_out"],), + "transition_type": (["horizontal slide", "vertical slide", "box", "circle", "horizontal door", "vertical door", "fade"],), + "transitioning_frames": ("INT", {"default": 1,"min": 0, "max": 4096, "step": 1}), + "blur_radius": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 100.0, "step": 0.1}), + "reverse": ("BOOLEAN", {"default": False}), + "device": (["CPU", "GPU"], {"default": "CPU"}), + }, + } + + #transitions from matteo's essential nodes + def transition(self, images, transitioning_frames, transition_type, interpolation, device, blur_radius, reverse): + if images.shape[0] == 1: + return images, + + gpu = model_management.get_torch_device() + + easing_function = easing_functions[interpolation] + + images_list = [] + pbar = ProgressBar(images.shape[0] - 1) + for i in range(images.shape[0] - 1): + frames = [] + image_1 = images[i] + image_2 = images[i + 1] + + if device == "GPU": + image_1 = image_1.to(gpu) + image_2 = image_2.to(gpu) + + if reverse: + image_1, image_2 = image_2, image_1 + + for frame in range(transitioning_frames): + t = frame / (transitioning_frames - 1) + alpha = easing_function(t) + alpha_tensor = torch.tensor(alpha, dtype=image_1.dtype, device=image_1.device) + frame_image = transition_images(image_1, image_2, alpha_tensor, transition_type, blur_radius, reverse) + frames.append(frame_image) + pbar.update(1) + + frames = torch.stack(frames).cpu() + images_list.append(frames) + images = torch.cat(images_list, dim=0) + + return images.cpu(), + +class ImageBatchJoinWithTransition: + RETURN_TYPES = ("IMAGE",) + FUNCTION = "transition_batches" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Transitions between two batches of images, starting at a specified index in the first batch. +During the transition, frames from both batches are blended frame-by-frame, so the video keeps playing. +""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images_1": ("IMAGE",), + "images_2": ("IMAGE",), + "start_index": ("INT", {"default": 0, "min": -10000, "max": 10000, "step": 1}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out", "bounce", "elastic", "glitchy", "exponential_ease_out"],), + "transition_type": (["horizontal slide", "vertical slide", "box", "circle", "horizontal door", "vertical door", "fade"],), + "transitioning_frames": ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1}), + "blur_radius": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), + "reverse": ("BOOLEAN", {"default": False}), + "device": (["CPU", "GPU"], {"default": "CPU"}), + }, + } + + def transition_batches(self, images_1, images_2, start_index, interpolation, transition_type, transitioning_frames, blur_radius, reverse, device): + if images_1.shape[0] == 0 or images_2.shape[0] == 0: + raise ValueError("Both input batches must have at least one image.") + + if start_index < 0: + start_index = images_1.shape[0] + start_index + if start_index < 0 or start_index > images_1.shape[0]: + raise ValueError("start_index is out of range.") + + gpu = model_management.get_torch_device() + easing_function = easing_functions[interpolation] + out_frames = [] + + # Add images from images_1 up to start_index + if start_index > 0: + out_frames.append(images_1[:start_index]) + + # Determine how many frames we can blend + max_transition = min(transitioning_frames, images_1.shape[0] - start_index, images_2.shape[0]) + + # Blend corresponding frames from both batches + for i in range(max_transition): + img1 = images_1[start_index + i] + img2 = images_2[i] + if device == "GPU": + img1 = img1.to(gpu) + img2 = img2.to(gpu) + if reverse: + img1, img2 = img2, img1 + t = i / (max_transition - 1) if max_transition > 1 else 1.0 + alpha = easing_function(t) + alpha_tensor = torch.tensor(alpha, dtype=img1.dtype, device=img1.device) + frame_image = transition_images(img1, img2, alpha_tensor, transition_type, blur_radius, reverse) + out_frames.append(frame_image.cpu().unsqueeze(0)) + + # Add remaining images from images_2 after transition + if images_2.shape[0] > max_transition: + out_frames.append(images_2[max_transition:]) + + # Concatenate all frames + out = torch.cat(out_frames, dim=0) + return (out.cpu(),) + +class ShuffleImageBatch: + RETURN_TYPES = ("IMAGE",) + FUNCTION = "shuffle" + CATEGORY = "KJNodes/image" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), + }, + } + + def shuffle(self, images, seed): + torch.manual_seed(seed) + B, H, W, C = images.shape + indices = torch.randperm(B) + shuffled_images = images[indices] + + return shuffled_images, + +class GetImageRangeFromBatch: + + RETURN_TYPES = ("IMAGE", "MASK", ) + FUNCTION = "imagesfrombatch" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Returns a range of images from a batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "start_index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}), + "num_frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), + }, + "optional": { + "images": ("IMAGE",), + "masks": ("MASK",), + } + } + + def imagesfrombatch(self, start_index, num_frames, images=None, masks=None): + chosen_images = None + chosen_masks = None + + # Process images if provided + if images is not None: + if start_index == -1: + start_index = max(0, len(images) - num_frames) + if start_index < 0 or start_index >= len(images): + raise ValueError("Start index is out of range") + end_index = min(start_index + num_frames, len(images)) + chosen_images = images[start_index:end_index] + + # Process masks if provided + if masks is not None: + if start_index == -1: + start_index = max(0, len(masks) - num_frames) + if start_index < 0 or start_index >= len(masks): + raise ValueError("Start index is out of range for masks") + end_index = min(start_index + num_frames, len(masks)) + chosen_masks = masks[start_index:end_index] + + return (chosen_images, chosen_masks,) + +class RandomImageFromBatch(io.ComfyNode): + @classmethod + def define_schema(cls): + template = io.MatchType.Template("input_type", [io.Image, io.Mask]) + return io.Schema( + node_id="RandomImageFromBatch", + display_name="Random Image From Batch", + search_aliases=["random", "mask", "sequence", "frame"], + category="KJNodes/image", + description="Picks a sequence of frames from an image or mask batch within a selected index range. " + "At randomness=0 the picks are evenly spaced across the range; at randomness=1 they are " + "uniformly random without replacement; values in between blend linearly. " + "Output is always sorted by batch index. Negative indices count from the end (-1 = last).", + inputs=[ + io.MatchType.Input("input", template=template, + tooltip="Image or mask batch to sample from."), + io.Int.Input("start_index", default=0, min=-4096, max=4096, + tooltip="Inclusive start of the sampling range. Negative values count from the end."), + io.Int.Input("end_index", default=-1, min=-4096, max=4096, + tooltip="Inclusive end of the sampling range. -1 means the last frame."), + io.Int.Input("num_frames", default=1, min=1, max=4096, + tooltip="How many frames to pick from the range."), + io.Float.Input("randomness", default=1.0, min=0.0, max=1.0, step=0.01, + tooltip="0 = evenly spaced across the range, 1 = uniformly random without replacement, " + "in-between = linear blend (jittered even spacing)."), + io.Int.Input("min_distance", default=0, min=0, max=4096, + tooltip="Minimum gap (in frames) between consecutive picks. 0 = no minimum. " + "Picks are pushed forward to satisfy this; later picks may clamp to the range end."), + io.Int.Input("max_distance", default=0, min=0, max=4096, + tooltip="Maximum gap (in frames) between consecutive picks. 0 = no maximum. " + "Picks are pulled in to satisfy this, which may compress the sequence toward the start."), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, step=1, + tooltip="Random seed for reproducible sampling. Ignored when randomness is 0."), + ], + outputs=[ + io.MatchType.Output(template=template, display_name="output"), + ], + ) + + @classmethod + def execute(cls, input, start_index, end_index, num_frames, randomness, min_distance, max_distance, seed) -> io.NodeOutput: + n = input.shape[0] + if n == 0: + raise ValueError("Input batch is empty.") + + s = start_index if start_index >= 0 else n + start_index + e = end_index if end_index >= 0 else n + end_index + s = max(0, min(s, n - 1)) + e = max(0, min(e, n - 1)) + if e < s: + s, e = e, s + range_size = e - s + 1 + + if num_frames == 1: + even = [(s + e) / 2] + else: + even = [s + i * (e - s) / (num_frames - 1) for i in range(num_frames)] + + if randomness <= 0: + picks_float = even + else: + rng = random.Random(seed) + if num_frames <= range_size: + random_picks = rng.sample(range(s, e + 1), num_frames) + else: + random_picks = [rng.randint(s, e) for _ in range(num_frames)] + random_picks.sort() + picks_float = [(1 - randomness) * ev + randomness * rp for ev, rp in zip(even, random_picks)] + + picks = sorted(max(s, min(e, int(round(p)))) for p in picks_float) + + if num_frames > 1 and (min_distance > 0 or max_distance > 0): + adjusted = [picks[0]] + for i in range(1, len(picks)): + prev = adjusted[-1] + target = picks[i] + if min_distance > 0 and target - prev < min_distance: + target = prev + min_distance + if max_distance > 0 and target - prev > max_distance: + target = prev + max_distance + adjusted.append(min(e, max(s, target))) + picks = adjusted + + idx = torch.tensor(picks, dtype=torch.long, device=input.device) + chosen = input.index_select(0, idx) + + return io.NodeOutput(chosen) + +class ImageBatchExtendWithOverlap: + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", ) + RETURN_NAMES = ("source_images", "start_images", "extended_images") + OUTPUT_TOOLTIPS = ( + "The original source images (passthrough)", + "The input images used as the starting point for extension", + "The extended images with overlap, if no new images are provided this will be empty", + ) + FUNCTION = "imagesfrombatch" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Helper node for video generation extension +First input source and overlap amount to get the starting frames for the extension. +Then on another copy of the node provide the newly generated frames and choose how to overlap them. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_images": ("IMAGE", {"tooltip": "The source images to extend"}), + "overlap": ("INT", {"default": 13,"min": 1, "max": 4096, "step": 1, "tooltip": "Number of overlapping frames between source and new images"}), + "overlap_side": (["source", "new_images"], {"default": "source", "tooltip": "Which side to overlap on"}), + "overlap_mode": (["cut", "linear_blend", "ease_in_out", "filmic_crossfade", "perceptual_crossfade"], {"default": "linear_blend", "tooltip": "Method to use for overlapping frames"}), + }, + "optional": { + "new_images": ("IMAGE", {"tooltip": "The new images to extend with"}), + } + } + + def imagesfrombatch(self, source_images, overlap, overlap_side, overlap_mode, new_images=None): + if overlap > len(source_images): + return source_images, source_images, source_images + + if new_images is not None: + if source_images.shape[1:3] != new_images.shape[1:3]: + raise ValueError(f"Source and new images must have the same shape: {source_images.shape[1:3]} vs {new_images.shape[1:3]}") + # Determine where to place the overlap + prefix = source_images[:-overlap] + if overlap_side == "source": + blend_src = source_images[-overlap:] + blend_dst = new_images[:overlap] + elif overlap_side == "new_images": + blend_src = new_images[:overlap] + blend_dst = source_images[-overlap:] + suffix = new_images[overlap:] + + if overlap_mode == "linear_blend": + # Vectorized version - process all frames at once + alpha = torch.linspace(0, 1, overlap + 2, device=blend_src.device, dtype=blend_src.dtype)[1:-1] + alpha = alpha.view(-1, 1, 1, 1) # Shape: [overlap, 1, 1, 1] + blended_images = (1 - alpha) * blend_src + alpha * blend_dst + extended_images = torch.cat((prefix, blended_images, suffix), dim=0) + + elif overlap_mode == "filmic_crossfade": + gamma = 2.2 + alpha = torch.linspace(0, 1, overlap + 2, device=blend_src.device, dtype=blend_src.dtype)[1:-1] + alpha = alpha.view(-1, 1, 1, 1) + linear_src = torch.pow(blend_src, gamma) + linear_dst = torch.pow(blend_dst, gamma) + blended = (1 - alpha) * linear_src + alpha * linear_dst + blended_images = torch.pow(blended, 1.0 / gamma) + extended_images = torch.cat((prefix, blended_images, suffix), dim=0) + + elif overlap_mode == "perceptual_crossfade": + import kornia + alpha = torch.linspace(0, 1, overlap + 2, device=blend_src.device, dtype=blend_src.dtype)[1:-1] + + src_nchw = blend_src.movedim(-1, 1) + dst_nchw = blend_dst.movedim(-1, 1) + lab_src = kornia.color.rgb_to_lab(src_nchw) + lab_dst = kornia.color.rgb_to_lab(dst_nchw) + + # Blend in LAB space + alpha = alpha.view(-1, 1, 1, 1) # [N,1,1,1] for broadcasting + blended_lab = (1 - alpha) * lab_src + alpha * lab_dst + + # Convert back to RGB and reshape + blended_rgb = kornia.color.lab_to_rgb(blended_lab) + blended_images = blended_rgb.movedim(1, -1) # [N,C,H,W] -> [N,H,W,C] + extended_images = torch.cat((prefix, blended_images, suffix), dim=0) + + elif overlap_mode == "ease_in_out": + # Vectorized ease_in_out + t = torch.linspace(0, 1, overlap + 2, device=blend_src.device, dtype=blend_src.dtype)[1:-1] + eased_t = 3 * t * t - 2 * t * t * t # ease_in_out formula + eased_t = eased_t.view(-1, 1, 1, 1) + blended_images = (1 - eased_t) * blend_src + eased_t * blend_dst + extended_images = torch.cat((prefix, blended_images, suffix), dim=0) + + elif overlap_mode == "cut": + extended_images = torch.cat((prefix, suffix), dim=0) + if overlap_side == "new_images": + extended_images = torch.cat((source_images, new_images[overlap:]), dim=0) + elif overlap_side == "source": + extended_images = torch.cat((source_images[:-overlap], new_images), dim=0) + else: + extended_images = torch.zeros((1, 64, 64, 3), device="cpu") + + start_images = source_images[-overlap:] + + return (source_images, start_images, extended_images) + +class GetLatentRangeFromBatch: + + RETURN_TYPES = ("LATENT", ) + FUNCTION = "latentsfrombatch" + CATEGORY = "KJNodes/latents" + DESCRIPTION = """ +Returns a range of latents from a batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "latents": ("LATENT",), + "start_index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}), + "num_frames": ("INT", {"default": 1,"min": -1, "max": 4096, "step": 1}), + }, + } + + def latentsfrombatch(self, latents, start_index, num_frames): + chosen_latents = None + samples = latents["samples"] + if len(samples.shape) == 4: + B, C, H, W = samples.shape + num_latents = B + elif len(samples.shape) == 5: + B, C, T, H, W = samples.shape + num_latents = T + + if start_index == -1: + start_index = max(0, num_latents - num_frames) + if start_index < 0 or start_index >= num_latents: + raise ValueError("Start index is out of range") + + end_index = num_latents if num_frames == -1 else min(start_index + num_frames, num_latents) + + if len(samples.shape) == 4: + chosen_latents = samples[start_index:end_index] + elif len(samples.shape) == 5: + chosen_latents = samples[:, :, start_index:end_index] + + return ({"samples": chosen_latents.contiguous(),},) + +class InsertLatentToIndex: + + RETURN_TYPES = ("LATENT", ) + FUNCTION = "insert" + CATEGORY = "KJNodes/latents" + DESCRIPTION = """ +Inserts a latent at the specified index into the original latent batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source": ("LATENT",), + "destination": ("LATENT",), + "index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}), + }, + } + + def insert(self, source, destination, index): + samples_destination = destination["samples"] + samples_source = source["samples"].to(samples_destination) + + if len(samples_source.shape) == 4: + B, C, H, W = samples_source.shape + num_latents = B + elif len(samples_source.shape) == 5: + B, C, T, H, W = samples_source.shape + num_latents = T + + if index >= num_latents or index < 0: + raise ValueError(f"Index {index} out of bounds for tensor with {num_latents} latents") + + if len(samples_source.shape) == 4: + joined_latents = torch.cat([ + samples_destination[:index], + samples_source, + samples_destination[index+1:] + ], dim=0) + else: + joined_latents = torch.cat([ + samples_destination[:, :, :index], + samples_source, + samples_destination[:, :, index+1:] + ], dim=2) + + return ({"samples": joined_latents,},) + +class ImageBatchFilter: + + RETURN_TYPES = ("IMAGE", "STRING",) + RETURN_NAMES = ("images", "removed_indices",) + FUNCTION = "filter" + CATEGORY = "KJNodes/image" + DESCRIPTION = "Removes empty images from a batch" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "empty_color": ("STRING", {"default": "0, 0, 0"}), + "empty_threshold": ("FLOAT", {"default": 0.01,"min": 0.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "replacement_image": ("IMAGE",), + } + } + + def filter(self, images, empty_color, empty_threshold, replacement_image=None): + B, H, W, C = images.shape + + input_images = images.clone() + + empty_color_list = [int(color.strip()) for color in empty_color.split(',')] + empty_color_tensor = torch.tensor(empty_color_list, dtype=torch.float32).to(input_images.device) + + color_diff = torch.abs(input_images - empty_color_tensor) + mean_diff = color_diff.mean(dim=(1, 2, 3)) + + empty_indices = mean_diff <= empty_threshold + empty_indices_string = ', '.join([str(i) for i in range(B) if empty_indices[i]]) + + if replacement_image is not None: + B_rep, H_rep, W_rep, C_rep = replacement_image.shape + replacement = replacement_image.clone() + if (H_rep != images.shape[1]) or (W_rep != images.shape[2]) or (C_rep != images.shape[3]): + replacement = common_upscale(replacement.movedim(-1, 1), W, H, "lanczos", "center").movedim(1, -1) + input_images[empty_indices] = replacement[0] + + return (input_images, empty_indices_string,) + else: + non_empty_images = input_images[~empty_indices] + return (non_empty_images, empty_indices_string,) + +class GetImagesFromBatchIndexed: + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "indexedimagesfrombatch" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Selects and returns the images at the specified indices as an image batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "indexes": ("STRING", {"default": "0, 1, 2", "multiline": True}), + }, + } + + def indexedimagesfrombatch(self, images, indexes): + + # Parse the indexes string into a list of integers + index_list = [int(index.strip()) for index in indexes.split(',')] + + # Convert list of indices to a PyTorch tensor + indices_tensor = torch.tensor(index_list, dtype=torch.long) + + # Select the images at the specified indices + chosen_images = images[indices_tensor] + + return (chosen_images,) + +class InsertImagesToBatchIndexed: + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "insertimagesfrombatch" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Inserts images at the specified indices into the original image batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "original_images": ("IMAGE",), + "images_to_insert": ("IMAGE",), + "indexes": ("STRING", {"default": "0, 1, 2", "multiline": True}), + }, + "optional": { + "mode": (["replace", "insert"],), + } + } + + def insertimagesfrombatch(self, original_images, images_to_insert, indexes, mode="replace"): + if indexes == "": + return (original_images,) + + input_images = original_images.clone() + + # Parse the indexes string into a list of integers + index_list = [int(index.strip()) for index in indexes.split(',')] + + # Convert list of indices to a PyTorch tensor + indices_tensor = torch.tensor(index_list, dtype=torch.long) + + # Ensure the images_to_insert is a tensor + if not isinstance(images_to_insert, torch.Tensor): + images_to_insert = torch.tensor(images_to_insert) + + if mode == "replace": + # Replace the images at the specified indices + for index, image in zip(indices_tensor, images_to_insert): + input_images[index] = image + else: + # Create a list to hold the new image sequence + new_images = [] + insert_offset = 0 + + for i in range(len(input_images) + len(indices_tensor)): + if insert_offset < len(indices_tensor) and i == indices_tensor[insert_offset]: + # Use modulo to cycle through images_to_insert + new_images.append(images_to_insert[insert_offset % len(images_to_insert)]) + insert_offset += 1 + else: + new_images.append(input_images[i - insert_offset]) + + # Convert the list back to a tensor + input_images = torch.stack(new_images, dim=0) + + return (input_images,) + +class PadImageBatchInterleaved: + + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("images", "masks",) + FUNCTION = "pad" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Inserts empty frames between the images in a batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "empty_frames_per_image": ("INT", {"default": 1,"min": 0, "max": 4096, "step": 1}), + "pad_frame_value": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 1.0, "step": 0.01}), + "add_after_last": ("BOOLEAN", {"default": False}), + }, + } + + def pad(self, images, empty_frames_per_image, pad_frame_value, add_after_last): + B, H, W, C = images.shape + + # Handle single frame case specifically + if B == 1: + total_frames = 1 + empty_frames_per_image if add_after_last else 1 + else: + # Original B images + (B-1) sets of empty frames between them + total_frames = B + (B-1) * empty_frames_per_image + # Add additional empty frames after the last image if requested + if add_after_last: + total_frames += empty_frames_per_image + + # Create new tensor with zeros (empty frames) + padded_batch = torch.ones((total_frames, H, W, C), + dtype=images.dtype, + device=images.device) * pad_frame_value + # Create mask tensor (1 for original frames, 0 for empty frames) + mask = torch.zeros((total_frames, H, W), + dtype=images.dtype, + device=images.device) + + # Fill in original images at their new positions + for i in range(B): + if B == 1: + # For single frame, just place it at the beginning + new_pos = 0 + else: + # Each image is separated by empty_frames_per_image blank frames + new_pos = i * (empty_frames_per_image + 1) + + padded_batch[new_pos] = images[i] + mask[new_pos] = 1.0 # Mark this as an original frame + + return (padded_batch, mask) + +class ReplaceImagesInBatch: + + RETURN_TYPES = ("IMAGE", "MASK",) + FUNCTION = "replace" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Replaces the images in a batch, starting from the specified start index, +with the replacement images. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "start_index": ("INT", {"default": 1,"min": 0, "max": 4096, "step": 1}), + }, + "optional": { + "original_images": ("IMAGE",), + "replacement_images": ("IMAGE",), + "original_masks": ("MASK",), + "replacement_masks": ("MASK",), + } + } + + def replace(self, original_images=None, replacement_images=None, start_index=1, original_masks=None, replacement_masks=None): + images = None + masks = None + + if original_images is not None and replacement_images is not None: + if start_index >= len(original_images): + raise ValueError("ReplaceImagesInBatch: Start index is out of range") + end_index = start_index + len(replacement_images) + if end_index > len(original_images): + raise ValueError("ReplaceImagesInBatch: End index is out of range") + + original_images_copy = original_images.clone() + if original_images_copy.shape[2] != replacement_images.shape[2] or original_images_copy.shape[3] != replacement_images.shape[3]: + replacement_images = common_upscale(replacement_images.movedim(-1, 1), original_images_copy.shape[1], original_images_copy.shape[2], "lanczos", "center").movedim(1, -1) + + original_images_copy[start_index:end_index] = replacement_images + images = original_images_copy + else: + images = torch.zeros((1, 64, 64, 3)) + + if original_masks is not None and replacement_masks is not None: + if start_index >= len(original_masks): + raise ValueError("ReplaceImagesInBatch: Start index is out of range") + end_index = start_index + len(replacement_masks) + if end_index > len(original_masks): + raise ValueError("ReplaceImagesInBatch: End index is out of range") + + original_masks_copy = original_masks.clone() + if original_masks_copy.shape[1] != replacement_masks.shape[1] or original_masks_copy.shape[2] != replacement_masks.shape[2]: + replacement_masks = common_upscale(replacement_masks.unsqueeze(1), original_masks_copy.shape[1], original_masks_copy.shape[2], "nearest-exact", "center").squeeze(0) + + original_masks_copy[start_index:end_index] = replacement_masks + masks = original_masks_copy + else: + masks = torch.zeros((1, 64, 64)) + + return (images, masks) + + +class ReverseImageBatch: + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "reverseimagebatch" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Reverses the order of the images in a batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + }, + } + + def reverseimagebatch(self, images): + reversed_images = torch.flip(images, [0]) + return (reversed_images, ) + +class ImageBatchMulti: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "image_1": ("IMAGE", ), + }, + "optional": { + "image_2": ("IMAGE", ), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "combine" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates an image batch from multiple images. +You can set how many inputs the node has, +with the **inputcount** and clicking update. +""" + + def combine(self, inputcount, **kwargs): + first = kwargs["image_1"] + h, w = first.shape[1], first.shape[2] + + # determine output shape + max_ch = first.shape[-1] + total_frames = first.shape[0] + for c in range(1, inputcount): + img = kwargs.get(f"image_{c + 1}") + if img is not None: + max_ch = max(max_ch, img.shape[-1]) + total_frames += img.shape[0] + else: + total_frames += first.shape[0] + + # pre-allocate output + out = torch.empty((total_frames, h, w, max_ch), dtype=first.dtype) + offset = 0 + + for c in range(inputcount): + img = kwargs.get(f"image_{c + 1}", torch.zeros((first.shape[0], h, w, max_ch), dtype=first.dtype)) + + if img.shape[1:3] != (h, w): + img = common_upscale(img.movedim(-1, 1), w, h, "bilinear", "center").movedim(1, -1) + + if img.shape[-1] < max_ch: + img = torch.nn.functional.pad(img, (0, max_ch - img.shape[-1]), mode='constant', value=1.0) + + n = img.shape[0] + out[offset:offset + n].copy_(img, non_blocking=True) + offset += n + del img + + return (out.cpu(),) + + +class ImageTensorList: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + }} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "append" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Creates an image list from the input images. +""" + + def append(self, image1, image2): + image_list = [] + if isinstance(image1, torch.Tensor) and isinstance(image2, torch.Tensor): + image_list = [image1, image2] + elif isinstance(image1, list) and isinstance(image2, torch.Tensor): + image_list = image1 + [image2] + elif isinstance(image1, torch.Tensor) and isinstance(image2, list): + image_list = [image1] + image2 + elif isinstance(image1, list) and isinstance(image2, list): + image_list = image1 + image2 + return image_list, + +class ImageAddMulti: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "image_1": ("IMAGE", ), + "image_2": ("IMAGE", ), + "blending": ( + [ 'add', + 'subtract', + 'multiply', + 'difference', + ], + { + "default": 'add' + }), + "blend_amount": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "add" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Add blends multiple images together. +You can set how many inputs the node has, +with the **inputcount** and clicking update. +""" + + def add(self, inputcount, blending, blend_amount, **kwargs): + image = kwargs["image_1"] + for c in range(1, inputcount): + new_image = kwargs[f"image_{c + 1}"] + if blending == "add": + image = torch.add(image * blend_amount, new_image * blend_amount) + elif blending == "subtract": + image = torch.sub(image * blend_amount, new_image * blend_amount) + elif blending == "multiply": + image = torch.mul(image * blend_amount, new_image * blend_amount) + elif blending == "difference": + image = torch.sub(image, new_image) + return (image,) + + +class ImageConcatMulti(io.ComfyNode): + @classmethod + def define_schema(cls): + # image_1 drives the output type; image_2 (and JS-added image_3+) can independently be IMAGE or MASK + type_template = io.MatchType.Template("multi_image_or_mask", allowed_types=[io.Image, io.Mask]) + return io.Schema( + node_id="ImageConcatMulti", + display_name="Image Concatenate Multi", + category="KJNodes/image", + description=( + "Creates an image from multiple images or masks.\n" + "Set the input count and click 'Update inputs' to add more slots.\n" + "The output type follows image_1; other inputs are converted to match." + ), + accept_all_inputs=True, # JS dynamically adds image_3..image_N beyond the declared inputs + inputs=[ + io.Int.Input("inputcount", default=2, min=2, max=1000, step=1), + io.MatchType.Input("image_1", template=type_template), + io.Combo.Input("direction", options=['right', 'down', 'left', 'up'], default='right'), + io.Boolean.Input("match_image_size", default=False), + io.MultiType.Input("image_2", types=[io.Image, io.Mask], optional=True), + ], + outputs=[ + io.MatchType.Output(template=type_template, display_name="output"), + ], + ) + + @classmethod + def execute(cls, inputcount, image_1, direction, match_image_size, image_2=None, **kwargs) -> io.NodeOutput: + kwargs["image_1"] = image_1 + if image_2 is not None: + kwargs["image_2"] = image_2 + image = image_1 + first_image_shape = image.shape + device = model_management.intermediate_device() + dtype = model_management.intermediate_dtype() + for c in range(1, inputcount): + key = f"image_{c + 1}" + new_image = kwargs[key] if key in kwargs else torch.zeros( + first_image_shape, dtype=dtype, device=device + ) + image = ImageConcanate.concatenate(image, new_image, direction, match_image_size, first_image_shape=first_image_shape) + return io.NodeOutput(image) + +class PreviewAnimation: + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + self.compress_level = 1 + + methods = {"default": 4, "fastest": 0, "slowest": 6} + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "fps": ("FLOAT", {"default": 8.0, "min": 0.01, "max": 1000.0, "step": 0.01}), + }, + "optional": { + "images": ("IMAGE", ), + "masks": ("MASK", ), + }, + } + + RETURN_TYPES = () + FUNCTION = "preview" + OUTPUT_NODE = True + CATEGORY = "KJNodes/image" + + def preview(self, fps, images=None, masks=None): + filename_prefix = "AnimPreview" + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + results = list() + + pil_images = [] + + if images is not None and masks is not None: + for image in images: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + pil_images.append(img) + for mask in masks: + if pil_images: + mask_np = mask.cpu().numpy() + mask_np = np.clip(mask_np * 255, 0, 255).astype(np.uint8) # Convert to values between 0 and 255 + mask_img = Image.fromarray(mask_np, mode='L') + img = pil_images.pop(0) # Remove and get the first image + img = img.convert("RGBA") # Convert base image to RGBA + + # Create a new RGBA image based on the grayscale mask + rgba_mask_img = Image.new("RGBA", img.size, (255, 255, 255, 255)) + rgba_mask_img.putalpha(mask_img) # Use the mask image as the alpha channel + + # Composite the RGBA mask onto the base image + composited_img = Image.alpha_composite(img, rgba_mask_img) + pil_images.append(composited_img) # Add the composited image back + + elif images is not None and masks is None: + for image in images: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + pil_images.append(img) + + elif masks is not None and images is None: + for mask in masks: + mask_np = 255. * mask.cpu().numpy() + mask_img = Image.fromarray(np.clip(mask_np, 0, 255).astype(np.uint8)) + pil_images.append(mask_img) + else: + logging.warning("PreviewAnimation: No images or masks provided") + return { "ui": { "images": results, "animated": (None,), "text": "empty" }} + + num_frames = len(pil_images) + + c = len(pil_images) + for i in range(0, c, num_frames): + file = f"{filename}_{counter:05}_.webp" + pil_images[i].save(os.path.join(full_output_folder, file), save_all=True, duration=int(1000.0/fps), append_images=pil_images[i + 1:i + num_frames], lossless=False, quality=50, method=0) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + animated = num_frames != 1 + return { "ui": { "images": results, "animated": (animated,), "text": [f"{num_frames}x{pil_images[0].size[0]}x{pil_images[0].size[1]}"] } } + +class ImageResizeKJ: + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "upscale_method": (s.upscale_methods,), + "keep_proportion": ("BOOLEAN", { "default": False }), + "divisible_by": ("INT", { "default": 2, "min": 0, "max": 512, "step": 1, }), + }, + "optional" : { + #"width_input": ("INT", { "forceInput": True}), + #"height_input": ("INT", { "forceInput": True}), + "get_image_size": ("IMAGE",), + "crop": (["disabled","center", 0], { "tooltip": "0 will do the default center crop, this is a workaround for the widget order changing with the new frontend, as in old workflows the value of this widget becomes 0 automatically" }), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT",) + RETURN_NAMES = ("IMAGE", "width", "height",) + FUNCTION = "resize" + CATEGORY = "KJNodes/image" + DEPRECATED = True + DESCRIPTION = """ +DEPRECATED! + +Due to ComfyUI frontend changes, this node should no longer be used, please check the +v2 of the node. This node is only kept to not completely break older workflows. + +""" + + def resize(self, image, width, height, keep_proportion, upscale_method, divisible_by, + width_input=None, height_input=None, get_image_size=None, crop="disabled"): + B, H, W, C = image.shape + + if width_input: + width = width_input + if height_input: + height = height_input + if get_image_size is not None: + _, height, width, _ = get_image_size.shape + + if keep_proportion and get_image_size is None: + # If one of the dimensions is zero, calculate it to maintain the aspect ratio + if width == 0 and height != 0: + ratio = height / H + width = round(W * ratio) + elif height == 0 and width != 0: + ratio = width / W + height = round(H * ratio) + elif width != 0 and height != 0: + # Scale based on which dimension is smaller in proportion to the desired dimensions + ratio = min(width / W, height / H) + width = round(W * ratio) + height = round(H * ratio) + else: + if width == 0: + width = W + if height == 0: + height = H + + if divisible_by > 1 and get_image_size is None: + width = width - (width % divisible_by) + height = height - (height % divisible_by) + + if crop == 0: #workaround for old workflows + crop = "center" + + image = image.movedim(-1,1) + image = common_upscale(image, width, height, upscale_method, crop) + image = image.movedim(1,-1) + + return(image, image.shape[2], image.shape[1],) + +class ImageResizeKJv2: + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos", "nvidia_rtx_vsr"] + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "upscale_method": (s.upscale_methods,), + "keep_proportion": (["stretch", "resize", "pad", "pad_edge", "pad_edge_pixel", "crop", "pillarbox_blur", "total_pixels"], { "default": False }), + "pad_color": ("STRING", { "default": "0, 0, 0", "tooltip": "Color to use for padding."}), + "crop_position": (["center", "top", "bottom", "left", "right"], { "default": "center" }), + "divisible_by": ("INT", { "default": 2, "min": 0, "max": 512, "step": 1, }), + }, + "optional" : { + "mask": ("MASK",), + "device": (["cpu", "gpu"],), + #"per_batch": ("INT", { "default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, "tooltip": "Process images in sub-batches to reduce memory usage. 0 disables sub-batching."}), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ("IMAGE", "INT", "INT", "MASK",) + RETURN_NAMES = ("IMAGE", "width", "height", "mask",) + FUNCTION = "resize" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ +Resizes the image to the specified width and height. +Size can be retrieved from the input. + +Keep proportions keeps the aspect ratio of the image, by +highest dimension. +""" + + def resize(self, image, width, height, keep_proportion, upscale_method, divisible_by, pad_color, crop_position, unique_id, device="cpu", mask=None, per_batch=64): + B, H, W, C = image.shape + + # Treat ComfyUI's 64x64 placeholder mask as no mask + if mask is not None and mask.shape[-2:] == (64, 64) and (H != 64 or W != 64): + mask = None + + # Scale mask to match image dimensions if they differ + if mask is not None and mask.shape[-2:] != (H, W): + mask = common_upscale(mask.unsqueeze(1), W, H, "bilinear", crop="disabled").squeeze(1) + + if device == "gpu": + if upscale_method == "lanczos": + raise ValueError("Lanczos is not supported on the GPU") + device = model_management.get_torch_device() + else: + device = torch.device("cpu") + + pillarbox_blur = keep_proportion == "pillarbox_blur" + + # Initialize padding variables + pad_left = pad_right = pad_top = pad_bottom = 0 + + if keep_proportion in ["resize", "total_pixels"] or keep_proportion.startswith("pad") or pillarbox_blur: + if keep_proportion == "total_pixels": + total_pixels = width * height + aspect_ratio = W / H + new_height = int(math.sqrt(total_pixels / aspect_ratio)) + new_width = int(math.sqrt(total_pixels * aspect_ratio)) + + # If one of the dimensions is zero, calculate it to maintain the aspect ratio + elif width == 0 and height == 0: + new_width = W + new_height = H + elif width == 0 and height != 0: + ratio = height / H + new_width = round(W * ratio) + new_height = height + elif height == 0 and width != 0: + ratio = width / W + new_width = width + new_height = round(H * ratio) + elif width != 0 and height != 0: + ratio = min(width / W, height / H) + new_width = round(W * ratio) + new_height = round(H * ratio) + else: + new_width = width + new_height = height + + if keep_proportion.startswith("pad") or pillarbox_blur: + # Calculate padding based on position + if crop_position == "center": + pad_left = (width - new_width) // 2 + pad_right = width - new_width - pad_left + pad_top = (height - new_height) // 2 + pad_bottom = height - new_height - pad_top + elif crop_position == "top": + pad_left = (width - new_width) // 2 + pad_right = width - new_width - pad_left + pad_top = 0 + pad_bottom = height - new_height + elif crop_position == "bottom": + pad_left = (width - new_width) // 2 + pad_right = width - new_width - pad_left + pad_top = height - new_height + pad_bottom = 0 + elif crop_position == "left": + pad_left = 0 + pad_right = width - new_width + pad_top = (height - new_height) // 2 + pad_bottom = height - new_height - pad_top + elif crop_position == "right": + pad_left = width - new_width + pad_right = 0 + pad_top = (height - new_height) // 2 + pad_bottom = height - new_height - pad_top + + width = new_width + height = new_height + else: + if width == 0: + width = W + if height == 0: + height = H + + if divisible_by > 1: + width = width - (width % divisible_by) + height = height - (height % divisible_by) + + # Preflight estimate (log-only when batching is active) + if per_batch != 0 and B > per_batch: + try: + bytes_per_elem = image.element_size() # typically 4 for float32 + est_total_bytes = B * height * width * C * bytes_per_elem + est_mb = est_total_bytes / (1024 * 1024) + msg = f"Resize v2estimated output ~{est_mb:.2f} MB; batching {per_batch}/{B}" + if unique_id and PromptServer is not None: + try: + PromptServer.instance.send_progress_text(msg, unique_id) + except Exception: + pass + else: + logging.info(f"[ImageResizeKJv2] estimated output ~{est_mb:.2f} MB; batching {per_batch}/{B}") + except Exception: + pass + + # NVIDIA RTX Video Super Resolution setup + nvvfx_sr = None + nvvfx_ctx = None + if upscale_method == "nvidia_rtx_vsr": + try: + import nvvfx + except ImportError: + raise ImportError("NVIDIA RTX Video Super Resolution is not available. Please install the nvidia-vfx library and ensure you have a compatible NVIDIA GPU.") + nvvfx_ctx = nvvfx.VideoSuperRes(nvvfx.effects.QualityLevel.ULTRA) + nvvfx_sr = nvvfx_ctx.__enter__() + nvvfx_sr.output_width = max(8, round(width / 8) * 8) + nvvfx_sr.output_height = max(8, round(height / 8) * 8) + nvvfx_sr.load() + + def _process_subbatch(in_image, in_mask, pad_left, pad_right, pad_top, pad_bottom): + # Avoid unnecessary clones; only move if needed + out_image = in_image if in_image.device == device else in_image.to(device) + out_mask = None if in_mask is None else (in_mask if in_mask.device == device else in_mask.to(device)) + + # Crop logic + if keep_proportion == "crop": + old_height = out_image.shape[-3] + old_width = out_image.shape[-2] + old_aspect = old_width / old_height + new_aspect = width / height + if old_aspect > new_aspect: + crop_w = round(old_height * new_aspect) + crop_h = old_height + else: + crop_w = old_width + crop_h = round(old_width / new_aspect) + if crop_position == "center": + x = (old_width - crop_w) // 2 + y = (old_height - crop_h) // 2 + elif crop_position == "top": + x = (old_width - crop_w) // 2 + y = 0 + elif crop_position == "bottom": + x = (old_width - crop_w) // 2 + y = old_height - crop_h + elif crop_position == "left": + x = 0 + y = (old_height - crop_h) // 2 + elif crop_position == "right": + x = old_width - crop_w + y = (old_height - crop_h) // 2 + out_image = out_image.narrow(-2, x, crop_w).narrow(-3, y, crop_h) + if out_mask is not None: + out_mask = out_mask.narrow(-1, x, crop_w).narrow(-2, y, crop_h) + + if upscale_method == "nvidia_rtx_vsr": + # Process each frame through RTX Video Super Resolution + frames_chw = out_image.movedim(-1, 1).cuda().contiguous() + upscaled_frames = [] + for j in range(frames_chw.shape[0]): + dlpack_out = nvvfx_sr.run(frames_chw[j]).image + upscaled_frames.append(torch.from_dlpack(dlpack_out).clone()) + out_image = torch.stack(upscaled_frames, dim=0).movedim(1, -1).cpu() + if out_mask is not None: + out_mask = common_upscale(out_mask.unsqueeze(1), width, height, "bilinear", crop="disabled").squeeze(1) + else: + out_image = common_upscale(out_image.movedim(-1,1), width, height, upscale_method, crop="disabled").movedim(1,-1) + if out_mask is not None: + if upscale_method == "lanczos": + out_mask = common_upscale(out_mask.unsqueeze(1).repeat(1, 3, 1, 1), width, height, upscale_method, crop="disabled").movedim(1,-1)[:, :, :, 0] + else: + out_mask = common_upscale(out_mask.unsqueeze(1), width, height, upscale_method, crop="disabled").squeeze(1) + + # Pad logic + if (keep_proportion.startswith("pad") or pillarbox_blur) and (pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0): + padded_width = width + pad_left + pad_right + padded_height = height + pad_top + pad_bottom + if divisible_by > 1: + width_remainder = padded_width % divisible_by + height_remainder = padded_height % divisible_by + if width_remainder > 0: + extra_width = divisible_by - width_remainder + pad_right += extra_width + if height_remainder > 0: + extra_height = divisible_by - height_remainder + pad_bottom += extra_height + + pad_mode = ( + "pillarbox_blur" if pillarbox_blur else + "edge" if keep_proportion == "pad_edge" else + "edge_pixel" if keep_proportion == "pad_edge_pixel" else + "color" + ) + out_image, out_mask = ImagePadKJ.pad(self, out_image, pad_left, pad_right, pad_top, pad_bottom, 0, pad_color, pad_mode, mask=out_mask) + + return out_image, out_mask + + # If batching disabled (per_batch==0) or batch fits, process whole batch + if per_batch == 0 or B <= per_batch: + out_image, out_mask = _process_subbatch(image, mask, pad_left, pad_right, pad_top, pad_bottom) + else: + chunks = [] + mask_chunks = [] if mask is not None else None + total_batches = (B + per_batch - 1) // per_batch + current_batch = 0 + for start_idx in range(0, B, per_batch): + current_batch += 1 + end_idx = min(start_idx + per_batch, B) + sub_img = image[start_idx:end_idx] + sub_mask = mask[start_idx:end_idx] if mask is not None else None + sub_out_img, sub_out_mask = _process_subbatch(sub_img, sub_mask, pad_left, pad_right, pad_top, pad_bottom) + chunks.append(sub_out_img.cpu()) + if mask is not None: + mask_chunks.append(sub_out_mask.cpu() if sub_out_mask is not None else None) + # Per-batch progress update + if unique_id and PromptServer is not None: + try: + PromptServer.instance.send_progress_text( + f"Resize v2batch {current_batch}/{total_batches} · images {end_idx}/{B}", + unique_id + ) + except Exception: + pass + else: + logging.info(f"[ImageResizeKJv2] batch {current_batch}/{total_batches} · images {end_idx}/{B}") + out_image = torch.cat(chunks, dim=0) + if mask is not None and any(m is not None for m in mask_chunks): + out_mask = torch.cat([m for m in mask_chunks if m is not None], dim=0) + else: + out_mask = None + + # Cleanup NVIDIA RTX VSR context + if nvvfx_ctx is not None: + nvvfx_ctx.__exit__(None, None, None) + + # Progress UI + if unique_id and PromptServer is not None: + try: + num_elements = out_image.numel() + element_size = out_image.element_size() + memory_size_mb = (num_elements * element_size) / (1024 * 1024) + PromptServer.instance.send_progress_text( + f"Output: {out_image.shape[0]} x {out_image.shape[2]} x {out_image.shape[1]} | {memory_size_mb:.2f}MB", + unique_id + ) + except Exception: + pass + + return (out_image.cpu(), out_image.shape[2], out_image.shape[1], out_mask.cpu() if out_mask is not None else torch.zeros(64,64, device=torch.device("cpu"), dtype=torch.float32)) + +class LoadAndResizeImage: + _color_channels = ["alpha", "red", "green", "blue"] + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f.name for f in pathlib.Path(input_dir).iterdir() if f.is_file()] + return {"required": + { + "image": (sorted(files), {"image_upload": True}), + "resize": ("BOOLEAN", { "default": False }), + "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + "repeat": ("INT", { "default": 1, "min": 1, "max": 4096, "step": 1, }), + "keep_proportion": ("BOOLEAN", { "default": False }), + "divisible_by": ("INT", { "default": 2, "min": 0, "max": 512, "step": 1, }), + "mask_channel": (s._color_channels, {"tooltip": "Channel to use for the mask output"}), + "background_color": ("STRING", { "default": "", "tooltip": "Fills the alpha channel with the specified color."}), + }, + } + + CATEGORY = "KJNodes/image" + RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING",) + RETURN_NAMES = ("image", "mask", "width", "height","image_path",) + FUNCTION = "load_image" + + def load_image(self, image, resize, width, height, repeat, keep_proportion, divisible_by, mask_channel, background_color): + image_path = folder_paths.get_annotated_filepath(image) + + img = node_helpers.pillow(Image.open, image_path) + img = ImageOps.exif_transpose(img) + + # Process the background_color using the helper function + if background_color: + color_list = string_to_color(background_color) + # Ensure we have RGBA (add alpha if only RGB) + if len(color_list) == 3: + bg_color_rgba = tuple(color_list) + (255,) + else: + bg_color_rgba = tuple(color_list) + else: + bg_color_rgba = None # No background color specified + + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ['MPO'] + + W, H = img.size + if resize: + if keep_proportion: + ratio = min(width / W, height / H) + width = round(W * ratio) + height = round(H * ratio) + else: + if width == 0: + width = W + if height == 0: + height = H + + if divisible_by > 1: + width = width - (width % divisible_by) + height = height - (height % divisible_by) + else: + width, height = W, H + + for frame in ImageSequence.Iterator(img): + frame = node_helpers.pillow(ImageOps.exif_transpose, frame) + + if frame.mode == 'I': + frame = frame.point(lambda i: i * (1 / 255)) + + if frame.mode == 'P': + frame = frame.convert("RGBA") + elif 'A' in frame.getbands(): + frame = frame.convert("RGBA") + + # Extract alpha channel if it exists + if 'A' in frame.getbands() and bg_color_rgba: + alpha_mask = np.array(frame.getchannel('A')).astype(np.float32) / 255.0 + alpha_mask = 1. - torch.from_numpy(alpha_mask) + bg_image = Image.new("RGBA", frame.size, bg_color_rgba) + # Composite the frame onto the background + frame = Image.alpha_composite(bg_image, frame) + else: + alpha_mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + + image = frame.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + if resize: + image = image.resize((width, height), Image.Resampling.BILINEAR) + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + c = mask_channel[0].upper() + if c in frame.getbands(): + if resize: + frame = frame.resize((width, height), Image.Resampling.BILINEAR) + mask = np.array(frame.getchannel(c)).astype(np.float32) / 255.0 + mask = torch.from_numpy(mask) + if c == 'A' and bg_color_rgba: + mask = alpha_mask + elif c == 'A': + mask = 1. - mask + else: + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + + output_images.append(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + if repeat > 1: + output_image = output_image.repeat(repeat, 1, 1, 1) + output_mask = output_mask.repeat(repeat, 1, 1) + + return (output_image, output_mask, width, height, image_path) + + + # @classmethod + # def IS_CHANGED(s, image, **kwargs): + # image_path = folder_paths.get_annotated_filepath(image) + # m = hashlib.sha256() + # with open(image_path, 'rb') as f: + # m.update(f.read()) + # return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, image): + if not folder_paths.exists_annotated_filepath(image): + return "Invalid image file: {}".format(image) + + return True + +class LoadImagesFromFolderKJ: + # Dictionary to store folder hashes + folder_hashes = {} + + @classmethod + def IS_CHANGED(cls, folder, **kwargs): + if folder and not os.path.isabs(folder) and args.base_directory: + folder = os.path.join(args.base_directory, folder) + if not folder or not os.path.isdir(folder): + return float("NaN") + + valid_extensions = ['.jpg', '.jpeg', '.png', '.webp', '.tga'] + include_subfolders = kwargs.get('include_subfolders', False) + + file_data = [] + if include_subfolders: + for root, _, files in os.walk(folder): + for file in files: + if any(file.lower().endswith(ext) for ext in valid_extensions): + path = os.path.join(root, file) + try: + mtime = os.path.getmtime(path) + file_data.append((path, mtime)) + except OSError: + pass + else: + for file in sorted(os.listdir(folder)): + if any(file.lower().endswith(ext) for ext in valid_extensions): + path = os.path.join(folder, file) + try: + mtime = os.path.getmtime(path) + file_data.append((path, mtime)) + except OSError: + pass + + file_data.sort() + + combined_hash = hashlib.md5() + combined_hash.update(folder.encode('utf-8')) + combined_hash.update(str(len(file_data)).encode('utf-8')) + + for path, mtime in file_data: + combined_hash.update(f"{path}:{mtime}".encode('utf-8')) + + current_hash = combined_hash.hexdigest() + + old_hash = cls.folder_hashes.get(folder) + cls.folder_hashes[folder] = current_hash + + if old_hash == current_hash: + return old_hash + + return current_hash + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "folder": ("STRING", {"default": ""}), + "width": ("INT", {"default": 1024, "min": -1, "step": 1}), + "height": ("INT", {"default": 1024, "min": -1, "step": 1}), + "keep_aspect_ratio": (["crop", "pad", "stretch",],), + }, + "optional": { + "image_load_cap": ("INT", {"default": 0, "min": 0, "step": 1}), + "start_index": ("INT", {"default": 0, "min": 0, "step": 1}), + "include_subfolders": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "INT", "STRING",) + RETURN_NAMES = ("image", "mask", "count", "image_path",) + FUNCTION = "load_images" + CATEGORY = "KJNodes/image" + DESCRIPTION = """Loads images from a folder into a batch, images are resized and loaded into a batch.""" + + def load_images(self, folder, width, height, image_load_cap, start_index, keep_aspect_ratio, include_subfolders=False): + if folder and not os.path.isabs(folder) and args.base_directory: + folder = os.path.join(args.base_directory, folder) + if not folder or not os.path.isdir(folder): + raise FileNotFoundError(f"Folder '{folder}' cannot be found.") + + valid_extensions = ['.jpg', '.jpeg', '.png', '.webp', '.tga'] + image_paths = [] + if include_subfolders: + for root, _, files in os.walk(folder): + for file in files: + if any(file.lower().endswith(ext) for ext in valid_extensions): + image_paths.append(os.path.join(root, file)) + else: + for file in sorted(os.listdir(folder)): + if any(file.lower().endswith(ext) for ext in valid_extensions): + image_paths.append(os.path.join(folder, file)) + + dir_files = sorted(image_paths) + + if len(dir_files) == 0: + raise FileNotFoundError(f"No files in directory '{folder}'.") + + # start at start_index + dir_files = dir_files[start_index:] + + images = [] + masks = [] + image_path_list = [] + + limit_images = False + if image_load_cap > 0: + limit_images = True + image_count = 0 + + pbar = ProgressBar(len(dir_files)) + + for image_path in dir_files: + if os.path.isdir(image_path): + continue + if limit_images and image_count >= image_load_cap: + break + i = Image.open(image_path) + i = ImageOps.exif_transpose(i) + + # Resize image to maximum dimensions + if width == -1 and height == -1: + width = i.size[0] + height = i.size[1] + if i.size != (width, height): + i = self.resize_with_aspect_ratio(i, width, height, keep_aspect_ratio) + + + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + if mask.shape != (height, width): + mask = torch.nn.functional.interpolate(mask.unsqueeze(0).unsqueeze(0), + size=(height, width), + mode='bilinear', + align_corners=False).squeeze() + else: + mask = torch.zeros((height, width), dtype=torch.float32, device="cpu") + + images.append(image) + masks.append(mask) + image_path_list.append(image_path) + image_count += 1 + pbar.update(1) + + if len(images) == 1: + return (images[0], masks[0], 1, image_path_list) + + elif len(images) > 1: + image1 = images[0] + mask1 = masks[0].unsqueeze(0) + + for image2 in images[1:]: + image1 = torch.cat((image1, image2), dim=0) + + for mask2 in masks[1:]: + mask1 = torch.cat((mask1, mask2.unsqueeze(0)), dim=0) + + return (image1, mask1, len(images), image_path_list) + def resize_with_aspect_ratio(self, img, width, height, mode): + if mode == "stretch": + return img.resize((width, height), Image.Resampling.LANCZOS) + + img_width, img_height = img.size + aspect_ratio = img_width / img_height + target_ratio = width / height + + if mode == "crop": + # Calculate dimensions for center crop + if aspect_ratio > target_ratio: + # Image is wider - crop width + new_width = int(height * aspect_ratio) + img = img.resize((new_width, height), Image.Resampling.LANCZOS) + left = (new_width - width) // 2 + return img.crop((left, 0, left + width, height)) + else: + # Image is taller - crop height + new_height = int(width / aspect_ratio) + img = img.resize((width, new_height), Image.Resampling.LANCZOS) + top = (new_height - height) // 2 + return img.crop((0, top, width, top + height)) + + elif mode == "pad": + pad_color = self.get_edge_color(img) + # Calculate dimensions for padding + if aspect_ratio > target_ratio: + # Image is wider - pad height + new_height = int(width / aspect_ratio) + img = img.resize((width, new_height), Image.Resampling.LANCZOS) + padding = (height - new_height) // 2 + padded = Image.new('RGBA', (width, height), pad_color) + padded.paste(img, (0, padding)) + return padded + else: + # Image is taller - pad width + new_width = int(height * aspect_ratio) + img = img.resize((new_width, height), Image.Resampling.LANCZOS) + padding = (width - new_width) // 2 + padded = Image.new('RGBA', (width, height), pad_color) + padded.paste(img, (padding, 0)) + return padded + def get_edge_color(self, img): + """Sample edges and return dominant color""" + width, height = img.size + img = img.convert('RGBA') + + # Create 1-pixel high/wide images from edges + top = img.crop((0, 0, width, 1)) + bottom = img.crop((0, height-1, width, height)) + left = img.crop((0, 0, 1, height)) + right = img.crop((width-1, 0, width, height)) + + # Combine edges into single image + edges = Image.new('RGBA', (width*2 + height*2, 1)) + edges.paste(top, (0, 0)) + edges.paste(bottom, (width, 0)) + edges.paste(left.resize((height, 1)), (width*2, 0)) + edges.paste(right.resize((height, 1)), (width*2 + height, 0)) + + # Get median color + stat = ImageStat.Stat(edges) + median = tuple(map(int, stat.median)) + return median + +class ImageGridtoBatch: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE", ), + "columns": ("INT", {"default": 3, "min": 1, "max": 8, "tooltip": "The number of columns in the grid."}), + "rows": ("INT", {"default": 0, "min": 1, "max": 8, "tooltip": "The number of rows in the grid. Set to 0 for automatic calculation."}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "decompose" + CATEGORY = "KJNodes/image" + DESCRIPTION = "Converts a grid of images to a batch of images." + + def decompose(self, image, columns, rows): + B, H, W, C = image.shape + + # Calculate cell width, rounding down + cell_width = W // columns + + if rows == 0: + # If rows is 0, calculate number of full rows + cell_height = H // columns + rows = H // cell_height + else: + # If rows is specified, adjust cell_height + cell_height = H // rows + + # Crop the image to fit full cells + image = image[:, :rows*cell_height, :columns*cell_width, :] + + # Reshape and permute the image to get the grid + image = image.view(B, rows, cell_height, columns, cell_width, C) + image = image.permute(0, 1, 3, 2, 4, 5).contiguous() + image = image.view(B, rows * columns, cell_height, cell_width, C) + + # Reshape to the final batch tensor + img_tensor = image.view(-1, cell_height, cell_width, C) + + return (img_tensor,) + +class SaveImageKJ: + def __init__(self): + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images to save."}), + "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}), + "output_folder": ("STRING", {"default": "output", "tooltip": "The folder to save the images to."}), + }, + "optional": { + "caption_file_extension": ("STRING", {"default": ".txt", "tooltip": "The extension for the caption file. Limited to plain-text/data formats."}), + "caption": ("STRING", {"forceInput": True, "tooltip": "string to save as .txt file"}), + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "KJNodes/image" + DESCRIPTION = "Saves the input images to your ComfyUI output directory." + + def save_images(self, images, output_folder, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, caption=None, caption_file_extension=".txt"): + filename_prefix += self.prefix_append + + if os.path.isabs(output_folder): + if not os.path.exists(output_folder): + os.makedirs(output_folder, exist_ok=True) + full_output_folder = output_folder + _, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_folder, images[0].shape[1], images[0].shape[0]) + else: + self.output_dir = folder_paths.get_output_directory() + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + + # sanitize caption extension: strip path components so it can't traverse out of the chosen folder, and allowlist to text/data formats + if caption is not None: + caption_file_extension = os.path.basename(caption_file_extension) + if caption_file_extension and not caption_file_extension.startswith("."): + caption_file_extension = "." + caption_file_extension + if caption_file_extension.lower() not in SaveStringKJ.ALLOWED_EXTENSIONS: + raise ValueError(f"Disallowed caption extension '{caption_file_extension}'. Allowed: {', '.join(SaveStringKJ.ALLOWED_EXTENSIONS)}") + + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + if not args.disable_metadata: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + base_file_name = f"{filename_with_batch_num}_{counter:05}_" + file = f"{base_file_name}.png" + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + if caption is not None: + txt_file = base_file_name + caption_file_extension + file_path = os.path.join(full_output_folder, txt_file) + + if os.path.commonpath((os.path.abspath(full_output_folder), os.path.abspath(file_path))) != os.path.abspath(full_output_folder): + raise ValueError(f"Refusing to write caption outside the target folder: {file_path}") + with open(file_path, "w", encoding="utf-8") as f: + f.write(caption) + + counter += 1 + + return file, + +class SaveStringKJ: + ALLOWED_EXTENSIONS = [".txt", ".caption", ".json", ".yaml", ".yml", ".md", ".csv", ".tsv", ".xml", ".log", ".ini", ".toml"] + + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "string": ("STRING", {"forceInput": True, "tooltip": "string to save as .txt file"}), + "filename_prefix": ("STRING", {"default": "text", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}), + "output_folder": ("STRING", {"default": "output", "tooltip": "Subfolder within the ComfyUI output directory to save to. Paths resolving outside the output directory are rejected."}), + }, + "optional": { + "file_extension": ("STRING", {"default": ".txt", "tooltip": "The extension for the saved file. Limited to plain-text/data formats."}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + FUNCTION = "save_string" + + OUTPUT_NODE = True + + CATEGORY = "KJNodes/misc" + DESCRIPTION = "Saves the input string to your ComfyUI output directory." + + def save_string(self, string, output_folder, filename_prefix="text", file_extension=".txt"): + filename_prefix += self.prefix_append + + output_dir = os.path.abspath(self.output_dir) + if output_folder and output_folder != "output": + sub = os.path.splitdrive(output_folder)[1].replace("\\", "/").lstrip("/") + target_dir = os.path.abspath(os.path.join(output_dir, sub)) + else: + target_dir = output_dir + + try: + inside = os.path.commonpath((output_dir, target_dir)) == output_dir + except ValueError: + inside = False + if not inside: + raise ValueError(f"output_folder must resolve within the ComfyUI output directory: {target_dir}") + os.makedirs(target_dir, exist_ok=True) + + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, target_dir) + + file_extension = os.path.basename(file_extension) + if file_extension and not file_extension.startswith("."): + file_extension = "." + file_extension + + if file_extension.lower() not in self.ALLOWED_EXTENSIONS: + raise ValueError(f"Disallowed file extension '{file_extension}'. Allowed: {', '.join(self.ALLOWED_EXTENSIONS)}") + + base_file_name = f"{filename_prefix}_{counter:05}_" + + txt_file = base_file_name + file_extension + file_path = os.path.join(full_output_folder, txt_file) + while os.path.exists(file_path): + counter += 1 + base_file_name = f"{filename_prefix}_{counter:05}_" + txt_file = base_file_name + file_extension + file_path = os.path.join(full_output_folder, txt_file) + + if os.path.commonpath((os.path.abspath(full_output_folder), os.path.abspath(file_path))) != os.path.abspath(full_output_folder): + raise ValueError(f"Refusing to write outside the target folder: {file_path}") + with open(file_path, 'w', encoding="utf-8") as f: + f.write(string) + + return file_path, + +class FastPreview: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", ), + "format": (["JPEG", "PNG"], {"default": "JPEG"}), + "max_size": ("INT", {"default": 768, "min": 128, "max": 4096, "step": 64, + "tooltip": "Maximum width or height for the preview. Images larger than this are downscaled before encoding."}), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "prompt_id": "PROMPT_ID", + }, + } + + RETURN_TYPES = () + FUNCTION = "preview" + CATEGORY = "KJNodes/experimental" + OUTPUT_NODE = True + DESCRIPTION = "Fast image preview using binary websocket, bypassing base64/JSON overhead." + + def preview(self, image, format, max_size, unique_id=None, prompt_id=None): + arr = image[0].cpu().mul(255).clamp(0, 255).byte().numpy() + h, w = arr.shape[:2] + + if w > max_size or h > max_size: + scale = max_size / max(w, h) + new_w, new_h = int(w * scale), int(h * scale) + if HAS_CV2: + arr = cv2.resize(arr, (new_w, new_h), interpolation=cv2.INTER_LINEAR) + pil_image = Image.fromarray(arr) + else: + pil_image = Image.fromarray(arr).resize((new_w, new_h), Image.BILINEAR) + else: + pil_image = Image.fromarray(arr) + + if format == "JPEG" and pil_image.mode != "RGB": + pil_image = pil_image.convert("RGB") + + if PromptServer is not None and unique_id is not None: + server = PromptServer.instance + client_supports_metadata = False + if hasattr(BinaryEventTypes, "PREVIEW_IMAGE_WITH_METADATA"): + try: + from comfy_api import feature_flags + client_supports_metadata = feature_flags.supports_feature( + server.sockets_metadata, server.client_id, "supports_preview_metadata" + ) + except Exception: + client_supports_metadata = False + + if client_supports_metadata: + server.send_sync( + BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, + ( + (format, pil_image, None), + { + "node_id": unique_id, + "display_node_id": unique_id, + "prompt_id": prompt_id or "", + }, + ), + server.client_id, + ) + else: + server.send_sync( + BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, + (format, pil_image, None), + server.client_id, + ) + + return {"ui": {"fast_preview": [True]}, "result": ()} + + +class FastPreviewBatch(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="FastPreviewBatch", + display_name="Fast Preview Batch", + category="KJNodes/experimental", + description="Encodes an image batch as an all-I-frame H.264 MP4 thumbnail strip " + "and shows it as an interactive grid. Click a tile to enlarge with " + "prev/next browsing. Avoids materializing N PNGs.", + inputs=[ + io.MultiType.Input("input", [io.Image, io.Mask], tooltip="Image or mask batch to preview."), + io.Int.Input("max_thumb_size", default=512, min=512, max=1024, step=8, + tooltip="Detail-view (mp4) thumbnail max side. Strip thumbs for the grid are auto-capped at 256."), + io.Int.Input("crf", default=25, min=0, max=51, step=1, + tooltip="H.264 CRF. Lower = higher quality / larger file."), + io.Int.Input("max_grid_frames", default=1024, min=1, max=4096, step=1, + tooltip="If batch exceeds this, frames are stride-sampled evenly."), + ], + is_output_node=True, + ) + + @classmethod + def execute(cls, input, max_thumb_size, crf, max_grid_frames) -> io.NodeOutput: + import av + import threading + import queue as _queue + if input.ndim == 3: + images = input.reshape((-1, 1, input.shape[-2], input.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + else: + images = input + B, H, W, _ = images.shape + + if B > max_grid_frames: + idx = torch.linspace(0, B - 1, max_grid_frames).round().long().tolist() + else: + idx = list(range(B)) + total = len(idx) + + scale = min(1.0, max_thumb_size / max(H, W)) + new_w = max(2, int(round(W * scale))) + new_h = max(2, int(round(H * scale))) + # yuv420p needs even dimensions + new_w -= new_w & 1 + new_h -= new_h & 1 + + # Strip thumbs serve the grid only; cap at 256 so the tiled JPEG stays well + # under any browser image-decode limit regardless of detail-view size. + STRIP_MAX = 256 + strip_scale = min(1.0, STRIP_MAX / max(new_h, new_w)) + strip_w = max(2, int(round(new_w * strip_scale))) + strip_h = max(2, int(round(new_h * strip_scale))) + + output_dir = folder_paths.get_temp_directory() + prefix = "kj_batch_preview_" + ''.join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(6)) + full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path( + prefix, output_dir, new_w, new_h + ) + file = f"{filename}_{counter:05}_.mp4" + filepath = os.path.join(full_output_folder, file) + strip_file = f"{filename}_{counter:05}_grid.jpg" + strip_path = os.path.join(full_output_folder, strip_file) + + # Square-ish tiling for the JS grid renderer. + strip_cols = max(1, int(math.ceil(math.sqrt(total)))) + strip_rows = int(math.ceil(total / strip_cols)) + strip_arr = np.zeros((strip_rows * strip_h, strip_cols * strip_w, 3), dtype=np.uint8) + + fps = 30 + container = av.open(filepath, mode="w") + try: + stream = container.add_stream("libx264", rate=Fraction(fps, 1)) + stream.width = new_w + stream.height = new_h + stream.pix_fmt = "yuv420p" + stream.options = {"crf": str(crf), "preset": "ultrafast", "g": "1", "tune": "fastdecode"} + + chunk_size = 32 + need_resize = (new_h, new_w) != (H, W) + need_strip_resize = (strip_h, strip_w) != (new_h, new_w) + mode = 'area' if scale < 1.0 else 'bilinear' + work_device = model_management.get_torch_device() + + def _to_numpy_nhwc_u8(t): + return (t.mul(255).clamp(0, 255) + .to(dtype=torch.uint8, device='cpu') + .permute(0, 2, 3, 1).contiguous().numpy()) + + # Producer: GPU resize + transfer; consumer (this thread): PyAV encode. + # PyTorch GPU ops, host transfers, and PyAV's libx264 call all release the + # GIL, so threading actually overlaps the two stages. + frame_queue = _queue.Queue(maxsize=2) + producer_error = [None] + + def producer(): + try: + for c_start in range(0, total, chunk_size): + c_idx = idx[c_start:c_start + chunk_size] + sel = (images[c_idx, ..., :3].permute(0, 3, 1, 2).contiguous() + .to(device=work_device, non_blocking=True)) + sel_video = F.interpolate(sel, size=(new_h, new_w), mode=mode) if need_resize else sel + sel_strip = F.interpolate(sel_video, size=(strip_h, strip_w), mode='area') if need_strip_resize else sel_video + video_frames = _to_numpy_nhwc_u8(sel_video) + strip_frames = video_frames if sel_strip is sel_video else _to_numpy_nhwc_u8(sel_strip) + del sel, sel_video, sel_strip + frame_queue.put((c_start, video_frames, strip_frames)) + except Exception as e: + producer_error[0] = e + finally: + frame_queue.put(None) + + producer_thread = threading.Thread(target=producer, daemon=True) + producer_thread.start() + + pbar = ProgressBar(total) + while True: + item = frame_queue.get() + if item is None: + break + c_start, video_frames, strip_frames = item + for i in range(video_frames.shape[0]): + global_idx = c_start + i + sr = global_idx // strip_cols + sc = global_idx % strip_cols + strip_arr[sr * strip_h:(sr + 1) * strip_h, sc * strip_w:(sc + 1) * strip_w] = strip_frames[i] + frame = av.VideoFrame.from_ndarray(video_frames[i], format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + pbar.update(1) + + producer_thread.join() + if producer_error[0] is not None: + raise producer_error[0] + + for packet in stream.encode(): + container.mux(packet) + finally: + container.close() + + Image.fromarray(strip_arr).save(strip_path, quality=85) + + return io.NodeOutput(ui={"kj_batch_preview": [{ + "filename": file, + "subfolder": subfolder, + "type": "temp", + "frame_count": total, + "fps": fps, + "thumb_w": new_w, + "thumb_h": new_h, + "strip_filename": strip_file, + "strip_cols": strip_cols, + "strip_cell_w": strip_w, + "strip_cell_h": strip_h, + }]}) + + +class ImageCropByMaskAndResize: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "mask": ("MASK", ), + "base_resolution": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + "padding": ("INT", { "default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "min_crop_resolution": ("INT", { "default": 128, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + "max_crop_resolution": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK", "BBOX", ) + RETURN_NAMES = ("images", "masks", "bbox",) + FUNCTION = "crop" + CATEGORY = "KJNodes/image" + + def crop_by_mask(self, mask, padding=0, min_crop_resolution=None, max_crop_resolution=None): + """ + Calculate bounding box from mask with proper padding boundary protection + Ensures crop region never exceeds original image boundaries + """ + iy, ix = (mask == 1).nonzero(as_tuple=True) + h0, w0 = mask.shape + + # Handle empty mask + if iy.numel() == 0: + x_c = w0 / 2.0 + y_c = h0 / 2.0 + width = 0 + height = 0 + else: + x_min = ix.min().item() + x_max = ix.max().item() + y_min = iy.min().item() + y_max = iy.max().item() + width = x_max - x_min + 1 # Include boundary pixels + height = y_max - y_min + 1 + x_c = (x_min + x_max) / 2.0 + y_c = (y_min + y_max) / 2.0 + + # Apply min/max resolution constraints + if min_crop_resolution: + width = max(width, min_crop_resolution) + height = max(height, min_crop_resolution) + if max_crop_resolution: + width = min(width, max_crop_resolution) + height = min(height, max_crop_resolution) + + # Critical: Limit padding expansion to available image space + # Calculate maximum possible padding for each direction + max_padding_x = min((w0 - width) // 2, padding) + max_padding_y = min((h0 - height) // 2, padding) + + # Apply constrained padding + final_width = width + 2 * max_padding_x + final_height = height + 2 * max_padding_y + + # Ensure final dimensions don't exceed image bounds + final_width = min(final_width, w0) + final_height = min(final_height, h0) + + # Calculate top-left corner with boundary protection + # Center the crop while respecting image boundaries + x0 = max(0, min(int(x_c - final_width / 2), w0 - final_width)) + y0 = max(0, min(int(y_c - final_height / 2), h0 - final_height)) + + return (x0, y0, final_width, final_height) + + def crop(self, image, mask, base_resolution, padding=0, min_crop_resolution=128, max_crop_resolution=512): + """ + Main crop and resize function with uniform target dimensions for all batch items + """ + mask = mask.round() + image_list = [] + mask_list = [] + bbox_list = [] + + # Step 1: Calculate individual bounding boxes + bbox_params = [] + aspect_ratios = [] + for i in range(image.shape[0]): + x0, y0, w, h = self.crop_by_mask(mask[i], padding, min_crop_resolution, max_crop_resolution) + bbox_params.append((x0, y0, w, h)) + aspect_ratios.append(w / h) + + # Step 2: Calculate uniform target dimensions based on maximum aspect ratio + max_w = max([w for x0, y0, w, h in bbox_params]) + max_h = max([h for x0, y0, w, h in bbox_params]) + max_aspect_ratio = max(aspect_ratios) + + # Round up to nearest multiple of 16 for stable processing + max_w = (max_w + 15) // 16 * 16 + max_h = (max_h + 15) // 16 * 16 + + # Determine target dimensions maintaining aspect ratio + if max_aspect_ratio > 1: + target_width = base_resolution + target_height = int(base_resolution / max_aspect_ratio) + else: + target_height = base_resolution + target_width = int(base_resolution * max_aspect_ratio) + + # Ensure target dimensions are multiples of 16 + target_width = (target_width + 15) // 16 * 16 + target_height = (target_height + 15) // 16 * 16 + + # Step 3: Process each image with uniform crop size + for i in range(image.shape[0]): + orig_x0, orig_y0, orig_w, orig_h = bbox_params[i] + + # Calculate center of original bounding box + x_center = orig_x0 + orig_w / 2 + y_center = orig_y0 + orig_h / 2 + + # Define uniform crop region centered on each image's bounding box + # This ensures all crops have exactly the same dimensions + x0_new = max(0, min(int(x_center - max_w / 2), image.shape[2] - max_w)) + y0_new = max(0, min(int(y_center - max_h / 2), image.shape[1] - max_h)) + x1_new = x0_new + max_w + y1_new = y0_new + max_h + + # Extract cropped regions + cropped_image = image[i][y0_new:y1_new, x0_new:x1_new, :] + cropped_mask = mask[i][y0_new:y1_new, x0_new:x1_new] + + # Resize to exact target dimensions + # Image with lanczos interpolation + cropped_image = cropped_image.unsqueeze(0).movedim(-1, 1) + cropped_image = common_upscale(cropped_image, target_width, target_height, "lanczos", "disabled") + cropped_image = cropped_image.movedim(1, -1).squeeze(0) + + # Mask with bilinear interpolation + cropped_mask = cropped_mask.unsqueeze(0).unsqueeze(0) + cropped_mask = common_upscale(cropped_mask, target_width, target_height, 'bilinear', "disabled") + cropped_mask = cropped_mask.squeeze(0).squeeze(0) + + image_list.append(cropped_image) + mask_list.append(cropped_mask) + bbox_list.append((x0_new, y0_new, x1_new, y1_new)) + + return (torch.stack(image_list), torch.stack(mask_list), bbox_list) + +class ImageCropByMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE", ), + "mask": ("MASK", ), + }, + } + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("image", ) + FUNCTION = "crop" + CATEGORY = "KJNodes/image" + DESCRIPTION = "Crops the input images based on the provided mask." + + def crop(self, image, mask): + B, H, W, C = image.shape + mask = mask.round() + + # Find bounding box for each batch + crops = [] + + for b in range(B): + # Get coordinates of non-zero elements + rows = torch.any(mask[min(b, mask.shape[0]-1)] > 0, dim=1) + cols = torch.any(mask[min(b, mask.shape[0]-1)] > 0, dim=0) + + # Find boundaries + y_min, y_max = torch.where(rows)[0][[0, -1]] + x_min, x_max = torch.where(cols)[0][[0, -1]] + + # Crop image and mask + crop = image[b:b+1, y_min:y_max+1, x_min:x_max+1, :] + crops.append(crop) + + # Stack results back together + cropped_images = torch.cat(crops, dim=0) + + return (cropped_images, ) + + + +class ImageUncropByMask: + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "destination": ("IMAGE",), + "source": ("IMAGE",), + "mask": ("MASK",), + "bbox": ("BBOX",), + }, + } + + CATEGORY = "KJNodes/image" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "uncrop" + + def uncrop(self, destination, source, mask, bbox=None): + + output_list = [] + + B, H, W, C = destination.shape + + for i in range(source.shape[0]): + x0, y0, x1, y1 = bbox[i] + bbox_height = y1 - y0 + bbox_width = x1 - x0 + + # Resize source image to match the bounding box dimensions + #resized_source = F.interpolate(source[i].unsqueeze(0).movedim(-1, 1), size=(bbox_height, bbox_width), mode='bilinear', align_corners=False) + resized_source = common_upscale(source[i].unsqueeze(0).movedim(-1, 1), bbox_width, bbox_height, "lanczos", "disabled") + resized_source = resized_source.movedim(1, -1).squeeze(0) + + # Resize mask to match the bounding box dimensions + resized_mask = common_upscale(mask[i].unsqueeze(0).unsqueeze(0), bbox_width, bbox_height, "bilinear", "disabled") + resized_mask = resized_mask.squeeze(0).squeeze(0) + + # Calculate padding values + pad_left = x0 + pad_right = W - x1 + pad_top = y0 + pad_bottom = H - y1 + + # Pad the resized source image and mask to fit the destination dimensions + padded_source = F.pad(resized_source, pad=(0, 0, pad_left, pad_right, pad_top, pad_bottom), mode='constant', value=0) + padded_mask = F.pad(resized_mask, pad=(pad_left, pad_right, pad_top, pad_bottom), mode='constant', value=0) + + # Ensure the padded mask has the correct shape + padded_mask = padded_mask.unsqueeze(2).expand(-1, -1, destination[i].shape[2]) + # Ensure the padded source has the correct shape + padded_source = padded_source.unsqueeze(2).expand(-1, -1, -1, destination[i].shape[2]).squeeze(2) + + # Combine the destination and padded source images using the mask + result = destination[i] * (1.0 - padded_mask) + padded_source * padded_mask + + output_list.append(result) + + + return (torch.stack(output_list),) + +class ImageCropByMaskBatch: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE", ), + "masks": ("MASK", ), + "width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + "height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, }), + "padding": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1, }), + "preserve_size": ("BOOLEAN", {"default": False}), + "bg_color": ("STRING", {"default": "0, 0, 0", "tooltip": "Color as RGB values in range 0-255 or 0.0-1.0, or color name or hex code"}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", ) + RETURN_NAMES = ("images", "masks",) + FUNCTION = "crop" + CATEGORY = "KJNodes/image" + DESCRIPTION = "Crops the input images based on the provided masks." + + def crop(self, image, masks, width, height, bg_color, padding, preserve_size): + B, H, W, C = image.shape + BM, HM, WM = masks.shape + mask_count = BM + if HM != H or WM != W: + masks = F.interpolate(masks.unsqueeze(1), size=(H, W), mode='nearest-exact').squeeze(1) + output_images = [] + output_masks = [] + + # Parse background color using helper function + color_list = string_to_color(bg_color) + bg_color = [x / 255.0 for x in color_list] + + # For each mask + for i in range(mask_count): + curr_mask = masks[i] + + # Find bounds + y_indices, x_indices = torch.nonzero(curr_mask, as_tuple=True) + if len(y_indices) == 0 or len(x_indices) == 0: + continue + + # Get exact bounds with padding + min_y = max(0, y_indices.min().item() - padding) + max_y = min(H, y_indices.max().item() + 1 + padding) + min_x = max(0, x_indices.min().item() - padding) + max_x = min(W, x_indices.max().item() + 1 + padding) + + # Ensure mask has correct shape for multiplication + curr_mask = curr_mask.unsqueeze(-1).expand(-1, -1, C) + + # Crop image and mask together + cropped_img = image[0, min_y:max_y, min_x:max_x, :] + cropped_mask = curr_mask[min_y:max_y, min_x:max_x, :] + + crop_h, crop_w = cropped_img.shape[0:2] + new_w = crop_w + new_h = crop_h + + if not preserve_size or crop_w > width or crop_h > height: + scale = min(width/crop_w, height/crop_h) + new_w = int(crop_w * scale) + new_h = int(crop_h * scale) + + # Resize RGB + resized_img = common_upscale(cropped_img.permute(2,0,1).unsqueeze(0), new_w, new_h, "lanczos", "disabled").squeeze(0).permute(1,2,0) + resized_mask = torch.nn.functional.interpolate( + cropped_mask.permute(2,0,1).unsqueeze(0), + size=(new_h, new_w), + mode='nearest' + ).squeeze(0).permute(1,2,0) + else: + resized_img = cropped_img + resized_mask = cropped_mask + + # Create empty tensors + new_img = torch.zeros((height, width, 3), dtype=image.dtype) + new_mask = torch.zeros((height, width), dtype=image.dtype) + + # Pad both + pad_x = (width - new_w) // 2 + pad_y = (height - new_h) // 2 + new_img[pad_y:pad_y+new_h, pad_x:pad_x+new_w, :] = resized_img + if len(resized_mask.shape) == 3: + resized_mask = resized_mask[:,:,0] # Take first channel if 3D + new_mask[pad_y:pad_y+new_h, pad_x:pad_x+new_w] = resized_mask + + output_images.append(new_img) + output_masks.append(new_mask) + + if not output_images: + return (torch.zeros((0, height, width, 3), dtype=image.dtype),) + + out_rgb = torch.stack(output_images, dim=0) + out_masks = torch.stack(output_masks, dim=0) + + # Apply mask to RGB + mask_expanded = out_masks.unsqueeze(-1).expand(-1, -1, -1, 3) + background_color = torch.tensor(bg_color, dtype=torch.float32, device=image.device) + out_rgb = out_rgb * mask_expanded + background_color * (1 - mask_expanded) + + return (out_rgb, out_masks) + +class ImagePadKJ: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE", ), + "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "extra_padding": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, }), + "pad_mode": (["edge", "edge_pixel", "color", "pillarbox_blur"],), + "color": ("STRING", {"default": "0, 0, 0", "tooltip": "Color as RGB values in range 0-255 or 0.0-1.0, or color name or hex code"}), + }, + "optional": { + "mask": ("MASK", ), + "target_width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, "forceInput": True}), + "target_height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, "forceInput": True}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", ) + RETURN_NAMES = ("images", "masks",) + FUNCTION = "pad" + CATEGORY = "KJNodes/image" + DESCRIPTION = "Pad the input image and optionally mask with the specified padding." + + def pad(self, image, left, right, top, bottom, extra_padding, color, pad_mode, mask=None, target_width=None, target_height=None): + B, H, W, C = image.shape + # Resize masks to image dimensions if necessary + if mask is not None: + BM, HM, WM = mask.shape + if HM != H or WM != W: + mask = F.interpolate(mask.unsqueeze(1), size=(H, W), mode='nearest-exact').squeeze(1) + + # Parse background color using helper function + color_list = string_to_color(color) + bg_color = [x / 255.0 for x in color_list] + if len(bg_color) == 1: + bg_color = bg_color * 3 # Grayscale to RGB + bg_color = torch.tensor(bg_color, dtype=image.dtype, device=image.device) + + # Calculate padding sizes with extra padding + if target_width is not None and target_height is not None: + if extra_padding > 0: + image = common_upscale(image.movedim(-1, 1), W - extra_padding, H - extra_padding, "lanczos", "disabled").movedim(1, -1) + B, H, W, C = image.shape + + padded_width = target_width + padded_height = target_height + pad_left = (padded_width - W) // 2 + pad_right = padded_width - W - pad_left + pad_top = (padded_height - H) // 2 + pad_bottom = padded_height - H - pad_top + else: + pad_left = left + extra_padding + pad_right = right + extra_padding + pad_top = top + extra_padding + pad_bottom = bottom + extra_padding + + padded_width = W + pad_left + pad_right + padded_height = H + pad_top + pad_bottom + + # Pillarbox blur mode + if pad_mode == "pillarbox_blur": + def _gaussian_blur_nchw(img_nchw, sigma_px): + if sigma_px <= 0: + return img_nchw + radius = max(1, int(3.0 * float(sigma_px))) + k = 2 * radius + 1 + x = torch.arange(-radius, radius + 1, device=img_nchw.device, dtype=img_nchw.dtype) + k1 = torch.exp(-(x * x) / (2.0 * float(sigma_px) * float(sigma_px))) + k1 = k1 / k1.sum() + kx = k1.view(1, 1, 1, k) + ky = k1.view(1, 1, k, 1) + c = img_nchw.shape[1] + kx = kx.repeat(c, 1, 1, 1) + ky = ky.repeat(c, 1, 1, 1) + img_nchw = F.conv2d(img_nchw, kx, padding=(0, radius), groups=c) + img_nchw = F.conv2d(img_nchw, ky, padding=(radius, 0), groups=c) + return img_nchw + + out_image = torch.zeros((B, padded_height, padded_width, C), dtype=image.dtype, device=image.device) + for b in range(B): + scale_fill = max(padded_width / float(W), padded_height / float(H)) if (W > 0 and H > 0) else 1.0 + bg_w = max(1, int(round(W * scale_fill))) + bg_h = max(1, int(round(H * scale_fill))) + src_b = image[b].movedim(-1, 0).unsqueeze(0) + bg = common_upscale(src_b, bg_w, bg_h, "bilinear", crop="disabled") + y0 = max(0, (bg_h - padded_height) // 2) + x0 = max(0, (bg_w - padded_width) // 2) + y1 = min(bg_h, y0 + padded_height) + x1 = min(bg_w, x0 + padded_width) + bg = bg[:, :, y0:y1, x0:x1] + if bg.shape[2] != padded_height or bg.shape[3] != padded_width: + pad_h = padded_height - bg.shape[2] + pad_w = padded_width - bg.shape[3] + pad_top_fix = max(0, pad_h // 2) + pad_bottom_fix = max(0, pad_h - pad_top_fix) + pad_left_fix = max(0, pad_w // 2) + pad_right_fix = max(0, pad_w - pad_left_fix) + bg = F.pad(bg, (pad_left_fix, pad_right_fix, pad_top_fix, pad_bottom_fix), mode="replicate") + sigma = max(1.0, 0.006 * float(min(padded_height, padded_width))) + bg = _gaussian_blur_nchw(bg, sigma_px=sigma) + if C >= 3: + r, g, bch = bg[:, 0:1], bg[:, 1:2], bg[:, 2:3] + luma = 0.2126 * r + 0.7152 * g + 0.0722 * bch + gray = torch.cat([luma, luma, luma], dim=1) + desat = 0.20 + rgb = torch.cat([r, g, bch], dim=1) + rgb = rgb * (1.0 - desat) + gray * desat + bg[:, 0:3, :, :] = rgb + dim = 0.35 + bg = torch.clamp(bg * dim, 0.0, 1.0) + out_image[b] = bg.squeeze(0).movedim(0, -1) + out_image[:, pad_top:pad_top+H, pad_left:pad_left+W, :] = image + # Mask handling for pillarbox_blur + if mask is not None: + fg_mask = mask + out_masks = torch.ones((B, padded_height, padded_width), dtype=image.dtype, device=image.device) + out_masks[:, pad_top:pad_top+H, pad_left:pad_left+W] = fg_mask + else: + out_masks = torch.ones((B, padded_height, padded_width), dtype=image.dtype, device=image.device) + out_masks[:, pad_top:pad_top+H, pad_left:pad_left+W] = 0.0 + return (out_image, out_masks) + + # Standard pad logic (edge/color) + out_image = torch.zeros((B, padded_height, padded_width, C), dtype=image.dtype, device=image.device) + for b in range(B): + if pad_mode == "edge": + # Pad with edge color (mean) + top_edge = image[b, 0, :, :] + bottom_edge = image[b, H-1, :, :] + left_edge = image[b, :, 0, :] + right_edge = image[b, :, W-1, :] + out_image[b, :pad_top, :, :] = top_edge.mean(dim=0) + out_image[b, pad_top+H:, :, :] = bottom_edge.mean(dim=0) + out_image[b, :, :pad_left, :] = left_edge.mean(dim=0) + out_image[b, :, pad_left+W:, :] = right_edge.mean(dim=0) + out_image[b, pad_top:pad_top+H, pad_left:pad_left+W, :] = image[b] + elif pad_mode == "edge_pixel": + # Pad with exact edge pixel values + for y in range(pad_top): + out_image[b, y, pad_left:pad_left+W, :] = image[b, 0, :, :] + for y in range(pad_top+H, padded_height): + out_image[b, y, pad_left:pad_left+W, :] = image[b, H-1, :, :] + for x in range(pad_left): + out_image[b, pad_top:pad_top+H, x, :] = image[b, :, 0, :] + for x in range(pad_left+W, padded_width): + out_image[b, pad_top:pad_top+H, x, :] = image[b, :, W-1, :] + out_image[b, :pad_top, :pad_left, :] = image[b, 0, 0, :] + out_image[b, :pad_top, pad_left+W:, :] = image[b, 0, W-1, :] + out_image[b, pad_top+H:, :pad_left, :] = image[b, H-1, 0, :] + out_image[b, pad_top+H:, pad_left+W:, :] = image[b, H-1, W-1, :] + out_image[b, pad_top:pad_top+H, pad_left:pad_left+W, :] = image[b] + else: + # Pad with specified background color + out_image[b, :, :, :] = bg_color.unsqueeze(0).unsqueeze(0) + out_image[b, pad_top:pad_top+H, pad_left:pad_left+W, :] = image[b] + + if mask is not None: + out_masks = torch.nn.functional.pad( + mask, + (pad_left, pad_right, pad_top, pad_bottom), + mode='replicate' + ) + else: + out_masks = torch.ones((B, padded_height, padded_width), dtype=image.dtype, device=image.device) + for m in range(B): + out_masks[m, pad_top:pad_top+H, pad_left:pad_left+W] = 0.0 + + return (out_image, out_masks) + +# extends https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite +class LoadVideosFromFolder: + @classmethod + def __init__(cls): + try: + cls.vhs_nodes = importlib.import_module("ComfyUI-VideoHelperSuite.videohelpersuite") + except ImportError: + try: + cls.vhs_nodes = importlib.import_module("comfyui-videohelpersuite.videohelpersuite") + except ImportError: + # Fallback to sys.modules search for Windows compatibility + import sys + vhs_module = None + for module_name in sys.modules: + if 'videohelpersuite' in module_name and 'videohelpersuite' in sys.modules[module_name].__dict__: + vhs_module = sys.modules[module_name] + break + + if vhs_module is None: + # Try direct access to the videohelpersuite submodule + for module_name in sys.modules: + if module_name.endswith('videohelpersuite'): + vhs_module = sys.modules[module_name] + break + + if vhs_module is not None: + cls.vhs_nodes = vhs_module + else: + raise ImportError("This node requires ComfyUI-VideoHelperSuite to be installed.") + + except ImportError: + raise ImportError("This node requires ComfyUI-VideoHelperSuite to be installed.") + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "video": ("STRING", {"default": "X://insert/path/"},), + "force_rate": ("FLOAT", {"default": 0, "min": 0, "max": 60, "step": 1, "disable": 0}), + "custom_width": ("INT", {"default": 0, "min": 0, "max": 4096, 'disable': 0}), + "custom_height": ("INT", {"default": 0, "min": 0, "max": 4096, 'disable': 0}), + "frame_load_cap": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1, "disable": 0}), + "skip_first_frames": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}), + "select_every_nth": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}), + "output_type": (["batch", "grid"], {"default": "batch"}), + "grid_max_columns": ("INT", {"default": 4, "min": 1, "max": 16, "step": 1, "disable": 1}), + "add_label": ( "BOOLEAN", {"default": False} ), + }, + "hidden": { + "force_size": "STRING", + "unique_id": "UNIQUE_ID" + }, + } + + CATEGORY = "KJNodes/misc" + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("IMAGE", ) + + FUNCTION = "load_video" + + def load_video(self, output_type, grid_max_columns, add_label=False, **kwargs): + if kwargs.get('video') and not os.path.isabs(kwargs['video']) and args.base_directory: + kwargs['video'] = os.path.join(args.base_directory, kwargs['video']) + + if self.vhs_nodes is None: + raise ImportError("This node requires ComfyUI-VideoHelperSuite to be installed.") + videos_list = [] + filenames = [] + for f in sorted(os.listdir(kwargs['video'])): + if os.path.isfile(os.path.join(kwargs['video'], f)): + file_parts = f.split('.') + if len(file_parts) > 1 and (file_parts[-1].lower() in ['webm', 'mp4', 'mkv', 'gif', 'mov']): + videos_list.append(os.path.join(kwargs['video'], f)) + filenames.append(f) + + kwargs.pop('video') + loaded_videos = [] + for idx, video in enumerate(videos_list): + video_tensor = self.vhs_nodes.load_video_nodes.load_video(video=video, **kwargs)[0] + if add_label: + # Add filename label above video (without extension) + if video_tensor.dim() == 4: + _, h, w, c = video_tensor.shape + else: + h, w, c = video_tensor.shape + # Remove extension from filename + label_text = filenames[idx].rsplit('.', 1)[0] + font_size = max(16, w // 20) + try: + font = ImageFont.truetype("arial.ttf", font_size) + except OSError: + font = ImageFont.load_default() + dummy_img = Image.new("RGB", (w, 10), (0,0,0)) + draw = ImageDraw.Draw(dummy_img) + text_bbox = draw.textbbox((0,0), label_text, font=font) + extra_padding = max(12, font_size // 2) # More padding under the font + label_height = text_bbox[3] - text_bbox[1] + extra_padding + label_img = Image.new("RGB", (w, label_height), (0,0,0)) + draw = ImageDraw.Draw(label_img) + draw.text((w//2 - (text_bbox[2]-text_bbox[0])//2, 4), label_text, font=font, fill=(255,255,255)) + label_np = np.asarray(label_img).astype(np.float32) / 255.0 + label_tensor = torch.from_numpy(label_np) + if c == 1: + label_tensor = label_tensor.mean(dim=2, keepdim=True) + elif c == 4: + alpha = torch.ones((label_height, w, 1), dtype=label_tensor.dtype) + label_tensor = torch.cat([label_tensor, alpha], dim=2) + if video_tensor.dim() == 4: + label_tensor = label_tensor.unsqueeze(0).expand(video_tensor.shape[0], -1, -1, -1) + video_tensor = torch.cat([label_tensor, video_tensor], dim=1) + else: + video_tensor = torch.cat([label_tensor, video_tensor], dim=0) + loaded_videos.append(video_tensor) + if output_type == "batch": + out_tensor = torch.cat(loaded_videos) + elif output_type == "grid": + rows = (len(loaded_videos) + grid_max_columns - 1) // grid_max_columns + # Pad the last row if needed + total_slots = rows * grid_max_columns + while len(loaded_videos) < total_slots: + loaded_videos.append(torch.zeros_like(loaded_videos[0])) + # Create grid by rows + row_tensors = [] + for row_idx in range(rows): + start_idx = row_idx * grid_max_columns + end_idx = start_idx + grid_max_columns + row_videos = loaded_videos[start_idx:end_idx] + # Pad all videos in this row to the same height + heights = [v.shape[1] for v in row_videos] + max_height = max(heights) + padded_row_videos = [] + for v in row_videos: + pad_height = max_height - v.shape[1] + if pad_height > 0: + # Pad (frames, H, W, C) or (H, W, C) + if v.dim() == 4: + pad = (0,0, 0,0, 0,pad_height, 0,0) # (C,W,H,F) + v = torch.nn.functional.pad(v, (0,0,0,0,0,pad_height,0,0)) + else: + v = torch.nn.functional.pad(v, (0,0,0,0,pad_height,0)) + padded_row_videos.append(v) + row_tensor = torch.cat(padded_row_videos, dim=2) # Concatenate horizontally + row_tensors.append(row_tensor) + out_tensor = torch.cat(row_tensors, dim=1) # Concatenate rows vertically + return out_tensor, + + @classmethod + def IS_CHANGED(s, video, **kwargs): + if s.vhs_nodes is not None: + return s.vhs_nodes.utils.hash_path(video) + return None + + +class EncodeVideoComponents(io.ComfyNode): + @classmethod + def define_schema(cls): + position_options = ["center", "top", "bottom", "left", "right"] + options = [ + io.DynamicCombo.Option(key="stretch", inputs=[]), + io.DynamicCombo.Option(key="resize", inputs=[]), + io.DynamicCombo.Option(key="total_pixels", inputs=[]), + io.DynamicCombo.Option(key="crop", inputs=[ + io.Combo.Input("crop_position", options=position_options, tooltip="Position to crop from."), + ]), + io.DynamicCombo.Option(key="pad", inputs=[ + io.String.Input("pad_color", default="0, 0, 0", tooltip="Color to use for padding."), + io.Combo.Input("pad_position", options=position_options, tooltip="Position to align the image within the padded area."), + ]), + io.DynamicCombo.Option(key="pad_edge", inputs=[ + io.Combo.Input("pad_position", options=position_options, tooltip="Position to align the image within the padded area."), + ]), + io.DynamicCombo.Option(key="pad_edge_pixel", inputs=[ + io.Combo.Input("pad_position", options=position_options, tooltip="Position to align the image within the padded area."), + ]), + io.DynamicCombo.Option(key="pillarbox_blur", inputs=[ + io.Combo.Input("pad_position", options=position_options, tooltip="Position to align the image within the padded area."), + ]), + ] + return io.Schema( + node_id="EncodeVideoComponents", + search_aliases=["video to latent", "encode video", "vae encode video"], + display_name="Encode Video Components", + category="KJNodes/image", + description="Extracts video frames, resizes them, and encodes with a VAE directly, avoiding storing the full image tensor.", + inputs=[ + io.Video.Input("video", tooltip="The video to extract and encode."), + io.Vae.Input("vae", tooltip="The VAE model to use for encoding."), + io.Int.Input("width", default=768, min=0, max=16384, step=2, tooltip="Target width for the frames before encoding. 0 = original width."), + io.Int.Input("height", default=512, min=0, max=16384, step=2, tooltip="Target height for the frames before encoding. 0 = original height."), + io.Int.Input("max_frames", default=0, min=0, max=999999, step=1, tooltip="Maximum number of frames. 0 = no limit."), + io.Combo.Input("upscale_method", options=["nearest-exact", "bilinear", "area", "bicubic", "lanczos"], default="lanczos", tooltip="Interpolation method for resizing."), + io.DynamicCombo.Input( + "keep_proportion", + options=options, + display_name="Keep Proportion", + tooltip="How to handle aspect ratio mismatch when resizing.", + ), + ], + outputs=[ + io.Latent.Output(display_name="latent"), + io.Audio.Output(display_name="audio"), + io.Float.Output(display_name="fps"), + io.Int.Output(display_name="frame_count", tooltip="Number pixel space frames after any possible cropping"), + ], + ) + + @staticmethod + def _compute_resize_params(mode, position, width, height, src_w, src_h): + """Compute target resize dimensions, crop region, and padding from keep_proportion mode.""" + if width == 0: + width = src_w + if height == 0: + height = src_h + pillarbox_blur = mode == "pillarbox_blur" + pad_left = pad_right = pad_top = pad_bottom = 0 + crop_region = None # (x, y, crop_w, crop_h) or None + + if mode in ["resize", "total_pixels"] or mode.startswith("pad") or pillarbox_blur: + if mode == "total_pixels": + total_pixels = width * height + aspect_ratio = src_w / src_h + new_height = int(math.sqrt(total_pixels / aspect_ratio)) + new_width = int(math.sqrt(total_pixels * aspect_ratio)) + else: + ratio = min(width / src_w, height / src_h) + new_width = round(src_w * ratio) + new_height = round(src_h * ratio) + + if mode.startswith("pad") or pillarbox_blur: + if position == "center": + pad_left = (width - new_width) // 2 + pad_right = width - new_width - pad_left + pad_top = (height - new_height) // 2 + pad_bottom = height - new_height - pad_top + elif position == "top": + pad_left = (width - new_width) // 2 + pad_right = width - new_width - pad_left + pad_top = 0 + pad_bottom = height - new_height + elif position == "bottom": + pad_left = (width - new_width) // 2 + pad_right = width - new_width - pad_left + pad_top = height - new_height + pad_bottom = 0 + elif position == "left": + pad_left = 0 + pad_right = width - new_width + pad_top = (height - new_height) // 2 + pad_bottom = height - new_height - pad_top + elif position == "right": + pad_left = width - new_width + pad_right = 0 + pad_top = (height - new_height) // 2 + pad_bottom = height - new_height - pad_top + + width = new_width + height = new_height + + if mode == "crop": + old_aspect = src_w / src_h + new_aspect = width / height + if old_aspect > new_aspect: + crop_w = round(src_h * new_aspect) + crop_h = src_h + else: + crop_w = src_w + crop_h = round(src_w / new_aspect) + if position == "center": + x = (src_w - crop_w) // 2 + y = (src_h - crop_h) // 2 + elif position == "top": + x = (src_w - crop_w) // 2 + y = 0 + elif position == "bottom": + x = (src_w - crop_w) // 2 + y = src_h - crop_h + elif position == "left": + x = 0 + y = (src_h - crop_h) // 2 + elif position == "right": + x = src_w - crop_w + y = (src_h - crop_h) // 2 + crop_region = (x, y, crop_w, crop_h) + + return width, height, crop_region, (pad_left, pad_right, pad_top, pad_bottom) + + @classmethod + def execute(cls, video, vae, width, height, max_frames, upscale_method, keep_proportion) -> io.NodeOutput: + import av + import itertools + + mode = keep_proportion["keep_proportion"] + position = keep_proportion.get("crop_position") or keep_proportion.get("pad_position", "center") + pad_color = keep_proportion.get("pad_color", "0, 0, 0") + target_dtype = vae.vae_dtype + + # Access VideoFromFile internals for efficient per-frame decode + source = video.get_stream_source() + start_time = getattr(video, '_VideoFromFile__start_time', 0) + duration = getattr(video, '_VideoFromFile__duration', 0) + + # Get frame count for progress bar, capped by max_frames + try: + total_frames = video.get_frame_count() + except (ValueError, AttributeError): + total_frames = 0 + if max_frames > 0 and total_frames > 0: + total_frames = min(total_frames, max_frames) + pbar = ProgressBar(total_frames) if total_frames > 0 else None + + # Lanczos requires PIL (CPU-only), all other methods use torch on GPU + use_gpu = upscale_method != "lanczos" + device = model_management.get_torch_device() if use_gpu else torch.device("cpu") + + # --- Decode video frames with per-frame resize + dtype cast --- + with av.open(source, mode='r') as container: + video_stream = container.streams.video[0] + start_pts = int(start_time / video_stream.time_base) + end_pts = int((start_time + duration) / video_stream.time_base) if duration else 0 + container.seek(start_pts, stream=video_stream) + + res_w, res_h, crop_region, padding = None, None, None, (0, 0, 0, 0) + frames = [] + for frame in container.decode(video_stream): + if frame.pts < start_pts: + continue + if duration and frame.pts >= end_pts: + break + if max_frames > 0 and len(frames) >= max_frames: + break + + if res_w is None: + src_h, src_w = frame.height, frame.width + res_w, res_h, crop_region, padding = cls._compute_resize_params( + mode, position, width, height, src_w, src_h + ) + + # Decode to tensor and normalize + img = torch.from_numpy(frame.to_ndarray(format='rgb24')).to(device=device, dtype=torch.float32) / 255.0 + + # Crop if needed (before resize) + if crop_region is not None: + cx, cy, cw, ch = crop_region + img = img[cy:cy+ch, cx:cx+cw, :] + + # Resize (GPU for torch-native methods, CPU/PIL for lanczos) + img = common_upscale( + img.unsqueeze(0).movedim(-1, 1), res_w, res_h, upscale_method, crop="disabled" + ).movedim(1, -1).squeeze(0).to(dtype=target_dtype, device="cpu") + + frames.append(img) + if pbar is not None: + pbar.update(1) + + frame_rate = video_stream.average_rate if video_stream.average_rate else 1 + + s = torch.stack(frames) if frames else torch.zeros(0, height, width, 3, dtype=target_dtype) + + # Pad logic (applied on the full stack since padding modes like pillarbox_blur need all frames) + pillarbox_blur = mode == "pillarbox_blur" + pad_left, pad_right, pad_top, pad_bottom = padding + if (mode.startswith("pad") or pillarbox_blur) and (pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0): + pad_mode = ( + "pillarbox_blur" if pillarbox_blur else + "edge" if mode == "pad_edge" else + "edge_pixel" if mode == "pad_edge_pixel" else + "color" + ) + s, _ = ImagePadKJ.pad(None, s, pad_left, pad_right, pad_top, pad_bottom, 0, pad_color, pad_mode) + + # Trim frames to a count valid for the VAE's temporal compression + try: + temporal_compress = vae.downscale_ratio[0] + temporal_decompress = vae.upscale_ratio[0] + valid_frames = temporal_decompress(temporal_compress(s.shape[0])) + if valid_frames < s.shape[0]: + logging.warning(f"[EncodeVideoComponents] Trimming {s.shape[0] - valid_frames} frames ({s.shape[0]} -> {valid_frames}) to match VAE temporal compression ratio") + s = s[:valid_frames] + except (TypeError, IndexError): + pass + + t = vae.encode(s) + + # --- Extract audio in a separate pass --- + audio = None + if isinstance(source, BytesIO): + source.seek(0) + with av.open(source, mode='r') as container: + if len(container.streams.audio): + audio_stream = container.streams.audio[-1] + if start_time > 0: + audio_start_pts = int(start_time / audio_stream.time_base) + container.seek(audio_start_pts, stream=audio_stream) + audio_frames = [] + resample = av.audio.resampler.AudioResampler(format='fltp').resample + aframes = itertools.chain.from_iterable( + map(resample, container.decode(audio_stream)) + ) + has_first_frame = False + for aframe in aframes: + offset_seconds = start_time - aframe.time + to_skip = int(offset_seconds * audio_stream.sample_rate) + if to_skip < aframe.samples: + has_first_frame = True + break + if has_first_frame: + audio_frames.append(aframe.to_ndarray()[..., to_skip:]) + for aframe in aframes: + if duration and aframe.time > start_time + duration: + break + audio_frames.append(aframe.to_ndarray()) + if audio_frames: + audio_data = np.concatenate(audio_frames, axis=1) + if duration: + audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)] + audio = { + "waveform": torch.from_numpy(audio_data).unsqueeze(0), + "sample_rate": int(audio_stream.sample_rate) if audio_stream.sample_rate else 1, + } + + return io.NodeOutput({"samples": t}, audio, float(frame_rate), s.shape[0]) + + +class DecodeAndSaveVideo(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="DecodeAndSaveVideo", + search_aliases=["video to latent", "decode video"], + display_name="Decode and Save Video", + category="KJNodes/image", + description="Decodes video frames and audio from latent representations, combines them, and saves as a video file, without keeping intermediate images in memory.", + inputs=[ + io.Latent.Input("video_latent", tooltip="The latent representation of the video frames."), + io.Latent.Input("audio_latent", optional=True, tooltip="The latent representation of the audio frames."), + io.Float.Input("fps", default=25.0, min=0.0, max=999.0, step=0.01, tooltip="Frame rate for the output video."), + io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."), + io.Combo.Input("format", options=Types.VideoContainer.as_input(), default="auto", tooltip="The format to save the video as."), + io.Combo.Input("codec", options=Types.VideoCodec.as_input(), default="auto", tooltip="The codec to use for the video."), + io.Vae.Input("video_vae", tooltip="The VAE model to use for encoding."), + io.Vae.Input("audio_vae", optional=True, tooltip="The VAE model to use for decoding audio."), + io.DynamicCombo.Input("tiling", options=[ + io.DynamicCombo.Option(key="disabled", inputs=[]), + io.DynamicCombo.Option(key="enabled", inputs=[ + io.Int.Input("tile_size", default=512, min=64, max=4096, step=32, tooltip="Size of the tiles to decode. Smaller tiles use less memory but take more time."), + io.Int.Input("overlap", default=64, min=0, max=4096, step=32, tooltip="Amount of overlap between tiles. Higher overlap can improve quality at the edges of tiles but uses more memory and takes more time."), + io.Int.Input("temporal_size", default=4096, min=8, max=4096, step=4, tooltip="Only used for video VAEs: Amount of frames to decode at a time. Higher value than number of frames = disabled"), + io.Int.Input("temporal_overlap", default=16, min=4, max=4096, step=4, tooltip="Only used for video VAEs: Amount of frames to overlap. Higher overlap can improve quality at the edges of temporal tiles but uses more memory and takes more time."), + ]), + ]), + ], + hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], + is_output_node=True, + ) + + @classmethod + def execute(cls, video_latent, video_vae, filename_prefix, format, codec, tiling, audio_latent=None, audio_vae=None, fps=25.0) -> io.NodeOutput: + if tiling["tiling"] == "enabled": + tile_size = tiling["tile_size"] + overlap = tiling["overlap"] + temporal_size = tiling["temporal_size"] + temporal_overlap = tiling["temporal_overlap"] + + if tile_size < overlap * 4: + overlap = tile_size // 4 + if temporal_size < temporal_overlap * 2: + temporal_overlap = temporal_overlap // 2 + temporal_compression = video_vae.temporal_compression_decode() + if temporal_compression is not None: + temporal_size = max(2, temporal_size // temporal_compression) + temporal_overlap = max(1, min(temporal_size // 2, temporal_overlap // temporal_compression)) + else: + temporal_size = None + temporal_overlap = None + + compression = video_vae.spacial_compression_decode() + + images = cls.decode_tiled(video_vae, video_latent["samples"], + tile_t=max(2, temporal_size), + tile_x=tile_size // compression, + tile_y=tile_size // compression, + overlap=(temporal_overlap if temporal_overlap is not None else 1, max(1, overlap // compression), max(1, overlap // compression)), + ).movedim(1, -1) + if len(images.shape) == 5: #Combine batches + images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + else: + images = cls.decode_video(video_vae, video_latent) + + if audio_latent is not None: + if audio_vae is None: + raise ValueError("Audio VAE must be provided if audio latent is provided.") + audio = cls.decode_audio(audio_latent, audio_vae) + else: + audio = None + + video = InputImpl.VideoFromComponents(Types.VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps))) + file, subfolder = cls.save_video(video, filename_prefix, format, codec) + + return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)])) + + @classmethod + def decode_video(cls, vae, samples): + samples_in = samples["samples"] + if samples_in.is_nested: + samples_in = samples_in.unbind()[0] + + vae.throw_exception_if_invalid() + pixel_samples = None + do_tile = False + if vae.latent_dim == 2 and samples_in.ndim == 5: + samples_in = samples_in[:, :, 0] + try: + memory_used = vae.memory_used_decode(samples_in.shape, vae.vae_dtype) + model_management.load_models_gpu([vae.patcher], memory_required=memory_used, force_full_load=True) + free_memory = vae.patcher.get_free_memory(vae.device) + batch_number = int(free_memory / memory_used) + batch_number = max(1, batch_number) + + for x in range(0, samples_in.shape[0], batch_number): + samples = samples_in[x:x+batch_number].to(vae.vae_dtype).to(vae.device) + out = vae.process_output(vae.first_stage_model.decode(samples).to(vae.output_device).to(torch.float16)) + if pixel_samples is None: + pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=vae.output_device, dtype=out.dtype) + pixel_samples[x:x+batch_number] = out + except Exception as e: + model_management.raise_non_oom(e) + logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.") + do_tile = True + + if do_tile: + dims = samples_in.ndim - 2 + if dims == 1 or cls.extra_1d_channel is not None: + pixel_samples = vae.decode_tiled_1d(samples_in) + elif dims == 2: + pixel_samples = vae.decode_tiled_2d(samples_in) + elif dims == 3: + tile = 256 // vae.spacial_compression_decode() + overlap = tile // 4 + pixel_samples = vae.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) + + pixel_samples = pixel_samples.to(vae.output_device).movedim(1,-1) + + if len(pixel_samples.shape) == 5: #Combine batches + pixel_samples = pixel_samples.reshape(-1, pixel_samples.shape[-3], pixel_samples.shape[-2], pixel_samples.shape[-1]) + return pixel_samples + + @classmethod + def decode_tiled(cls, vae, samples, tile_t=999, tile_x=32, tile_y=32, overlap=(1, 8, 8)): + vae.throw_exception_if_invalid() + memory_used = vae.memory_used_decode(samples.shape, vae.vae_dtype) + model_management.load_models_gpu([vae.patcher], memory_required=memory_used, force_full_load=vae.disable_offload) + decode_fn = lambda a: vae.first_stage_model.decode(a.to(vae.vae_dtype).to(vae.device)).to(torch.float16) + return vae.process_output(tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, + upscale_amount=vae.upscale_ratio, out_channels=vae.output_channels, index_formulas=vae.upscale_index_formula, output_device=vae.output_device)) + + + @classmethod + def decode_audio(cls, samples, audio_vae): + audio_latent = samples["samples"] + if audio_latent.is_nested: + audio_latent = audio_latent.unbind()[-1] + audio = audio_vae.decode(audio_latent) + # Post-PR #13486: audio_vae is a comfy.sd.VAE wrapper returning channels-last (BTC). + # Pre-PR: audio_vae is a raw AudioVAE returning channels-first (BCT). + if hasattr(audio_vae, "first_stage_model"): + audio = audio.movedim(-1, 1) + audio = audio.to(audio_latent.device) + output_audio_sample_rate = getattr( + audio_vae, + "audio_sample_rate_output", + getattr(audio_vae, "output_sample_rate", None), + ) + if output_audio_sample_rate is None: + output_audio_sample_rate = getattr( + getattr(audio_vae, "first_stage_model", None), "output_sample_rate", 44100 + ) + return {"waveform": audio, "sample_rate": int(output_audio_sample_rate)} + + @classmethod + def save_video(cls, video, filename_prefix, format, codec) -> io.NodeOutput: + width, height = video.get_dimensions() + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( + filename_prefix, + folder_paths.get_output_directory(), + width, + height + ) + saved_metadata = None + if not args.disable_metadata: + metadata = {} + if cls.hidden.extra_pnginfo is not None: + metadata.update(cls.hidden.extra_pnginfo) + if cls.hidden.prompt is not None: + metadata["prompt"] = cls.hidden.prompt + if len(metadata) > 0: + saved_metadata = metadata + file = f"{filename}_{counter:05}_.{Types.VideoContainer.get_extension(format)}" + video.save_to( + os.path.join(full_output_folder, file), + format=Types.VideoContainer(format), + codec=codec, + metadata=saved_metadata + ) + return file, subfolder + + +class PreviewImageOrMask(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PreviewImageOrMask", + display_name="Preview Image Or Mask", + category="KJNodes/misc", + description="Previews the input images or masks.", + search_aliases=["output"], + inputs=[ + io.MultiType.Input("input", [io.Image, io.Mask], tooltip="The image or mask to preview."), + ], + hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], + is_output_node=True, + ) + + @classmethod + def execute(cls, input) -> io.NodeOutput: + if input.ndim == 3: + return io.NodeOutput(ui=ui.PreviewMask(input, cls=cls)) + return io.NodeOutput(ui=ui.PreviewImage(input, cls=cls)) + diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/image_transform_node.py b/custom_nodes/ComfyUI-KJNodes/nodes/image_transform_node.py new file mode 100644 index 0000000000000000000000000000000000000000..e2dc0f0ec8927182aa23d2a6247f1307226a8995 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/image_transform_node.py @@ -0,0 +1,515 @@ +import numpy as np +import torch +import random +import math +import os +import json + +from PIL import Image + +from comfy.utils import common_upscale +from comfy_api.latest import io +import folder_paths +from nodes import MAX_RESOLUTION + +from ..utility.utility import string_to_color + +def _upscale_mask(mask, width, height, method, crop): + if method == "lanczos": + return common_upscale(mask.unsqueeze(1).repeat(1, 3, 1, 1), width, height, method, crop).movedim(1, -1)[:, :, :, 0] + return common_upscale(mask.unsqueeze(1), width, height, method, crop).squeeze(1) + + +def _resize_single_channel(tensor, width, height): + """Resize a 3D (B,H,W) tensor using bilinear interpolation.""" + return common_upscale(tensor.unsqueeze(1), width, height, "bilinear", "disabled").squeeze(1) + + +def _pad_inputs(): + """Shared pad_top/bottom/left/right input definitions for extra_padding options.""" + return [ + io.Int.Input("pad_top", default=0, min=0, max=MAX_RESOLUTION, step=1, tooltip="Padding pixels on top."), + io.Int.Input("pad_bottom", default=0, min=0, max=MAX_RESOLUTION, step=1, tooltip="Padding pixels on bottom."), + io.Int.Input("pad_left", default=0, min=0, max=MAX_RESOLUTION, step=1, tooltip="Padding pixels on left."), + io.Int.Input("pad_right", default=0, min=0, max=MAX_RESOLUTION, step=1, tooltip="Padding pixels on right."), + ] + + +def _apply_padding(tensor, pad_top, pad_bottom, pad_left, pad_right, mode, edge_mode="clamp", fill_rgb=None): + """Apply padding to a BHWC tensor. Returns the padded tensor. + mode: 'color' or 'edge' + edge_mode: 'clamp', 'repeat', 'mirror' (only used when mode='edge') + fill_rgb: list of [r, g, b] float values 0-1 (only used when mode='color') + """ + h, w = tensor.shape[1], tensor.shape[2] + new_h = h + pad_top + pad_bottom + new_w = w + pad_left + pad_right + + if mode == "color": + fill = fill_rgb or [0.0, 0.0, 0.0] + padded = torch.zeros(tensor.shape[0], new_h, new_w, tensor.shape[3], device=tensor.device, dtype=tensor.dtype) + for c in range(min(3, tensor.shape[3])): + padded[:, :, :, c] = fill[c] + padded[:, pad_top:pad_top+h, pad_left:pad_left+w, :] = tensor + return padded + + # mode == "edge" + if edge_mode == "clamp": + padded = torch.zeros(tensor.shape[0], new_h, new_w, tensor.shape[3], device=tensor.device, dtype=tensor.dtype) + padded[:, pad_top:pad_top+h, pad_left:pad_left+w, :] = tensor + if pad_top > 0: + padded[:, :pad_top, pad_left:pad_left+w, :] = tensor[:, 0:1, :, :].expand(-1, pad_top, -1, -1) + if pad_bottom > 0: + padded[:, pad_top+h:, pad_left:pad_left+w, :] = tensor[:, -1:, :, :].expand(-1, pad_bottom, -1, -1) + if pad_left > 0: + padded[:, :, :pad_left, :] = padded[:, :, pad_left:pad_left+1, :].expand(-1, -1, pad_left, -1) + if pad_right > 0: + padded[:, :, pad_left+w:, :] = padded[:, :, pad_left+w-1:pad_left+w, :].expand(-1, -1, pad_right, -1) + return padded + elif edge_mode == "repeat": + tiles_x = (new_w + w - 1) // w + 1 + tiles_y = (new_h + h - 1) // h + 1 + tiled = tensor.repeat(1, tiles_y, tiles_x, 1) + # Offset so original content lands at (pad_top, pad_left) in output + off_x = (w - pad_left % w) % w + off_y = (h - pad_top % h) % h + return tiled[:, off_y:off_y+new_h, off_x:off_x+new_w, :] + elif edge_mode == "mirror": + flipped_h = tensor.flip(2) + flipped_v = tensor.flip(1) + flipped_hv = tensor.flip(1).flip(2) + mirror_block = torch.cat([ + torch.cat([tensor, flipped_h], dim=2), + torch.cat([flipped_v, flipped_hv], dim=2), + ], dim=1) + mb_h, mb_w = mirror_block.shape[1], mirror_block.shape[2] + tiles_x = (new_w + mb_w - 1) // mb_w + 1 + tiles_y = (new_h + mb_h - 1) // mb_h + 1 + tiled = mirror_block.repeat(1, tiles_y, tiles_x, 1) + # Offset so original content lands at (pad_top, pad_left) in output + off_x = (mb_w - pad_left % mb_w) % mb_w + off_y = (mb_h - pad_top % mb_h) % mb_h + return tiled[:, off_y:off_y+new_h, off_x:off_x+new_w, :] + return tensor + + +class ImageTransformKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ImageTransformKJ", + display_name="Image Transform KJ", + category="KJNodes/image", + search_aliases=["resize", "crop", "pad", "upscale", "keep proportion", "bbox", "bounding box", "transform", "rotate", "mirror"], + is_experimental=True, + description=""" +Interactive image transform node: crop, resize, pad, and rotate. +Connect an image input — the preview appears automatically. + +Cropping: +Click + drag to draw a crop region. +Drag inside to move, drag edges/corners to resize. +Right-click to delete a region. +Ctrl to snap to grid. +Shift + resize to constrain aspect ratio. +Alt + resize to resize symmetrically. + +Padding: +Shift + drag to adjust padding position. + +Rotate button enables rotation cross (drag to rotate, right-click to reset). +Set target_width/height to resize output (0 = keep original). +Use keep_proportion to control how the image fits the target. +Use extra_padding to add padding with color or edge fill (clamp/repeat/mirror).""", + inputs=[ + io.MatchType.Input("image", io.MatchType.Template("img_or_mask", [io.Image, io.Mask]), tooltip="The image or mask to transform."), + io.Mask.Input("mask", optional=True, tooltip="Optional mask to transform alongside the image."), + io.Int.Input("target_width", default=0, min=0, max=MAX_RESOLUTION, step=1, tooltip="Target output width. 0 = keep original dimensions."), + io.Int.Input("target_height", default=0, min=0, max=MAX_RESOLUTION, step=1, tooltip="Target output height. 0 = keep original dimensions."), + io.Combo.Input("upscale_method", options=["nearest-exact", "bilinear", "area", "bicubic", "lanczos"], default="lanczos", tooltip="Interpolation method for resizing."), + io.DynamicCombo.Input("keep_proportion", options=[ + io.DynamicCombo.Option(key="keep_long_edge", inputs=[]), + io.DynamicCombo.Option(key="keep_short_edge", inputs=[]), + io.DynamicCombo.Option(key="total_pixels", inputs=[]), + io.DynamicCombo.Option(key="stretch", inputs=[]), + io.DynamicCombo.Option(key="crop", inputs=[]), + io.DynamicCombo.Option(key="pad_color", inputs=[ + io.Float.Input("pad_x", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Horizontal position of content within padding (0=left, 0.5=center, 1=right). Shift+drag content in preview to adjust."), + io.Float.Input("pad_y", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Vertical position of content within padding (0=top, 0.5=center, 1=bottom). Shift+drag content in preview to adjust."), + ]), + io.DynamicCombo.Option(key="pad_edge", inputs=[ + io.Combo.Input("edge_mode", options=["clamp", "repeat", "mirror"], default="clamp", + tooltip="clamp: extend edge pixels. repeat: tile the image. mirror: tile with mirroring."), + io.Float.Input("pad_x", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Horizontal position of content within padding (0=left, 0.5=center, 1=right). Shift+drag content in preview to adjust."), + io.Float.Input("pad_y", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Vertical position of content within padding (0=top, 0.5=center, 1=bottom). Shift+drag content in preview to adjust."), + ]), + io.DynamicCombo.Option(key="multiplier", inputs=[ + io.Float.Input("width_mult", default=1.0, min=0.01, max=16.0, step=0.05, + tooltip="Multiply the crop width by this factor."), + io.Float.Input("height_mult", default=1.0, min=0.01, max=16.0, step=0.05, + tooltip="Multiply the crop height by this factor."), + ]), + ]), + io.Int.Input("divisible_by", default=2, min=0, max=512, step=1), + io.DynamicCombo.Input("extra_padding", options=[ + io.DynamicCombo.Option(key="disabled", inputs=[]), + io.DynamicCombo.Option(key="pad_color", inputs=_pad_inputs()), + io.DynamicCombo.Option(key="pad_edge", inputs=_pad_inputs() + [ + io.Combo.Input("edge_mode", options=["clamp", "repeat", "mirror"], default="clamp", + tooltip="clamp: extend edge pixels. repeat: tile the image. mirror: tile with mirroring."), + ]), + io.DynamicCombo.Option(key="pad_crop_color", inputs=_pad_inputs()), + io.DynamicCombo.Option(key="pad_crop_edge", inputs=_pad_inputs() + [ + io.Combo.Input("edge_mode", options=["clamp", "repeat", "mirror"], default="clamp", + tooltip="clamp: extend edge pixels. repeat: tile the image. mirror: tile with mirroring."), + ]), + ]), + io.DynamicCombo.Input("invert_crop", options=[ + io.DynamicCombo.Option(key="disabled", inputs=[]), + io.DynamicCombo.Option(key="enabled", inputs=[]), + ]), + io.String.Input("bboxes", default="", socketless=True, advanced=True), + ], + outputs=[ + io.MatchType.Output(io.MatchType.Template("img_or_mask", [io.Image, io.Mask]), id="cropped", display_name="output", is_output_list=True), + io.Mask.Output("cropped_mask", display_name="output_mask", is_output_list=True), + io.BBOX.Output("bbox", display_name="bbox", is_output_list=True), + io.Mask.Output("bbox_mask", display_name="bbox_mask", is_output_list=True), + io.Int.Output("width", display_name="width", tooltip="Width of the output image."), + io.Int.Output("height", display_name="height", tooltip="Height of the output image."), + ], + ) + + + @classmethod + def execute(cls, image, target_width, target_height, upscale_method, keep_proportion, divisible_by, + extra_padding, invert_crop, bboxes, mask=None): + # Unpack DynamicCombos + edge_mode = keep_proportion.get("edge_mode", "clamp") + pad_x = keep_proportion.get("pad_x", 0.5) + pad_y = keep_proportion.get("pad_y", 0.5) + width_mult = keep_proportion.get("width_mult", 1.0) + height_mult = keep_proportion.get("height_mult", 1.0) + keep_proportion = keep_proportion["keep_proportion"] + extra_top = extra_padding.get("pad_top", 0) + extra_bottom = extra_padding.get("pad_bottom", 0) + extra_left = extra_padding.get("pad_left", 0) + extra_right = extra_padding.get("pad_right", 0) + extra_edge_mode = extra_padding.get("edge_mode", "clamp") + extra_pad_mode = extra_padding.get("extra_padding", "disabled") + invert_crop = invert_crop["invert_crop"] + + + # Parse fill color from bboxes JSON (shared color picker) + fill_color_rgb = [0, 0, 0] + if bboxes: + try: + _parsed_tmp = json.loads(bboxes) + if isinstance(_parsed_tmp, dict) and "fillColor" in _parsed_tmp: + fill_color_rgb = string_to_color(_parsed_tmp["fillColor"]) + except (json.JSONDecodeError, Exception): + pass + fill_rgb = [c / 255.0 for c in fill_color_rgb[:3]] + + # Handle mask input (3D) by converting to image-like 4D tensor + input_is_mask = image.ndim == 3 + if input_is_mask: + image = image.unsqueeze(-1).repeat(1, 1, 1, 3) + + # Save input image as temp preview file for JS canvas + temp_dir = folder_paths.get_temp_directory() + pil_img = Image.fromarray((image[0].cpu().numpy() * 255).astype(np.uint8)) + preview_filename = f"crop_preview_{random.randint(0, 0xFFFFFF):06x}.webp" + pil_img.save(os.path.join(temp_dir, preview_filename), format="WEBP", quality=80) + preview_ui = {"preview_filename": [preview_filename]} + + img_height = image.shape[1] + img_width = image.shape[2] + + # Parse bboxes and rotation + bbox_list = [] + rotation = 0.0 + if bboxes: + try: + parsed = json.loads(bboxes) + # New format: { bboxes: [...], rotation: N } + if isinstance(parsed, dict): + bbox_list = [b for b in parsed.get("bboxes", []) if b and all(k in b for k in ("startX", "startY", "endX", "endY"))] + rotation = parsed.get("rotation", 0.0) + # Legacy format: [bbox, bbox, ...] + elif isinstance(parsed, list): + bbox_list = [b for b in parsed if b and all(k in b for k in ("startX", "startY", "endX", "endY"))] + except json.JSONDecodeError: + pass + + # Content mask tracks which pixels are actual image content (1=content, 0=fill) + content_mask = torch.ones(1, img_height, img_width, device=image.device) + + # Apply rotation before cropping + if rotation != 0: + from torchvision.transforms.functional import rotate as tv_rotate + import torch.nn.functional as F + # Use shared fill color for rotation corners (unless edge mode) + rot_fill = fill_rgb + + is_edge_mode = extra_pad_mode in ("pad_edge", "pad_crop_edge") or keep_proportion == "pad_edge" + if is_edge_mode: + h, w = image.shape[1], image.shape[2] + pad_amt = max(h, w) + img_chw = image.movedim(-1, 1) + img_padded = F.pad(img_chw, [pad_amt, pad_amt, pad_amt, pad_amt], mode='replicate') + img_rotated = tv_rotate(img_padded, -rotation, expand=False, fill=rot_fill) + ch, cw = img_rotated.shape[2], img_rotated.shape[3] + cy, cx = ch // 2, cw // 2 + image = img_rotated[:, :, cy - h // 2:cy - h // 2 + h, cx - w // 2:cx - w // 2 + w].movedim(1, -1) + if mask is not None: + mask_padded = F.pad(mask.unsqueeze(1), [pad_amt, pad_amt, pad_amt, pad_amt], mode='replicate') + mask_rotated = tv_rotate(mask_padded, -rotation, expand=False, fill=[0.0]) + mask = mask_rotated[:, :, cy - h // 2:cy - h // 2 + h, cx - w // 2:cx - w // 2 + w].squeeze(1) + # Content mask: rotate the same way (no padding — just rotate and crop) + cm_padded = F.pad(content_mask.unsqueeze(1), [pad_amt, pad_amt, pad_amt, pad_amt], mode='constant', value=0) + cm_rotated = tv_rotate(cm_padded, -rotation, expand=False, fill=[0.0]) + content_mask = cm_rotated[:, :, cy - h // 2:cy - h // 2 + h, cx - w // 2:cx - w // 2 + w].squeeze(1) + else: + image = tv_rotate(image.movedim(-1, 1), -rotation, expand=True, fill=rot_fill).movedim(1, -1) + if mask is not None: + mask = tv_rotate(mask.unsqueeze(1), -rotation, expand=True, fill=[0.0]).squeeze(1) + # Content mask: rotate with expand, fill=0 + content_mask = tv_rotate(content_mask.unsqueeze(1), -rotation, expand=True, fill=[0.0]).squeeze(1) + img_height = image.shape[1] + img_width = image.shape[2] + + # Normalize mask dimensions to match image + if mask is not None: + if mask.shape[-2] != img_height or mask.shape[-1] != img_width: + if mask.shape[-2] == img_width and mask.shape[-1] == img_height: + mask = mask.transpose(-2, -1) + else: + mask = _resize_single_channel(mask, img_width, img_height) + + # "Pad first" modes: apply extra padding to the full image before cropping + # Skip for keep_proportion pad modes — those handle extra padding via target subtraction + is_pad_first = extra_pad_mode in ("pad_color", "pad_edge") + kp_is_pad_mode = keep_proportion in ("pad_color", "pad_edge") + if is_pad_first and not kp_is_pad_mode and (extra_top > 0 or extra_bottom > 0 or extra_left > 0 or extra_right > 0): + pad_mode = "color" if extra_pad_mode == "pad_color" else "edge" + padded_img = _apply_padding(image, extra_top, extra_bottom, extra_left, extra_right, pad_mode, extra_edge_mode, fill_rgb) + + image = padded_img + img_height = image.shape[1] + img_width = image.shape[2] + # Expand content mask and user mask + cm_new = torch.zeros(1, img_height, img_width, device=content_mask.device) + cm_new[:, extra_top:extra_top+content_mask.shape[1], extra_left:extra_left+content_mask.shape[2]] = content_mask + content_mask = cm_new + if mask is not None: + m_new = torch.zeros(mask.shape[0], img_height, img_width, device=mask.device, dtype=mask.dtype) + m_new[:, extra_top:extra_top+mask.shape[1], extra_left:extra_left+mask.shape[2]] = mask + mask = m_new + + # If no bboxes, treat the full image as a single bbox + if not bbox_list: + bbox_list = [None] + + all_cropped = [] + all_cropped_masks = [] + all_bbox_tuples = [] + all_bbox_masks = [] + + for bbox_data in bbox_list: + has_bbox = bbox_data is not None + + if has_bbox: + preview_width = bbox_data.get("previewWidth", 0) + preview_height = bbox_data.get("previewHeight", 0) + sx = img_width / preview_width if preview_width > 0 else 1.0 + sy = img_height / preview_height if preview_height > 0 else 1.0 + + x_min = int(min(bbox_data["startX"], bbox_data["endX"]) * sx) + y_min = int(min(bbox_data["startY"], bbox_data["endY"]) * sy) + x_max = int(max(bbox_data["startX"], bbox_data["endX"]) * sx) + y_max = int(max(bbox_data["startY"], bbox_data["endY"]) * sy) + + x_min = max(0, min(x_min, img_width - 1)) + y_min = max(0, min(y_min, img_height - 1)) + x_max = max(x_min + 1, min(x_max, img_width)) + y_max = max(y_min + 1, min(y_max, img_height)) + + cropped = image[:, y_min:y_max, x_min:x_max, :] + cropped_content_mask = content_mask[:, y_min:y_max, x_min:x_max] + all_bbox_tuples.append((x_min, y_min, x_max - x_min, y_max - y_min)) + bm = torch.zeros(1, img_height, img_width) + bm[0, y_min:y_max, x_min:x_max] = 1.0 + all_bbox_masks.append(bm) + cropped_mask = mask[:, y_min:y_max, x_min:x_max] if mask is not None else None + else: + cropped = image + cropped_content_mask = content_mask + all_bbox_tuples.append((0, 0, img_width, img_height)) + all_bbox_masks.append(torch.ones(1, img_height, img_width)) + cropped_mask = mask + x_min, y_min, x_max, y_max = 0, 0, img_width, img_height + + # Multiplier mode: compute target from crop dims * multiplier + if keep_proportion == "multiplier": + crop_h, crop_w = cropped.shape[1], cropped.shape[2] + tw = round(crop_w * width_mult) + th = round(crop_h * height_mult) + target_width = tw + target_height = th + + # Resize cropped image if target dimensions are set + if target_width > 0 or target_height > 0: + crop_h, crop_w = cropped.shape[1], cropped.shape[2] + tw = target_width if target_width > 0 else crop_w + th = target_height if target_height > 0 else crop_h + + # Subtract extra padding from target so content + padding = original target + # For pad-first + non-pad keep_proportion, padding is on the source (don't subtract) + # For pad modes or pad-crop, subtract so padding is in the output + has_extra = extra_top > 0 or extra_bottom > 0 or extra_left > 0 or extra_right > 0 + kp_is_pad = keep_proportion in ("pad_color", "pad_edge") + if has_extra and (kp_is_pad or not is_pad_first): + if target_width > 0: + tw = max(1, tw - extra_left - extra_right) + if target_height > 0: + th = max(1, th - extra_top - extra_bottom) + + if keep_proportion == "keep_long_edge": + ratio = min(tw / crop_w, th / crop_h) + tw = round(crop_w * ratio) + th = round(crop_h * ratio) + elif keep_proportion == "keep_short_edge": + ratio = max(tw / crop_w, th / crop_h) + tw = round(crop_w * ratio) + th = round(crop_h * ratio) + elif keep_proportion == "total_pixels": + total_pixels = tw * th + aspect_ratio = crop_w / crop_h + th = int(math.sqrt(total_pixels / aspect_ratio)) + tw = int(math.sqrt(total_pixels * aspect_ratio)) + elif keep_proportion == "crop": + ratio = max(tw / crop_w, th / crop_h) + scale_w = round(crop_w * ratio) + scale_h = round(crop_h * ratio) + samples = common_upscale(cropped.movedim(-1, 1), scale_w, scale_h, upscale_method, "center") + cropped = samples.movedim(1, -1) + if cropped_mask is not None: + cropped_mask = _upscale_mask(cropped_mask, scale_w, scale_h, upscale_method, "center") + cropped_content_mask = _resize_single_channel(cropped_content_mask, scale_w, scale_h) + cx = (scale_w - tw) // 2 + cy = (scale_h - th) // 2 + cropped = cropped[:, cy:cy+th, cx:cx+tw, :] + if cropped_mask is not None: + cropped_mask = cropped_mask[:, cy:cy+th, cx:cx+tw] + cropped_content_mask = cropped_content_mask[:, cy:cy+th, cx:cx+tw] + elif keep_proportion in ("pad_color", "pad_edge"): + ratio = min(tw / crop_w, th / crop_h) + scale_w = round(crop_w * ratio) + scale_h = round(crop_h * ratio) + samples = common_upscale(cropped.movedim(-1, 1), scale_w, scale_h, upscale_method, "disabled") + resized = samples.movedim(1, -1) + # pad_x/pad_y position across full target (not just content area) + full_tw = target_width if target_width > 0 else crop_w + full_th = target_height if target_height > 0 else crop_h + pad_left = round((full_tw - scale_w) * pad_x) + pad_top = round((full_th - scale_h) * pad_y) + pad_right = full_tw - pad_left - scale_w + pad_bottom = full_th - pad_top - scale_h + tw = full_tw + th = full_th + + pad_mode = "edge" if keep_proportion == "pad_edge" else "color" + cropped = _apply_padding(resized, pad_top, pad_bottom, pad_left, pad_right, pad_mode, edge_mode, fill_rgb) + if cropped_mask is not None: + mask_resized = _upscale_mask(cropped_mask, scale_w, scale_h, upscale_method, "disabled") + mask_padded = torch.zeros(mask_resized.shape[0], th, tw, device=mask_resized.device, dtype=mask_resized.dtype) + mask_padded[:, pad_top:pad_top+scale_h, pad_left:pad_left+scale_w] = mask_resized + cropped_mask = mask_padded + # Update content mask for padding area + cm_resized = _resize_single_channel(cropped_content_mask, scale_w, scale_h) + cm_padded = torch.zeros(1, th, tw, device=cropped_content_mask.device) + cm_padded[:, pad_top:pad_top+scale_h, pad_left:pad_left+scale_w] = cm_resized + cropped_content_mask = cm_padded + + if divisible_by > 1: + tw = tw - (tw % divisible_by) + th = th - (th % divisible_by) + + if tw > 0 and th > 0: + if keep_proportion in ("stretch", "keep_long_edge", "keep_short_edge", "total_pixels", "multiplier"): + cropped = common_upscale(cropped.movedim(-1, 1), tw, th, upscale_method, "disabled").movedim(1, -1) + if cropped_mask is not None: + cropped_mask = _upscale_mask(cropped_mask, tw, th, upscale_method, "disabled") + cropped_content_mask = _resize_single_channel(cropped_content_mask, tw, th) + else: + cropped = cropped[:, :th, :tw, :] + if cropped_mask is not None: + cropped_mask = cropped_mask[:, :th, :tw] + cropped_content_mask = cropped_content_mask[:, :th, :tw] + + # Enforce divisible_by even when no target dimensions are set + elif divisible_by > 1: + final_w = cropped.shape[2] - (cropped.shape[2] % divisible_by) + final_h = cropped.shape[1] - (cropped.shape[1] % divisible_by) + if final_w != cropped.shape[2] or final_h != cropped.shape[1]: + cropped = cropped[:, :final_h, :final_w, :] + if cropped_mask is not None: + cropped_mask = cropped_mask[:, :final_h, :final_w] + cropped_content_mask = cropped_content_mask[:, :final_h, :final_w] + + # Apply extra padding (skip for pad-first and keep_proportion pad modes which handle it above) + kp_handles_ep = keep_proportion in ("pad_color", "pad_edge") + if not is_pad_first and not kp_handles_ep and (extra_top > 0 or extra_bottom > 0 or extra_left > 0 or extra_right > 0): + h_cur, w_cur = cropped.shape[1], cropped.shape[2] + pad_mode = "edge" if extra_pad_mode == "pad_crop_edge" else "color" + cropped = _apply_padding(cropped, extra_top, extra_bottom, extra_left, extra_right, pad_mode, extra_edge_mode, fill_rgb) + new_h, new_w = cropped.shape[1], cropped.shape[2] + if cropped_mask is not None: + padded_mask = torch.zeros(cropped_mask.shape[0], new_h, new_w, device=cropped_mask.device, dtype=cropped_mask.dtype) + padded_mask[:, extra_top:extra_top+h_cur, extra_left:extra_left+w_cur] = cropped_mask + cropped_mask = padded_mask + cm_h, cm_w = cropped_content_mask.shape[-2], cropped_content_mask.shape[-1] + if cm_h != h_cur or cm_w != w_cur: + cropped_content_mask = _resize_single_channel(cropped_content_mask, w_cur, h_cur) + cm_ep = torch.zeros(1, new_h, new_w, device=cropped_content_mask.device) + cm_ep[:, extra_top:extra_top+h_cur, extra_left:extra_left+w_cur] = cropped_content_mask + cropped_content_mask = cm_ep + + # If no mask was provided, output a zeros mask matching the cropped image + if cropped_mask is None: + cropped_mask = torch.zeros(1, cropped.shape[1], cropped.shape[2]) + + # Apply fill mask — marks filled/padded areas as 1 in the output mask + # Combines with incoming mask: 1 where either input mask is 1 OR area is filled + if cropped_content_mask is not None: + out_h, out_w = cropped_mask.shape[1], cropped_mask.shape[2] + cm_h, cm_w = cropped_content_mask.shape[1], cropped_content_mask.shape[2] + if cm_h != out_h or cm_w != out_w: + cropped_content_mask = _resize_single_channel(cropped_content_mask, out_w, out_h) + # fill_mask: 1 where filled, 0 where content + fill_mask = 1.0 - cropped_content_mask.clamp(0, 1) + # Combine: output mask is max of incoming mask and fill mask + cropped_mask = torch.max(cropped_mask, fill_mask) + + # Invert crop: output area outside the bbox instead of inside + if invert_crop == "enabled" and has_bbox: + inverted = image.clone() + for c in range(min(3, inverted.shape[3])): + inverted[:, y_min:y_max, x_min:x_max, c] = fill_rgb[c] + cropped = inverted + + # Convert back to mask if input was a mask + if input_is_mask: + cropped = cropped[:, :, :, 0] + + all_cropped.append(cropped) + all_cropped_masks.append(cropped_mask) + + width, height = all_cropped[0].shape[2], all_cropped[0].shape[1] + + return io.NodeOutput(all_cropped, all_cropped_masks, all_bbox_tuples, all_bbox_masks, width, height, ui=preview_ui) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/lora_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/lora_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..e2de4fd9c261dcbfb79e847bad29de3b38f88cce --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/lora_nodes.py @@ -0,0 +1,673 @@ +import torch +import comfy.model_management +import comfy.utils +import folder_paths +import os +import logging +from tqdm import tqdm +import numpy as np +from comfy_api.latest import io + +device = comfy.model_management.get_torch_device() + +CLAMP_QUANTILE = 0.99 + +def extract_lora(diff, key, rank, algorithm, lora_type, lowrank_iters=7, adaptive_param=1.0, clamp_quantile=True): + """ + Extracts LoRA weights from a weight difference tensor using SVD. + """ + conv2d = (len(diff.shape) == 4) + kernel_size = None if not conv2d else diff.size()[2:4] + conv2d_3x3 = conv2d and kernel_size != (1, 1) + out_dim, in_dim = diff.size()[0:2] + + if conv2d: + if conv2d_3x3: + diff = diff.flatten(start_dim=1) + else: + diff = diff.squeeze() + + diff_float = diff.float() + if algorithm == "svd_lowrank": + U, S, V = torch.svd_lowrank(diff_float, q=min(rank, in_dim, out_dim), niter=lowrank_iters) + U = U @ torch.diag(S) + Vh = V.t() + else: + #torch.linalg.svdvals() + U, S, Vh = torch.linalg.svd(diff_float) + # Flexible rank selection logic like locon: https://github.com/KohakuBlueleaf/LyCORIS/blob/main/tools/extract_locon.py + if "adaptive" in lora_type: + if lora_type == "adaptive_ratio": + min_s = torch.max(S) * adaptive_param + lora_rank = torch.sum(S > min_s).item() + elif lora_type == "adaptive_energy": + energy = torch.cumsum(S**2, dim=0) + total_energy = torch.sum(S**2) + threshold = adaptive_param * total_energy # e.g., adaptive_param=0.95 for 95% + lora_rank = torch.sum(energy < threshold).item() + 1 + elif lora_type == "adaptive_quantile": + s_cum = torch.cumsum(S, dim=0) + min_cum_sum = adaptive_param * torch.sum(S) + lora_rank = torch.sum(s_cum < min_cum_sum).item() + elif lora_type == "adaptive_fro": + S_squared = S.pow(2) + S_fro_sq = float(torch.sum(S_squared)) + sum_S_squared = torch.cumsum(S_squared, dim=0) / S_fro_sq + lora_rank = int(torch.searchsorted(sum_S_squared, adaptive_param**2)) + 1 + lora_rank = max(1, min(lora_rank, len(S))) + else: + pass # Will print after capping + + # Cap adaptive rank by the specified max rank + lora_rank = min(lora_rank, rank) + + # Calculate and print actual fro percentage retained after capping + if lora_type == "adaptive_fro": + S_squared = S.pow(2) + s_fro = torch.sqrt(torch.sum(S_squared)) + s_red_fro = torch.sqrt(torch.sum(S_squared[:lora_rank])) + fro_percent = float(s_red_fro / s_fro) + logging.info(f"{key} Extracted LoRA rank: {lora_rank}, Frobenius retained: {fro_percent:.1%}") + else: + logging.info(f"{key} Extracted LoRA rank: {lora_rank}") + else: + lora_rank = rank + + lora_rank = max(1, lora_rank) + lora_rank = min(out_dim, in_dim, lora_rank) + + U = U[:, :lora_rank] + S = S[:lora_rank] + U = U @ torch.diag(S) + Vh = Vh[:lora_rank, :] + + if clamp_quantile: + dist = torch.cat([U.flatten(), Vh.flatten()]) + if dist.numel() > 100_000: + # Sample 100,000 elements for quantile estimation + idx = torch.randperm(dist.numel(), device=dist.device)[:100_000] + dist_sample = dist[idx] + hi_val = torch.quantile(dist_sample, CLAMP_QUANTILE) + else: + hi_val = torch.quantile(dist, CLAMP_QUANTILE) + low_val = -hi_val + + U = U.clamp(low_val, hi_val) + Vh = Vh.clamp(low_val, hi_val) + if conv2d: + U = U.reshape(out_dim, lora_rank, 1, 1) + Vh = Vh.reshape(lora_rank, in_dim, kernel_size[0], kernel_size[1]) + return (U, Vh) + + +def calc_lora_model(model_diff, rank, prefix_model, prefix_lora, output_sd, lora_type, algorithm, lowrank_iters, out_dtype, bias_diff=False, adaptive_param=1.0, clamp_quantile=True): + # Get key names from module structure without materializing weights + sd_keys = [] + for name, _ in model_diff.model.named_parameters(): + if prefix_model is None or name.startswith(prefix_model): + sd_keys.append(name) + for name, _ in model_diff.model.named_buffers(): + if prefix_model is None or name.startswith(prefix_model): + sd_keys.append(name) + + total_keys = len([k for k in sd_keys if k.endswith(".weight") or (bias_diff and k.endswith(".bias"))]) + progress_bar = tqdm(total=total_keys, desc=f"Extracting LoRA ({prefix_lora.strip('.')})") + comfy_pbar = comfy.utils.ProgressBar(total_keys) + + # Process one weight at a time to minimize memory usage + for k in sd_keys: + if k.endswith(".weight"): + # Patch and retrieve single weight + weight_diff = model_diff.patch_weight_to_device(k, return_weight=True) + if weight_diff is None: + progress_bar.update(1) + comfy_pbar.update(1) + continue + if weight_diff.ndim == 5: + logging.info(f"Skipping 5D tensor for key {k}") + del weight_diff + progress_bar.update(1) + comfy_pbar.update(1) + continue + if lora_type != "full": + if weight_diff.ndim < 2: + if bias_diff: + output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().to(out_dtype).cpu() + del weight_diff + progress_bar.update(1) + comfy_pbar.update(1) + continue + try: + out = extract_lora(weight_diff.to(device), k, rank, algorithm, lora_type, lowrank_iters=lowrank_iters, adaptive_param=adaptive_param, clamp_quantile=clamp_quantile) + output_sd["{}{}.lora_up.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[0].contiguous().to(out_dtype).cpu() + output_sd["{}{}.lora_down.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[1].contiguous().to(out_dtype).cpu() + except Exception as e: + logging.warning(f"Could not generate lora weights for key {k}, error {e}") + else: + output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().to(out_dtype).cpu() + del weight_diff + progress_bar.update(1) + comfy_pbar.update(1) + + elif bias_diff and k.endswith(".bias"): + weight = model_diff.patch_weight_to_device(k, return_weight=True) + if weight is not None: + output_sd["{}{}.diff_b".format(prefix_lora, k[len(prefix_model):-5])] = weight.contiguous().to(out_dtype).cpu() + del weight + progress_bar.update(1) + comfy_pbar.update(1) + + progress_bar.close() + del model_diff + comfy.model_management.soft_empty_cache() + return output_sd + + +class LoraExtractKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoraExtractKJ", + category="KJNodes/lora", + is_output_node=True, + inputs=[ + io.MultiType.Input("finetuned", [io.Model, io.Clip], tooltip="The finetuned model or clip to extract LoRA from."), + io.MultiType.Input("original", [io.Model, io.Clip], tooltip="The original base model or clip to diff against."), + io.String.Input("filename_prefix", default="loras/ComfyUI_extracted_lora"), + io.Int.Input("rank", default=64, min=1, max=4096, step=1, tooltip="The rank to use for standard LoRA, or maximum rank limit for adaptive methods."), + io.Combo.Input("lora_type", options=["standard", "full", "adaptive_ratio", "adaptive_quantile", "adaptive_energy", "adaptive_fro"]), + io.Combo.Input("algorithm", options=["svd_linalg", "svd_lowrank"], default="svd_lowrank", tooltip="SVD algorithm to use, svd_lowrank is faster but less accurate."), + io.Int.Input("lowrank_iters", default=7, min=1, max=100, step=1, tooltip="The number of subspace iterations for lowrank SVD algorithm."), + io.Combo.Input("output_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), + io.Boolean.Input("bias_diff", default=True), + io.Float.Input("adaptive_param", default=0.15, min=0.0, max=1.0, step=0.01, tooltip="For ratio mode, this is the ratio of the maximum singular value. For quantile mode, this is the quantile of the singular values. For fro mode, this is the Frobenius norm retention ratio."), + io.Boolean.Input("clamp_quantile", default=False), + ], + ) + + + @classmethod + def execute(cls, finetuned, original, filename_prefix, rank, lora_type, algorithm, lowrank_iters, output_dtype, bias_diff, adaptive_param, clamp_quantile) -> io.NodeOutput: + if algorithm == "svd_lowrank" and lora_type != "standard": + raise ValueError("svd_lowrank algorithm is only supported for standard LoRA extraction.") + + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[output_dtype] + + output_dir = folder_paths.get_output_directory() + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir) + + output_sd = {} + + is_clip = hasattr(finetuned, "patcher") + + if is_clip: + clip_diff = finetuned.clone() + kp = original.get_key_patches() + kp = {k: v for k, v in kp.items() if not k.endswith(".position_ids") and not k.endswith(".logit_scale")} + clip_diff.add_patches(kp, -1.0, 1.0) + output_sd = calc_lora_model(clip_diff.patcher, rank, "", "text_encoders.", output_sd, lora_type, algorithm, lowrank_iters, dtype, bias_diff=bias_diff, adaptive_param=adaptive_param, clamp_quantile=clamp_quantile) + else: + m = finetuned.clone() + kp = original.get_key_patches("diffusion_model.") + m.add_patches(kp, -1.0, 1.0) + output_sd = calc_lora_model(m, rank, "diffusion_model.", "diffusion_model.", output_sd, lora_type, algorithm, lowrank_iters, dtype, bias_diff=bias_diff, adaptive_param=adaptive_param, clamp_quantile=clamp_quantile) + + if "adaptive" in lora_type: + rank_str = f"{lora_type}_{adaptive_param:.2f}" + else: + rank_str = rank + output_checkpoint = f"{filename}_rank_{rank_str}_{output_dtype}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(full_output_folder, output_checkpoint) + + comfy.utils.save_torch_file(output_sd, output_checkpoint, metadata=None) + return io.NodeOutput() + +class LoraReduceRank(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoraReduceRankKJ", + display_name="LoraReduceRank", + category="KJNodes/lora", + description="Resize a LoRA model by reducing its rank. Based on kohya's sd-scripts: https://github.com/kohya-ss/sd-scripts/blob/main/networks/resize_lora.py", + is_output_node=True, + is_experimental=True, + inputs=[ + io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"), tooltip="The name of the LoRA."), + io.Int.Input("new_rank", default=8, min=1, max=4096, step=1, tooltip="The new rank to resize the LoRA. Acts as max rank when using dynamic_method."), + io.Combo.Input("dynamic_method", options=["disabled", "sv_ratio", "sv_cumulative", "sv_fro", "sv_knee"], default="disabled", tooltip="Method to use for dynamically determining new alphas and dims. sv_knee finds the elbow point in the singular value curve."), + io.Float.Input("dynamic_param", default=0.2, min=0.0, max=2.0, step=0.01, tooltip="Parameter for dynamic methods. For sv_knee: sensitivity (1.0=standard knee, <1.0=more aggressive/lower rank, >1.0=more conservative)."), + io.Combo.Input("output_dtype", options=["match_original", "fp16", "bf16", "fp32"], default="match_original", tooltip="Data type to save the LoRA as."), + io.Boolean.Input("verbose", default=True), + ], + ) + + @classmethod + def execute(cls, lora_name, new_rank, dynamic_method, dynamic_param, output_dtype, verbose) -> io.NodeOutput: + lora_path = folder_paths.get_full_path("loras", lora_name) + lora_sd, metadata = comfy.utils.load_torch_file(lora_path, return_metadata=True) + + if output_dtype == "fp16": + save_dtype = torch.float16 + elif output_dtype == "bf16": + save_dtype = torch.bfloat16 + elif output_dtype == "fp32": + save_dtype = torch.float32 + elif output_dtype == "match_original": + first_weight_key = next(k for k in lora_sd if k.endswith(".weight") and isinstance(lora_sd[k], torch.Tensor)) + save_dtype = lora_sd[first_weight_key].dtype + + new_lora_sd = {} + for k, v in lora_sd.items(): + new_lora_sd[k.replace(".default", "")] = v + del lora_sd + logging.info("Resizing Lora...") + output_sd, old_dim, new_alpha, rank_list = resize_lora_model(new_lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose) + + if metadata is None: + metadata = {} + + comment = metadata.get("ss_training_comment", "") + + if dynamic_method == "disabled": + metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {new_rank}; {comment}" + metadata["ss_network_dim"] = str(new_rank) + metadata["ss_network_alpha"] = str(new_alpha) + else: + metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from {old_dim}; {comment}" + metadata["ss_network_dim"] = "Dynamic" + metadata["ss_network_alpha"] = "Dynamic" + + for key in list(output_sd.keys()): + value = output_sd[key] + if type(value) == torch.Tensor and value.dtype.is_floating_point and value.dtype != save_dtype: + output_sd[key] = value.to(save_dtype) + + output_dir = folder_paths.get_output_directory() + output_filename_prefix = "loras/" + lora_name + + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(output_filename_prefix, output_dir) + output_dtype_str = f"_{output_dtype}" if output_dtype != "match_original" else "" + average_rank = str(int(np.mean(rank_list))) + rank_str = new_rank if dynamic_method == "disabled" else f"dynamic_{average_rank}" + output_checkpoint = f"{filename.replace('.safetensors', '')}_resized_from_{old_dim}_to_{rank_str}{output_dtype_str}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(full_output_folder, output_checkpoint) + logging.info(f"Saving resized LoRA to {output_checkpoint}") + + comfy.utils.save_torch_file(output_sd, output_checkpoint, metadata=metadata) + return io.NodeOutput() + +# Convert LoRA to different rank approximation (should only be used to go to lower rank) +# This code is based off the extract_lora_from_models.py file which is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py +# Thanks to cloneofsimo + +# This version is based on +# https://github.com/kohya-ss/sd-scripts/blob/main/networks/resize_lora.py + +MIN_SV = 1e-6 + +LORA_DOWN_UP_FORMATS = [ + ("lora_down", "lora_up"), # sd-scripts LoRA + ("lora_A", "lora_B"), # PEFT LoRA + ("down", "up"), # ControlLoRA +] + +# Indexing functions +def index_sv_cumulative(S, target): + original_sum = float(torch.sum(S)) + cumulative_sums = torch.cumsum(S, dim=0) / original_sum + index = int(torch.searchsorted(cumulative_sums, target)) + 1 + index = max(1, min(index, len(S) - 1)) + + return index + + +def index_sv_fro(S, target): + S_squared = S.pow(2) + S_fro_sq = float(torch.sum(S_squared)) + sum_S_squared = torch.cumsum(S_squared, dim=0) / S_fro_sq + index = int(torch.searchsorted(sum_S_squared, target**2)) + 1 + index = max(1, min(index, len(S) - 1)) + + return index + + +def index_sv_ratio(S, target): + max_sv = S[0] + min_sv = max_sv / target + index = int(torch.sum(S > min_sv).item()) + index = max(1, min(index, len(S) - 1)) + + return index + + +def index_sv_knee(S, sensitivity=1.0): + """Find the knee/elbow point in the singular value curve. + Uses the Kneedle method: normalizes the curve to [0,1] on both axes, + then finds the point with maximum distance from a reference line. + + sensitivity controls the aggressiveness: + 1.0 = standard knee point + < 1.0 = more aggressive (lower rank), e.g. 0.5 roughly halves the knee rank + > 1.0 = more conservative (higher rank) + The reference line tilts based on sensitivity: at lower values, + the line favors keeping fewer singular values.""" + n = len(S) + if n <= 2: + return 1 + + S_np = S.cpu().float().numpy() + + # Normalize x and y to [0, 1] + x = np.linspace(0, 1, n) + y = (S_np - S_np[-1]) / (S_np[0] - S_np[-1] + 1e-10) + + # Reference line from (0, 1) to (1, 1-sensitivity) + # At sensitivity=1.0: line goes (0,1)->(1,0), standard kneedle + # At sensitivity=0.5: line goes (0,1)->(1,0.5), steeper = more aggressive + y_line = 1.0 - sensitivity * x + + # Signed distance: positive means curve is above the line (keep) + distances = y - y_line + + # Find the point of maximum positive distance + index = int(np.argmax(distances)) + index = max(1, min(index, n - 1)) + return index + + +# Modified from Kohaku-blueleaf's extract/merge functions +def _svd_extract(weight_2d, lora_rank, dynamic_method, dynamic_param, device, scale=1): + """Shared SVD extraction for both linear and conv weights. + Dynamic mode: single full SVD (need all singular values for rank selection). + Disabled mode: svd_lowrank only (rank is known, much faster for large matrices).""" + weight_2d = weight_2d.to(device) + if weight_2d.dtype != torch.float32: + weight_2d = weight_2d.float() + + if dynamic_method and dynamic_method != "disabled": + # Full SVD: we need all singular values for dynamic rank selection, + # and we reuse U/Vh directly — one SVD, no wasted work + U, S, Vh = torch.linalg.svd(weight_2d, full_matrices=False) + param_dict = rank_resize(S, lora_rank, dynamic_method, dynamic_param, scale) + lora_rank = param_dict["new_rank"] + U = U[:, :lora_rank] + S = S[:lora_rank] + Vh = Vh[:lora_rank, :] + else: + # Randomized lowrank SVD: only compute top-k, much faster when rank << min(m,n) + U, S, V = torch.svd_lowrank(weight_2d, q=lora_rank, niter=7) + Vh = V.t() + param_dict = {"new_rank": lora_rank, "new_alpha": float(scale * lora_rank)} + del V + + sqrt_S = torch.diag(torch.sqrt(S)) + U = U @ sqrt_S + Vh = sqrt_S @ Vh + + return U, Vh, param_dict + + +def extract_conv(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1): + out_size, in_size, kernel_size, _ = weight.size() + + U, Vh, param_dict = _svd_extract( + weight.reshape(out_size, -1), lora_rank, dynamic_method, dynamic_param, device, scale + ) + lora_rank = param_dict["new_rank"] + + param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kernel_size, kernel_size).cpu() + param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1).cpu() + del U, Vh, weight + return param_dict + + +def extract_linear(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1): + out_size, in_size = weight.size() + + U, Vh, param_dict = _svd_extract( + weight, lora_rank, dynamic_method, dynamic_param, device, scale + ) + lora_rank = param_dict["new_rank"] + + param_dict["lora_down"] = Vh.reshape(lora_rank, in_size).cpu() + param_dict["lora_up"] = U.reshape(out_size, lora_rank).cpu() + del U, Vh, weight + return param_dict + + +def merge_conv(lora_down, lora_up, device): + in_rank, in_size, kernel_size, k_ = lora_down.shape + out_size, out_rank, _, _ = lora_up.shape + assert in_rank == out_rank and kernel_size == k_, f"rank {in_rank} {out_rank} or kernel {kernel_size} {k_} mismatch" + + lora_down = lora_down.to(device) + lora_up = lora_up.to(device) + + merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1) + weight = merged.reshape(out_size, in_size, kernel_size, kernel_size) + del lora_up, lora_down + return weight + + +def merge_linear(lora_down, lora_up, device): + in_rank, in_size = lora_down.shape + out_size, out_rank = lora_up.shape + assert in_rank == out_rank, f"rank {in_rank} {out_rank} mismatch" + + lora_down = lora_down.to(device) + lora_up = lora_up.to(device) + + weight = lora_up @ lora_down + del lora_up, lora_down + return weight + + +def merge_conv3d(lora_down, lora_up, device): + in_rank, in_size, kD, kH, kW = lora_down.shape + out_size, out_rank, _, _, _ = lora_up.shape + assert in_rank == out_rank, f"rank {in_rank} {out_rank} mismatch" + + lora_down = lora_down.to(device) + lora_up = lora_up.to(device) + + merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank, -1) + weight = merged.reshape(out_size, in_size, kD, kH, kW) + del lora_up, lora_down + return weight + + +def extract_conv3d(weight, lora_rank, dynamic_method, dynamic_param, device, scale=1): + out_size, in_size, kD, kH, kW = weight.size() + + U, Vh, param_dict = _svd_extract( + weight.reshape(out_size, -1), lora_rank, dynamic_method, dynamic_param, device, scale + ) + lora_rank = param_dict["new_rank"] + + param_dict["lora_down"] = Vh.reshape(lora_rank, in_size, kD, kH, kW).cpu() + param_dict["lora_up"] = U.reshape(out_size, lora_rank, 1, 1, 1).cpu() + del U, Vh, weight + return param_dict + + +# Calculate new rank + + +def rank_resize(S, rank, dynamic_method, dynamic_param, scale=1): + param_dict = {} + + if dynamic_method == "sv_ratio": + # Calculate new dim and alpha based off ratio + new_rank = index_sv_ratio(S, dynamic_param) + 1 + new_alpha = float(scale * new_rank) + + elif dynamic_method == "sv_cumulative": + # Calculate new dim and alpha based off cumulative sum + new_rank = index_sv_cumulative(S, dynamic_param) + 1 + new_alpha = float(scale * new_rank) + + elif dynamic_method == "sv_fro": + # Calculate new dim and alpha based off sqrt sum of squares + new_rank = index_sv_fro(S, dynamic_param) + 1 + new_alpha = float(scale * new_rank) + + elif dynamic_method == "sv_knee": + # Knee/elbow detection in singular value curve + new_rank = index_sv_knee(S, dynamic_param) + 1 + new_alpha = float(scale * new_rank) + else: + new_rank = rank + new_alpha = float(scale * new_rank) + + if S[0] <= MIN_SV: # Zero matrix, set dim to 1 + new_rank = 1 + new_alpha = float(scale * new_rank) + elif new_rank > rank: # cap max rank at rank + new_rank = rank + new_alpha = float(scale * new_rank) + + # Calculate resize info + s_sum = torch.sum(torch.abs(S)) + s_rank = torch.sum(torch.abs(S[:new_rank])) + + S_squared = S.pow(2) + s_fro = torch.sqrt(torch.sum(S_squared)) + s_red_fro = torch.sqrt(torch.sum(S_squared[:new_rank])) + fro_percent = float(s_red_fro / s_fro) + + param_dict["new_rank"] = new_rank + param_dict["new_alpha"] = new_alpha + param_dict["sum_retained"] = (s_rank) / s_sum + param_dict["fro_retained"] = fro_percent + param_dict["max_ratio"] = S[0] / S[new_rank - 1] + + return param_dict + + +def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose): + max_old_rank = None + new_alpha = None + verbose_str = "\n" + fro_list = [] + rank_list = [] + + if dynamic_method: + logging.info(f"Dynamically determining new alphas and dims based off {dynamic_method}: {dynamic_param}, max rank is {new_rank}") + + lora_down_weight = None + lora_up_weight = None + + o_lora_sd = lora_sd.copy() + block_down_name = None + block_up_name = None + + total_keys = len([k for k in lora_sd if k.endswith(".weight")]) + + pbar = comfy.utils.ProgressBar(total_keys) + for key, value in tqdm(lora_sd.items(), leave=True, desc="Resizing LoRA weights"): + key_parts = key.split(".") + block_down_name = None + for _format in LORA_DOWN_UP_FORMATS: + # Currently we only match lora_down_name in the last two parts of key + # because ("down", "up") are general words and may appear in block_down_name + if len(key_parts) >= 2 and _format[0] == key_parts[-2]: + block_down_name = ".".join(key_parts[:-2]) + lora_down_name = "." + _format[0] + lora_up_name = "." + _format[1] + weight_name = "." + key_parts[-1] + break + if len(key_parts) >= 1 and _format[0] == key_parts[-1]: + block_down_name = ".".join(key_parts[:-1]) + lora_down_name = "." + _format[0] + lora_up_name = "." + _format[1] + weight_name = "" + break + + if block_down_name is None: + # This parameter is not lora_down + continue + + # Now weight_name can be ".weight" or "" + # Find corresponding lora_up and alpha + block_up_name = block_down_name + lora_down_weight = value + lora_up_weight = lora_sd.get(block_up_name + lora_up_name + weight_name, None) + lora_alpha = lora_sd.get(block_down_name + ".alpha", None) + + weights_loaded = lora_down_weight is not None and lora_up_weight is not None + + if weights_loaded: + + conv2d = len(lora_down_weight.size()) == 4 + conv3d = len(lora_down_weight.size()) == 5 + old_rank = lora_down_weight.size()[0] + max_old_rank = max(max_old_rank or 0, old_rank) + + # Skip if merged weight would be too large (>100k elements in any dimension) + if conv2d: + in_rank, in_size, kernel_size, _ = lora_down_weight.shape + out_size, out_rank, _, _ = lora_up_weight.shape + merged_size = out_size * in_size * kernel_size * kernel_size + elif conv3d: + in_rank, in_size, kD, kH, kW = lora_down_weight.shape + out_size, out_rank, _, _, _ = lora_up_weight.shape + merged_size = out_size * in_size * kD * kH * kW + else: + in_rank, in_size = lora_down_weight.shape + out_size, out_rank = lora_up_weight.shape + merged_size = out_size * in_size + + if merged_size > 100_000_000: # Skip if >100M elements + logging.warning(f"Skipping {block_down_name}: merged weight too large ({merged_size:,} elements)") + tqdm.write(f"SKIPPED: {block_down_name} - too large ({merged_size:,} elements)") + pbar.update(1) + continue + + if lora_alpha is None: + scale = 1.0 + else: + scale = lora_alpha / old_rank + + if conv2d: + full_weight_matrix = merge_conv(lora_down_weight, lora_up_weight, device) + param_dict = extract_conv(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale) + elif conv3d: + full_weight_matrix = merge_conv3d(lora_down_weight, lora_up_weight, device) + param_dict = extract_conv3d(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale) + else: + full_weight_matrix = merge_linear(lora_down_weight, lora_up_weight, device) + param_dict = extract_linear(full_weight_matrix, new_rank, dynamic_method, dynamic_param, device, scale) + + if verbose and "fro_retained" in param_dict: + max_ratio = param_dict["max_ratio"] + sum_retained = param_dict["sum_retained"] + fro_retained = param_dict["fro_retained"] + if not np.isnan(fro_retained): + fro_list.append(float(fro_retained)) + log_str = f"{block_down_name:75} | sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}, new dim: {param_dict['new_rank']}" + tqdm.write(log_str) + verbose_str += log_str + + if verbose and dynamic_method: + verbose_str += f", dynamic | dim: {param_dict['new_rank']}, alpha: {param_dict['new_alpha']}\n" + else: + verbose_str += "\n" + + new_alpha = param_dict["new_alpha"] + o_lora_sd[block_down_name + lora_down_name + weight_name] = param_dict["lora_down"].to(save_dtype).contiguous() + o_lora_sd[block_up_name + lora_up_name + weight_name] = param_dict["lora_up"].to(save_dtype).contiguous() + o_lora_sd[block_down_name + ".alpha"] = torch.tensor(param_dict["new_alpha"]).to(save_dtype) + + block_down_name = None + block_up_name = None + lora_down_weight = None + lora_up_weight = None + weights_loaded = False + rank_list.append(param_dict["new_rank"]) + del param_dict + pbar.update(1) + + if verbose: + logging.info(f"Average Frobenius norm retention: {np.mean(fro_list):.2%} | std: {np.std(fro_list):0.3f}") + return o_lora_sd, max_old_rank, new_alpha, rank_list diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/ltxv_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/ltxv_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..44b63f6ebac38f5c63f565f247db5e8e486831f4 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/ltxv_nodes.py @@ -0,0 +1,2136 @@ +from comfy_extras.nodes_lt import get_noise_mask, LTXVAddGuide, _append_guide_attention_entry +import types +import math +from typing import Tuple +import comfy +from comfy_api.latest import io +import numpy as np +import torch +import logging +import comfy.model_management as mm +import comfy.ldm.modules.attention as _comfy_attn +from comfy.ldm.lightricks.model import apply_rotary_emb as _apply_rope +try: + from comfy.ldm.lightricks.model import GuideAttentionMask as _GuideAttentionMask, _attention_with_guide_mask as _ltx_attn_with_guide_mask +except ImportError: + _GuideAttentionMask = None + _ltx_attn_with_guide_mask = None +device = mm.get_torch_device() +import latent_preview + +class LTXVAddGuideMulti(LTXVAddGuide): + + @classmethod + def define_schema(cls): + options = [] + for num_guides in range(1, 21): # 1 to 20 guides + guide_inputs = [] + for i in range(1, num_guides + 1): + guide_inputs.extend([ + io.Image.Input(f"image_{i}"), + io.Int.Input( + f"frame_idx_{i}", + default=0, + min=-9999, + max=9999, + tooltip=f"Frame index for guide {i}.", + ), + io.Float.Input(f"strength_{i}", default=1.0, min=0.0, max=10.0, step=0.01, tooltip=f"Strength for guide {i}."), + ]) + options.append(io.DynamicCombo.Option( + key=str(num_guides), + inputs=guide_inputs + )) + + return io.Schema( + node_id="LTXVAddGuideMulti", + category="KJNodes/ltxv", + description="Add multiple guide images at specified frame indices with strengths, uses DynamicCombo which requires ComfyUI 0.8.1 and frontend 1.33.4 or later.", + inputs=[ + io.Conditioning.Input("positive", tooltip="Positive conditioning to which guide keyframe info will be added"), + io.Conditioning.Input("negative", tooltip="Negative conditioning to which guide keyframe info will be added"), + io.Vae.Input("vae", tooltip="Video VAE used to encode the guide images"), + io.Latent.Input("latent", tooltip="Video latent, guides are added to the end of this latent"), + io.DynamicCombo.Input( + "num_guides", + options=options, + display_name="Number of Guides", + tooltip="Select how many guide images to use", + ), + ], + outputs=[ + io.Conditioning.Output(display_name="positive"), + io.Conditioning.Output(display_name="negative"), + io.Latent.Output(display_name="latent", tooltip="Video latent with added guides"), + ], + ) + + @classmethod + def execute(cls, positive, negative, vae, latent, num_guides) -> io.NodeOutput: + scale_factors = vae.downscale_index_formula + latent_image = latent["samples"] + noise_mask = get_noise_mask(latent) + + _, _, latent_length, latent_height, latent_width = latent_image.shape + + # num_guides is a dict containing the inputs from the selected option + # e.g., {'image_1': tensor, 'frame_idx_1': 0, 'strength_1': 1.0, 'image_2': tensor, 'frame_idx_2': 20, 'strength_2': 0.8, ...} + + image_keys = sorted([k for k in num_guides.keys() if k.startswith('image_')]) + + for img_key in image_keys: + i = img_key.split('_')[1] + + img = num_guides[f"image_{i}"] + f_idx = num_guides[f"frame_idx_{i}"] + strength = num_guides[f"strength_{i}"] + + image_1, t = cls.encode(vae, latent_width, latent_height, img, scale_factors) + + frame_idx, latent_idx = cls.get_latent_index(positive, latent_length, len(image_1), f_idx, scale_factors) + assert latent_idx + t.shape[2] <= latent_length, "Conditioning frames exceed the length of the latent sequence." + + positive, negative, latent_image, noise_mask = cls.append_keyframe( + positive, + negative, + frame_idx, + latent_image, + noise_mask, + t, + strength, + scale_factors, + ) + + # Track this guide for per-reference attention control. + pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4] + guide_latent_shape = list(t.shape[2:]) # [F, H, W] + positive, negative = _append_guide_attention_entry(positive, negative, pre_filter_count, guide_latent_shape, strength=strength) + + return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask}) + +class LTXVAddGuidesFromBatch(LTXVAddGuide): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVAddGuidesFromBatch", + category="conditioning/ltxv", + description="Adds multiple guide images from a batch to the latent at corresponding frame indices. Non-black images in the batch are used as guides.", + inputs=[ + io.Conditioning.Input("positive"), + io.Conditioning.Input("negative"), + io.Vae.Input("vae"), + io.Latent.Input("latent"), + io.Image.Input("images", tooltip="Batch of images - non-black images will be used as guides"), + io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Strength for all guides."), + ], + outputs=[ + io.Conditioning.Output(display_name="positive"), + io.Conditioning.Output(display_name="negative"), + io.Latent.Output(display_name="latent"), + ], + ) + + @classmethod + def execute(cls, positive, negative, vae, latent, images, strength) -> io.NodeOutput: + scale_factors = vae.downscale_index_formula + latent_image = latent["samples"] + noise_mask = get_noise_mask(latent) + + _, _, latent_length, latent_height, latent_width = latent_image.shape + + # Process each image in the batch + batch_size = images.shape[0] + + for i in range(batch_size): + img = images[i:i+1] + + # Check if image is not black and use batch index as frame index + if img.max() > 0.001: + f_idx = i + + image_1, t = cls.encode(vae, latent_width, latent_height, img, scale_factors) + + frame_idx, latent_idx = cls.get_latent_index(positive, latent_length, len(image_1), f_idx, scale_factors) + + if latent_idx + t.shape[2] <= latent_length: + positive, negative, latent_image, noise_mask = cls.append_keyframe( + positive, + negative, + frame_idx, + latent_image, + noise_mask, + t, + strength, + scale_factors, + ) + + # Track this guide for per-reference attention control. + pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4] + guide_latent_shape = list(t.shape[2:]) # [F, H, W] + positive, negative = _append_guide_attention_entry(positive, negative, pre_filter_count, guide_latent_shape, strength=strength) + else: + print(f"Warning: Skipping guide at index {i} - conditioning frames exceed latent sequence length") + + return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask}) + + +class LTXVAudioVideoMask(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVAudioVideoMask", + category="KJNodes/ltxv", + description="Creates noise masks for video and audio latents based on specified time ranges. New content is generated within these masked regions", + inputs=[ + io.Latent.Input("video_latent", optional=True), + io.Latent.Input("audio_latent", optional=True), + io.Float.Input("video_fps", default=25, min=0.0, max=100.0, step=0.01), + io.Float.Input("video_start_time", default=0.0, min=0.0, max=10000.0, step=0.01, tooltip="Start time in seconds for the video mask."), + io.Float.Input("video_end_time", default=5.0, min=0.0, max=10000.0, step=0.01, tooltip="End time in seconds for the video mask."), + io.Float.Input("audio_start_time", default=0.0, min=0.0, max=10000.0, step=0.01, tooltip="Start time in seconds for the audio mask."), + io.Float.Input("audio_end_time", default=5.0, min=0.0, max=10000.0, step=0.01, tooltip="End time in seconds for the audio mask."), + io.Combo.Input( + "max_length", + options=["truncate", "pad", "partial"], + default="truncate", + tooltip="'truncate': cut latent to end_time length. 'pad': extend latent to end_time. 'partial': mask range within existing latent.", + ), + io.Combo.Input("existing_mask_mode", options=["add", "subtract", "overwrite"], optional=True, default="add", tooltip="How to combine with existing noise masks if present. 'add' will take the max of existing and new mask, 'overwrite' will replace with new mask. 'subtract' will set the masked region to 0 instead of 1, effectively unmasking it."), + ], + outputs=[ + io.Latent.Output(display_name="video_latent"), + io.Latent.Output(display_name="audio_latent"), + ], + ) + + @classmethod + def execute(cls, video_fps, video_start_time, video_end_time, audio_start_time, audio_end_time, max_length="truncate", existing_mask_mode="add", video_latent=None, audio_latent=None,) -> io.NodeOutput: + + time_scale_factor = 8 + mel_hop_length = 160 + sampling_rate = 16000 + latent_downsample_factor = 4 + audio_latents_per_second = (sampling_rate / mel_hop_length / latent_downsample_factor) # 25 + + if video_latent is not None: + video_latent_frame_count = video_latent["samples"].shape[2] + + video_pixel_frame_start_raw = int(round(video_start_time * video_fps)) + video_pixel_frame_end_raw = int(round(video_end_time * video_fps)) + + # Calculate required latent frames based on end time + required_latent_frames = (video_pixel_frame_end_raw - 1) // time_scale_factor + 1 + + # Handle different max_length modes + if max_length == "pad" and required_latent_frames > video_latent_frame_count: + # Pad video latent if required frames exceed current length + pad_frames = required_latent_frames - video_latent_frame_count + padding = torch.zeros( + video_latent["samples"].shape[0], + video_latent["samples"].shape[1], + pad_frames, + video_latent["samples"].shape[3], + video_latent["samples"].shape[4], + dtype=video_latent["samples"].dtype, + device=video_latent["samples"].device + ) + video_samples = torch.cat([video_latent["samples"], padding], dim=2) + video_latent_frame_count = video_samples.shape[2] + elif max_length == "truncate": + # Truncate to the end_time + video_samples = video_latent["samples"][:, :, :required_latent_frames] + video_latent_frame_count = video_samples.shape[2] + else: # partial + video_samples = video_latent["samples"] + + # Now calculate indices based on potentially padded latent + video_pixel_frame_count = (video_latent_frame_count - 1) * time_scale_factor + 1 + xp = np.array([0] + list(range(1, video_pixel_frame_count + time_scale_factor, time_scale_factor))) + + # video_frame_index_start = index of the value in xp rounding up + video_latent_frame_index_start = np.searchsorted(xp, video_pixel_frame_start_raw, side="left") + # video_frame_index_end = index of the value in xp rounding down + video_latent_frame_index_end = np.searchsorted(xp, video_pixel_frame_end_raw, side="right") - 1 + + video_latent_frame_index_start = max(0, video_latent_frame_index_start) + video_latent_frame_index_end = min(video_latent_frame_index_end, video_latent_frame_count) + + # Get existing noise mask if present, otherwise create new one + if "noise_mask" in video_latent and video_latent["noise_mask"] is not None and existing_mask_mode != "overwrite": + video_mask = video_latent["noise_mask"].clone() + # Adjust mask size based on mode + if max_length == "pad" and video_samples.shape[2] > video_latent["samples"].shape[2]: + # Pad the mask if we padded the samples + mask_padding = torch.zeros( + video_mask.shape[0], + video_mask.shape[1], + video_samples.shape[2] - video_mask.shape[2], + video_mask.shape[3], + video_mask.shape[4], + dtype=video_mask.dtype, + device=video_mask.device + ) + video_mask = torch.cat([video_mask, mask_padding], dim=2) + elif max_length == "truncate": + # Truncate the mask to match truncated samples + video_mask = video_mask[:, :, :video_samples.shape[2]] + else: + video_mask = torch.zeros_like(video_samples)[:, :1] + + video_mask[:, :, video_latent_frame_index_start:video_latent_frame_index_end] = 1.0 if existing_mask_mode != "subtract" else 0.0 + # ensure all padded frames are also masked + if max_length == "pad" and video_samples.shape[2] > video_latent["samples"].shape[2]: + video_mask[:, :, video_latent["samples"].shape[2]:] = 1.0 if existing_mask_mode != "subtract" else 0.0 + video_latent = video_latent.copy() + video_latent["samples"] = video_samples + video_latent["noise_mask"] = video_mask + + if audio_latent is not None: + audio_latent_frame_count = audio_latent["samples"].shape[2] + + audio_latent_frame_index_start = int(round(audio_start_time * audio_latents_per_second)) + audio_latent_frame_index_end = int(round(audio_end_time * audio_latents_per_second)) + 1 + + # Handle different max_length modes + if max_length == "pad" and audio_latent_frame_index_end > audio_latent_frame_count: + # Pad audio latent if end index exceeds current length + pad_frames = audio_latent_frame_index_end - audio_latent_frame_count + padding = torch.zeros( + audio_latent["samples"].shape[0], + audio_latent["samples"].shape[1], + pad_frames, + audio_latent["samples"].shape[3], + dtype=audio_latent["samples"].dtype, + device=audio_latent["samples"].device + ) + audio_samples = torch.cat([audio_latent["samples"], padding], dim=2) + audio_latent_frame_count = audio_samples.shape[2] + elif max_length == "truncate": + # Truncate to the end_time + audio_samples = audio_latent["samples"][:, :, :audio_latent_frame_index_end] + audio_latent_frame_count = audio_samples.shape[2] + else: # partial + audio_samples = audio_latent["samples"] + + audio_latent_frame_index_start = max(0, audio_latent_frame_index_start) + audio_latent_frame_index_end = min(audio_latent_frame_index_end, audio_latent_frame_count) + + # Get existing noise mask if present, otherwise create new one + if "noise_mask" in audio_latent and audio_latent["noise_mask"] is not None and existing_mask_mode != "overwrite": + audio_mask = audio_latent["noise_mask"].clone() + # Adjust mask size based on mode + if max_length == "pad" and audio_samples.shape[2] > audio_latent["samples"].shape[2]: + # Pad the mask if we padded the samples + mask_padding = torch.zeros( + audio_mask.shape[0], + audio_mask.shape[1], + audio_samples.shape[2] - audio_mask.shape[2], + audio_mask.shape[3], + dtype=audio_mask.dtype, + device=audio_mask.device + ) + audio_mask = torch.cat([audio_mask, mask_padding], dim=2) + elif max_length == "truncate": + # Truncate the mask to match truncated samples + audio_mask = audio_mask[:, :, :audio_samples.shape[2]] + else: + audio_mask = torch.zeros_like(audio_samples) + + audio_mask[:, :, audio_latent_frame_index_start:audio_latent_frame_index_end] = 1.0 + # ensure all padded frames are also masked + if max_length == "pad" and audio_samples.shape[2] > audio_latent["samples"].shape[2]: + audio_mask[:, :, audio_latent["samples"].shape[2]:] = 1.0 + audio_latent = audio_latent.copy() + audio_latent["samples"] = audio_samples + audio_latent["noise_mask"] = audio_mask + + return io.NodeOutput(video_latent, audio_latent) + +def _compute_attention(self, query, context, attn_precision=None, transformer_options={}, mask=None): + """Compute attention and return the result. Cleans up intermediate tensors. + + When `mask` is non-None, bypass wrap_attn and call attention_pytorch directly — + sage and similar backends drop arbitrary masks. + """ + k = self.k_norm(self.to_k(context)).to(query.dtype) + v = self.to_v(context).to(query.dtype) + if mask is None: + x = comfy.ldm.modules.attention.optimized_attention(query, k, v, heads=self.heads, attn_precision=attn_precision, transformer_options=transformer_options).flatten(2) + else: + x = comfy.ldm.modules.attention.attention_pytorch(query, k, v, heads=self.heads, mask=mask, attn_precision=attn_precision, _inside_attn_wrapper=True, transformer_options=transformer_options).flatten(2) + del k, v + return x + +def nag_attention(self, query, context_positive, nag_context, attn_precision=None, transformer_options={}, mask=None): + # mask (e.g. PromptRelay's temporal routing) only applies to the real-context, not NAG context + x_positive = _compute_attention(self, query, context_positive, attn_precision, transformer_options, mask=mask) + x_negative = _compute_attention(self, query, nag_context, attn_precision, transformer_options) + return x_positive, x_negative + +def normalized_attention_guidance(self, x_positive, x_negative): + if self.inplace: + nag_guidance = x_negative.mul_(self.nag_scale - 1).neg_().add_(x_positive, alpha=self.nag_scale) + del x_negative + else: + nag_guidance = x_negative * (self.nag_scale - 1) + del x_negative + nag_guidance = (x_positive * self.nag_scale).sub_(nag_guidance) + + norm_positive = torch.norm(x_positive, p=1, dim=-1, keepdim=True) + norm_guidance = torch.norm(nag_guidance, p=1, dim=-1, keepdim=True) + + scale = norm_guidance / norm_positive + torch.nan_to_num_(scale, nan=10.0) + mask = scale > self.nag_tau + del scale + + adjustment = (norm_positive * self.nag_tau) / (norm_guidance + 1e-7) + del norm_positive, norm_guidance + + nag_guidance.mul_(torch.where(mask, adjustment, 1.0)) + del mask, adjustment + + if self.inplace: + return nag_guidance.sub_(x_positive).mul_(self.nag_alpha).add_(x_positive) + else: + nag_guidance.mul_(self.nag_alpha) + return nag_guidance.add_(x_positive * (1 - self.nag_alpha)) + +#region NAG +def ltxv_crossattn_forward_nag(self, x, context, mask=None, transformer_options={}, **kwargs): + + if mask is None: + mask_provider = transformer_options.get("promptrelay_mask_fn") + if mask_provider is not None: + mask = mask_provider(x.shape[1], context.shape[1], x.dtype, x.device, transformer_options) + + # Single or [pos, neg] pair + if context.shape[0] == 1: + x_pos, context_pos = x, context + x_neg, context_neg = None, None + else: + x_pos, x_neg = torch.chunk(x, 2, dim=0) + context_pos, context_neg = torch.chunk(context, 2, dim=0) + + # Positive + q_pos = self.q_norm(self.to_q(x_pos)) + del x_pos + + x_positive, x_negative = nag_attention(self, q_pos, context_pos, self.nag_context, attn_precision=self.attn_precision, transformer_options=transformer_options, mask=mask) + del context_pos, q_pos + + x_pos_out = normalized_attention_guidance(self, x_positive, x_negative) + del x_positive, x_negative + + # Negative + if x_neg is not None and context_neg is not None: + q_neg = self.q_norm(self.to_q(x_neg)) + k_neg = self.k_norm(self.to_k(context_neg)) + v_neg = self.to_v(context_neg) + + x_neg_out = comfy.ldm.modules.attention.optimized_attention(q_neg, k_neg, v_neg, heads=self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options) + out = torch.cat([x_pos_out, x_neg_out], dim=0) + else: + out = x_pos_out + + if self.to_gate_logits is not None: + gate_logits = self.to_gate_logits(x) # (B, T, H) + b, t, _ = out.shape + out = out.view(b, t, self.heads, self.dim_head) + gates = 2.0 * torch.sigmoid(gate_logits) # zero-init -> identity + out = out * gates.unsqueeze(-1) + out = out.view(b, t, self.heads * self.dim_head) + + return self.to_out(out) + + +class LTXVCrossAttentionPatch: + def __init__(self, context, nag_scale, nag_alpha, nag_tau, inplace=True): + self.nag_context = context + self.nag_scale = nag_scale + self.nag_alpha = nag_alpha + self.nag_tau = nag_tau + self.inplace = inplace + + def __get__(self, obj, objtype=None): + # Create bound method with stored parameters + def wrapped_attention(self_module, *args, **kwargs): + self_module.nag_context = self.nag_context + self_module.nag_scale = self.nag_scale + self_module.nag_alpha = self.nag_alpha + self_module.nag_tau = self.nag_tau + self_module.inplace = self.inplace + + return ltxv_crossattn_forward_nag(self_module, *args, **kwargs) + return types.MethodType(wrapped_attention, obj) + +class LTX2_NAG(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTX2_NAG", + display_name="LTX2 NAG", + category="KJNodes/ltxv", + description="https://github.com/ChenDarYen/Normalized-Attention-Guidance", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + io.Float.Input("nag_scale", default=11.0, min=0.0, max=100.0, step=0.001, tooltip="Strength of negative guidance effect"), + io.Float.Input("nag_alpha", default=0.25, min=0.0, max=1.0, step=0.001, tooltip="Mixing coefficient in that controls the balance between the normalized guided representation and the original positive representation."), + io.Float.Input("nag_tau", default=2.5, min=0.0, max=10.0, step=0.001, tooltip="Clipping threshold that controls how much the guided attention can deviate from the positive attention."), + io.Conditioning.Input("nag_cond_video", optional=True), + io.Conditioning.Input("nag_cond_audio", optional=True), + io.Boolean.Input("inplace", default=True, optional=True, tooltip="If true, modifies tensors in place to save memory. Leads to different numerical results which may change the output slightly."), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, nag_scale, nag_alpha, nag_tau, nag_cond_video=None, nag_cond_audio=None, inplace=True) -> io.NodeOutput: + if nag_scale == 0: + return io.NodeOutput(model) + + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = model.model.manual_cast_dtype + if dtype is None: + dtype = model.model.diffusion_model.dtype + + model_clone = model.clone() + + diffusion_model = model_clone.get_model_object("diffusion_model") + img_dim = diffusion_model.inner_dim + audio_dim = diffusion_model.audio_inner_dim + + context_video = context_audio = None + + if nag_cond_video is not None: + context_video = nag_cond_video[0][0].to(device, dtype) + vid_split = getattr(diffusion_model, "cross_attention_dim", None) + if vid_split is not None and context_video.shape[-1] == vid_split + diffusion_model.audio_cross_attention_dim: + context_video = context_video[:, :, :vid_split] + if diffusion_model.caption_proj_before_connector and diffusion_model.caption_projection_first_linear: + diffusion_model.caption_projection.to(device) + context_video = diffusion_model.caption_projection(context_video) + diffusion_model.caption_projection.to(offload_device) + if hasattr(diffusion_model, "video_embeddings_connector"): + diffusion_model.video_embeddings_connector.to(device) + context_video = diffusion_model.video_embeddings_connector(context_video)[0] + diffusion_model.video_embeddings_connector.to(offload_device) + context_video = context_video.view(1, -1, img_dim) + for idx, block in enumerate(diffusion_model.transformer_blocks): + patched_attn2 = LTXVCrossAttentionPatch(context_video, nag_scale, nag_alpha, nag_tau, inplace=inplace).__get__(block.attn2, block.__class__) + model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.attn2.forward", patched_attn2) + + if nag_cond_audio is not None and diffusion_model.audio_caption_projection is not None: + context_audio = nag_cond_audio[0][0].to(device, dtype) + vid_split = getattr(diffusion_model, "cross_attention_dim", None) + if vid_split is not None and context_audio.shape[-1] == vid_split + diffusion_model.audio_cross_attention_dim: + context_audio = context_audio[:, :, vid_split:] + if diffusion_model.caption_proj_before_connector and diffusion_model.caption_projection_first_linear: + diffusion_model.audio_caption_projection.to(device) + context_audio = diffusion_model.audio_caption_projection(context_audio) + diffusion_model.audio_caption_projection.to(offload_device) + if hasattr(diffusion_model, "audio_embeddings_connector"): + diffusion_model.audio_embeddings_connector.to(device) + context_audio = diffusion_model.audio_embeddings_connector(context_audio)[0] + diffusion_model.audio_embeddings_connector.to(offload_device) + context_audio = context_audio.view(1, -1, audio_dim) + for idx, block in enumerate(diffusion_model.transformer_blocks): + patched_audio_attn2 = LTXVCrossAttentionPatch(context_audio, nag_scale, nag_alpha, nag_tau, inplace=inplace).__get__(block.audio_attn2, block.__class__) + model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.audio_attn2.forward", patched_audio_attn2) + + return io.NodeOutput(model_clone) + + +def ffn_chunked_forward(self, x): + if x.shape[1] > self.dim_threshold: + chunk_size = x.shape[1] // self.num_chunks + for i in range(self.num_chunks): + start_idx = i * chunk_size + end_idx = (i + 1) * chunk_size if i < self.num_chunks - 1 else x.shape[1] + x[:, start_idx:end_idx] = self.net(x[:, start_idx:end_idx]) + return x + else: + return self.net(x) + +class LTXVffnChunkPatch: + def __init__(self, num_chunks, dim_threshold=4096): + self.num_chunks = num_chunks + self.dim_threshold = dim_threshold + + def __get__(self, obj, objtype=None): + def wrapped_forward(self_module, *args, **kwargs): + self_module.num_chunks = self.num_chunks + self_module.dim_threshold = self.dim_threshold + return ffn_chunked_forward(self_module, *args, **kwargs) + return types.MethodType(wrapped_forward, obj) + +class LTXVChunkFeedForward(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVChunkFeedForward", + display_name="LTXV Chunk FeedForward", + category="KJNodes/ltxv", + description="EXPERIMENTAL AND MAY CHANGE THE MODEL OUTPUT!! Chunks feedforward activations to reduce peak VRAM usage.", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + io.Int.Input("chunks", default=2, min=1, max=100, step=1, tooltip="Number of chunks to split the feedforward activations into to reduce peak VRAM usage."), + io.Int.Input("dim_threshold", default=4096, min=0, max=16384, step=256, tooltip="Dimension threshold above which to apply chunking."), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, chunks, dim_threshold) -> io.NodeOutput: + if chunks == 1: + return io.NodeOutput(model) + + model_clone = model.clone() + diffusion_model = model_clone.get_model_object("diffusion_model") + + for idx, block in enumerate(diffusion_model.transformer_blocks): + patched_attn2 = LTXVffnChunkPatch(chunks, dim_threshold).__get__(block.ff, block.__class__) + model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.ff.forward", patched_attn2) + + return io.NodeOutput(model_clone) + + + +#borrowed VideoHelperSuite https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite/blob/main/videohelpersuite/latent_preview.py +import server +from threading import Thread +import torch.nn.functional as F +import time +import struct +from PIL import Image +from io import BytesIO +serv = server.PromptServer.instance + +class WrappedPreviewer(): + def __init__(self, latent_rgb_factors, latent_rgb_factors_bias, rate=8, taeltx=None): + self.first_preview = True + self.taeltx = taeltx + self.last_time = 0 + self.c_index = 0 + self.rate = rate + self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu").transpose(0, 1) + self.latent_rgb_factors_bias = torch.tensor(latent_rgb_factors_bias, device="cpu") if latent_rgb_factors_bias is not None else None + + def decode_latent_to_preview_image(self, preview_format, x0): + if x0.ndim == 5: + #Keep batch major + x0 = x0.movedim(2,1) + x0 = x0.reshape((-1,)+x0.shape[-3:]) + num_images = x0.size(0) + new_time = time.time() + num_previews = int((new_time - self.last_time) * self.rate) + self.last_time = self.last_time + num_previews/self.rate + if num_previews > num_images: + num_previews = num_images + elif num_previews <= 0: + return None + if self.first_preview: + self.first_preview = False + serv.send_sync('VHS_latentpreview', {'length':num_images, 'rate': self.rate, 'id': serv.last_node_id}) + self.last_time = new_time + 1/self.rate + if self.c_index + num_previews > num_images: + x0 = x0.roll(-self.c_index, 0)[:num_previews] + else: + x0 = x0[self.c_index:self.c_index + num_previews] + Thread(target=self.process_previews, args=(x0, self.c_index, + num_images)).run() + self.c_index = (self.c_index + num_previews) % num_images + return None + def process_previews(self, image_tensor, ind, leng): + max_size = 512 + min_size = 256 + image_tensor = self.decode_latent_to_preview(image_tensor) + + if image_tensor.size(1) < min_size or image_tensor.size(2) < min_size: + image_tensor = F.interpolate(image_tensor.movedim(-1,0), scale_factor=4, mode='nearest').movedim(0,-1) + + if image_tensor.size(1) > max_size or image_tensor.size(2) > max_size: + image_tensor = image_tensor.movedim(-1,0) + if image_tensor.size(2) < image_tensor.size(3): + height = (max_size * image_tensor.size(2)) // image_tensor.size(3) + image_tensor = F.interpolate(image_tensor, (height,max_size), mode='nearest') + else: + width = (max_size * image_tensor.size(3)) // image_tensor.size(2) + image_tensor = F.interpolate(image_tensor, (max_size, width), mode='nearest') + image_tensor = image_tensor.movedim(0,-1) + + previews_ubyte = (image_tensor.clamp(0, 1).mul(0xFF)).to(device="cpu", dtype=torch.uint8) + + # Send VHS preview + for preview in previews_ubyte: + i = Image.fromarray(preview.numpy()) + message = BytesIO() + message.write((1).to_bytes(length=4, byteorder='big')*2) + message.write(ind.to_bytes(length=4, byteorder='big')) + message.write(struct.pack('16p', (serv.last_node_id or "").encode('ascii'))) + i.save(message, format="JPEG", quality=95, compress_level=1) + #NOTE: send sync already uses call_soon_threadsafe + serv.send_sync(server.BinaryEventTypes.PREVIEW_IMAGE, + message.getvalue(), serv.client_id) + if self.taeltx is not None: + ind = (ind + 1) % ((leng-1) * 8 + 1) + else: + ind = (ind + 1) % leng + + def decode_latent_to_preview(self, x0): + if self.taeltx is not None: + x0 = x0.unsqueeze(0).to(dtype=self.taeltx.first_stage_model.decoder[1].weight.dtype, device=device) + x_sample = self.taeltx.first_stage_model.decode(x0)[0].permute(1, 2, 3, 0) + return x_sample + else: + self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device) + if self.latent_rgb_factors_bias is not None: + self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to(dtype=x0.dtype, device=x0.device) + latent_image = F.linear(x0.movedim(1, -1), self.latent_rgb_factors, + bias=self.latent_rgb_factors_bias) + + #low = latent_image.quantile(0.01) + #high = latent_image.quantile(0.99) + #latent_image = ((latent_image - low) / (high - low)).clamp(0, 1) + latent_image = torch.sigmoid(latent_image) + return latent_image + + +def get_ltx_rgb_factors(is_23): + """Return (latent_rgb_factors, latent_rgb_factors_bias) for LTX (is_23=False) or + LTX2/LTXAV (is_23=True). Exposed so other nodes can build a WrappedPreviewer with + the right LTX factors without duplicating the 128-row tables.""" + return _get_ltx_rgb_factors_impl(is_23) + + +def _get_ltx_rgb_factors_impl(is_23): + if not is_23: + latent_rgb_factors = [ + [ 0.0350, 0.0159, 0.0132], + [ 0.0025, -0.0021, -0.0003], + [ 0.0286, 0.0028, 0.0020], + [ 0.0280, -0.0114, -0.0202], + [-0.0186, 0.0073, 0.0092], + [ 0.0027, 0.0097, -0.0113], + [-0.0069, -0.0032, -0.0024], + [-0.0323, -0.0370, -0.0457], + [ 0.0174, 0.0164, 0.0106], + [-0.0097, 0.0061, 0.0035], + [-0.0130, -0.0042, -0.0012], + [-0.0102, -0.0002, -0.0091], + [-0.0025, 0.0063, 0.0161], + [ 0.0003, 0.0037, 0.0108], + [ 0.0152, 0.0082, 0.0143], + [ 0.0317, 0.0203, 0.0312], + [-0.0092, -0.0233, -0.0119], + [-0.0405, -0.0226, -0.0023], + [ 0.0376, 0.0397, 0.0352], + [ 0.0171, -0.0043, -0.0095], + [ 0.0482, 0.0341, 0.0213], + [ 0.0031, -0.0046, -0.0018], + [-0.0486, -0.0383, -0.0294], + [-0.0071, -0.0272, -0.0123], + [ 0.0320, 0.0218, 0.0289], + [ 0.0327, 0.0088, -0.0116], + [-0.0098, -0.0240, -0.0111], + [ 0.0094, -0.0116, 0.0021], + [ 0.0309, 0.0092, 0.0165], + [-0.0065, -0.0077, -0.0107], + [ 0.0179, 0.0114, 0.0038], + [-0.0018, -0.0030, -0.0026], + [-0.0002, 0.0076, -0.0029], + [-0.0131, -0.0059, -0.0170], + [ 0.0055, 0.0066, -0.0038], + [ 0.0154, 0.0063, 0.0090], + [ 0.0186, 0.0175, 0.0188], + [-0.0166, -0.0381, -0.0428], + [ 0.0121, 0.0015, -0.0153], + [ 0.0118, 0.0050, 0.0019], + [ 0.0125, 0.0259, 0.0231], + [ 0.0046, 0.0130, 0.0081], + [ 0.0271, 0.0250, 0.0250], + [-0.0054, -0.0347, -0.0326], + [-0.0438, -0.0262, -0.0228], + [-0.0191, -0.0256, -0.0173], + [-0.0205, -0.0058, 0.0042], + [ 0.0404, 0.0434, 0.0346], + [-0.0242, -0.0177, -0.0146], + [ 0.0161, 0.0223, 0.0168], + [-0.0240, -0.0320, -0.0299], + [-0.0019, 0.0043, 0.0008], + [-0.0060, -0.0133, -0.0244], + [-0.0048, -0.0225, -0.0167], + [ 0.0267, 0.0133, 0.0152], + [ 0.0222, 0.0167, 0.0028], + [ 0.0015, -0.0062, 0.0013], + [-0.0241, -0.0178, -0.0079], + [ 0.0040, -0.0081, -0.0097], + [-0.0064, 0.0133, -0.0011], + [-0.0204, -0.0231, -0.0304], + [ 0.0011, -0.0011, 0.0145], + [-0.0283, -0.0259, -0.0260], + [ 0.0038, 0.0171, -0.0029], + [ 0.0637, 0.0424, 0.0409], + [ 0.0092, 0.0163, 0.0188], + [ 0.0082, 0.0055, -0.0179], + [-0.0177, -0.0286, -0.0147], + [ 0.0171, 0.0242, 0.0398], + [-0.0129, 0.0095, -0.0071], + [-0.0154, 0.0036, 0.0128], + [-0.0081, -0.0009, 0.0118], + [-0.0067, -0.0178, -0.0230], + [-0.0022, -0.0125, -0.0003], + [-0.0032, -0.0039, -0.0022], + [-0.0005, -0.0127, -0.0131], + [-0.0143, -0.0157, -0.0165], + [-0.0262, -0.0263, -0.0270], + [ 0.0063, 0.0127, 0.0178], + [ 0.0092, 0.0133, 0.0150], + [-0.0106, -0.0068, 0.0032], + [-0.0214, -0.0022, 0.0171], + [-0.0104, -0.0266, -0.0362], + [ 0.0021, 0.0048, -0.0005], + [ 0.0345, 0.0431, 0.0402], + [-0.0275, -0.0110, -0.0195], + [ 0.0203, 0.0251, 0.0224], + [ 0.0016, -0.0037, -0.0094], + [ 0.0241, 0.0198, 0.0114], + [-0.0003, 0.0027, 0.0141], + [ 0.0012, -0.0052, -0.0084], + [ 0.0057, -0.0028, -0.0163], + [-0.0488, -0.0545, -0.0509], + [-0.0076, -0.0025, -0.0014], + [-0.0249, -0.0142, -0.0367], + [ 0.0136, 0.0041, 0.0135], + [ 0.0007, 0.0034, -0.0053], + [-0.0068, -0.0109, 0.0029], + [ 0.0006, -0.0237, -0.0094], + [-0.0149, -0.0177, -0.0131], + [-0.0105, 0.0039, 0.0216], + [ 0.0242, 0.0200, 0.0180], + [-0.0339, -0.0153, -0.0195], + [ 0.0104, 0.0151, 0.0120], + [-0.0043, 0.0089, 0.0047], + [ 0.0157, -0.0030, 0.0008], + [ 0.0126, 0.0102, -0.0040], + [ 0.0040, 0.0114, 0.0137], + [ 0.0423, 0.0473, 0.0436], + [-0.0128, -0.0066, -0.0152], + [-0.0337, -0.0087, -0.0026], + [-0.0052, 0.0235, 0.0291], + [ 0.0079, 0.0154, 0.0260], + [-0.0539, -0.0377, -0.0358], + [-0.0188, 0.0062, -0.0035], + [-0.0186, 0.0041, -0.0083], + [ 0.0045, -0.0049, 0.0053], + [ 0.0172, 0.0071, 0.0042], + [-0.0003, -0.0078, -0.0096], + [-0.0209, -0.0132, -0.0135], + [-0.0074, 0.0017, 0.0099], + [-0.0038, 0.0070, 0.0014], + [-0.0013, -0.0017, 0.0073], + [ 0.0030, 0.0105, 0.0105], + [ 0.0154, -0.0168, -0.0235], + [-0.0108, -0.0038, 0.0047], + [-0.0298, -0.0347, -0.0436], + [-0.0206, -0.0189, -0.0139] + ] + latent_rgb_factors_bias = [0.2796, 0.1101, -0.0047] + else: + latent_rgb_factors = [[0.002269406570121646, -0.02110900916159153, -0.009850316680967808], [-0.016038373112678528, -0.012462412007153034, -0.01112896017730236], [0.025274179875850677, 0.011209743097424507, 0.025426799431443214], [0.04690725728869438, 0.041542328894138336, 0.03568895906209946], [-0.02388044260442257, -0.0018645941745489836, 0.01858334057033062], [0.03720448538661003, 0.0220357533544302, 0.027937663719058037], [-0.07273884862661362, -0.09326262027025223, -0.11579664051532745], [-0.063837431371212, 0.00026216846890747547, 0.03042735904455185], [0.02903873845934868, 0.042082373052835464, 0.030649805441498756], [0.03777873516082764, 0.0322984978556633, -0.005671461578458548], [-0.0075670829974114895, -0.012113905511796474, -0.01638956367969513], [0.026524530723690987, 0.060518112033605576, 0.059549521654844284], [0.10093028098344803, 0.10073262453079224, 0.0505094900727272], [0.03725508227944374, 0.015382086858153343, 0.005786076188087463], [-0.03139607608318329, -0.01690264232456684, -0.0013519978383556008], [-0.027200624346733093, -0.02517341822385788, -0.008874989114701748], [0.024963486939668655, 0.04293748363852501, 0.05582639202475548], [-0.0364827960729599, -0.026975594460964203, -0.021950015798211098], [0.027655167505145073, 0.025136707350611687, 0.043967027217149734], [0.035822272300720215, 0.013104500249028206, 0.01113432738929987], [0.05353763327002525, 0.013606574386358261, -0.018720127642154694], [-0.013587888330221176, -0.01689346879720688, -0.027842802926898003], [0.059415675699710846, 0.03734271228313446, 0.04562298208475113], [-0.02946414425969124, -0.038338612765073776, 0.001805233070626855], [0.03921474143862724, 0.0651894062757492, 0.10681862384080887], [-0.00744189927354455, 0.007951526902616024, 0.020728807896375656], [-0.04038553684949875, -0.05215264856815338, -0.07213657349348068], [-0.004655141849070787, 0.01305423304438591, 0.026104029268026352], [0.03434251993894577, 0.018448110669851303, 0.013096392154693604], [0.0022075253073126078, -0.0011812079465016723, 0.0002940484555438161], [-0.00043441299931146204, 0.02366728149354458, 0.035889431834220886], [-0.030657343566417694, -0.024926183745265007, -0.012355240061879158], [-0.018955843523144722, -0.017360301688313484, -0.008214764297008514], [-0.01113052573055029, -0.01201171800494194, -0.002986249281093478], [0.018902746960520744, 0.01758778840303421, 0.026414571329951286], [-0.019977254793047905, -0.01605399139225483, -0.019136475399136543], [-0.00300968368537724, -0.017609693109989166, -0.013655650429427624], [0.0022096361499279737, 0.017998533323407173, 0.01815750263631344], [0.05186990648508072, 0.03285299986600876, 0.016072165220975876], [0.012626334093511105, 0.0013884707586839795, -0.012077193707227707], [-0.0037861645687371492, -0.013902144506573677, -0.01911942847073078], [-0.014163163490593433, -0.00513274222612381, -0.014303527772426605], [-0.010461323894560337, 0.009658926166594028, 0.01644069515168667], [-0.008665377274155617, 0.002501955023035407, -0.009703717194497585], [-0.03404829278588295, -0.02546044997870922, -0.014914450235664845], [0.04997691139578819, 0.06592527031898499, 0.073111392557621], [0.027394814416766167, 0.024555068463087082, 0.019957970827817917], [-0.027501430362462997, -0.01673700101673603, -0.03089248389005661], [-0.018696032464504242, -0.0020940247923135757, 0.015244065783917904], [-0.0062704551964998245, -0.0067006442695856094, -0.007532030809670687], [0.014871004968881607, 0.009914354421198368, 0.020960720255970955], [0.03662937879562378, 0.04413224756717682, 0.04220828413963318], [-0.011242181062698364, -0.013539309613406658, -0.016438307240605354], [-0.014854325912892818, 0.0038217694964259863, -0.002461288822814822], [-0.014826249331235886, 0.0009719038498587906, -0.012078499421477318], [-0.029396841302514076, -0.01432017982006073, 0.013018904253840446], [0.02755064144730568, 0.028369395062327385, 0.01640605367720127], [0.12049165368080139, 0.1395745575428009, 0.14566579461097717], [0.019721267744898796, 0.009739740751683712, 0.0023876908235251904], [-0.007320966571569443, 0.0065013207495212555, 0.01603059470653534], [0.007391378283500671, -0.0073603675700724125, -0.01770283281803131], [0.02984853833913803, 0.012391146272420883, 0.010563627816736698], [-0.013479884713888168, -0.008637298829853535, -0.013457189314067364], [0.04127075523138046, 0.03032625839114189, 0.024770958349108696], [-0.06524652987718582, -0.012209279462695122, 0.02087211236357689], [-0.1179763451218605, -0.060323599725961685, -0.07592175155878067], [-0.07122819870710373, -0.04385707899928093, -0.022124603390693665], [-0.04682473465800285, -0.022610662505030632, -0.010107148438692093], [-0.0054328180849552155, -0.010368981398642063, -0.008167334832251072], [0.029181398451328278, 0.030588403344154358, 0.028090540319681168], [0.016619984060525894, 0.004931286443024874, -0.006450849585235119], [0.01035264041274786, 0.002237115055322647, 0.0013903985964134336], [-0.04313831403851509, -0.061772625893354416, -0.08946335315704346], [0.0150345079600811, 0.007781678810715675, 0.0011013159528374672], [-0.013585779815912247, 0.008117705583572388, 0.020367907360196114], [-0.172962948679924, -0.16406646370887756, -0.1668281853199005], [0.0083833709359169, 0.0015236001927405596, -0.01731627807021141], [0.021939430385828018, 0.018004458397626877, 0.014768349006772041], [0.008083095774054527, -0.013463049195706844, -0.022061636671423912], [0.024328550323843956, 0.0128010343760252, 0.014966367743909359], [0.05850301682949066, 0.027980001643300056, 0.02225641906261444], [0.09690416604280472, 0.06929530203342438, 0.03253814950585365], [0.048208240419626236, 0.025294817984104156, 0.023508133366703987], [-0.026432134211063385, -0.040383171290159225, -0.03950457274913788], [-0.021598653867840767, -0.017070941627025604, -0.010933087207376957], [0.011645167134702206, 0.002806191798299551, 0.003779367310926318], [0.10478592664003372, 0.08954174816608429, 0.06555330753326416], [0.015151776373386383, -0.016160616651177406, -0.024905217811465263], [0.019659176468849182, 0.008487952873110771, 0.002426224760711193], [-0.05173315480351448, -0.026337839663028717, -0.02127116546034813], [0.016987523064017296, 0.006270893849432468, 0.0015798212261870503], [0.007938026450574398, -0.005250005517154932, -0.020408453419804573], [0.013017759658396244, 0.01654384844005108, 0.04163840040564537], [-0.009886542335152626, -0.026848411187529564, -0.03070281818509102], [0.01108171883970499, 0.01827266439795494, -0.007332107983529568], [-0.0285995751619339, -0.031727731227874756, -0.03370537981390953], [0.005299570970237255, 0.05678633600473404, 0.02825017087161541], [-0.055322226136922836, -0.09084303677082062, -0.12999044358730316], [0.01844066195189953, 0.031044499948620796, 0.021148500964045525], [-0.004471115302294493, 0.005830412730574608, 0.00911418441683054], [-0.04053843766450882, -0.016424428671598434, -0.0010634599020704627], [0.03858831524848938, 0.007309338077902794, -0.005618985276669264], [0.01423253770917654, -0.0055681923404335976, 3.394074519746937e-05], [0.11455483734607697, 0.14653916656970978, 0.1488018035888672], [-0.005231931805610657, -0.0033921014983206987, -0.000995257287286222], [0.01449565589427948, 0.019586293026804924, 0.04565812274813652], [-0.005179048050194979, -0.011201606132090092, -0.0008710073889233172], [-0.015361929312348366, 0.00778581015765667, -0.008238887414336205], [-0.1147838830947876, -0.09109023958444595, -0.050579313188791275], [0.09037500619888306, 0.09597006440162659, 0.10811734944581985], [0.001873677596449852, -0.01772197335958481, -0.07681205868721008], [-0.020383257418870926, -0.016072455793619156, -0.01077069528400898], [-0.060444317758083344, -0.05499502643942833, -0.06153025105595589], [-0.016717270016670227, 0.026493264362215996, 0.021835654973983765], [0.008203534409403801, 0.00418612826615572, 0.013867748901247978], [0.0789225772023201, 0.05467747151851654, 0.016568133607506752], [-0.15149451792240143, -0.1526806503534317, -0.14325062930583954], [0.00538366474211216, 0.010192245244979858, -0.00449327751994133], [-0.004906965419650078, -0.005569908302277327, -0.02096559666097164], [0.024530155584216118, 0.010962833650410175, 0.0034586559049785137], [0.03551010414958, 0.017310436815023422, 0.007064413744956255], [0.11111932247877121, 0.09825586527585983, 0.08827318251132965], [-0.051722846925258636, -0.047595202922821045, -0.03763044252991676], [-0.02975175902247429, -0.02153967320919037, -0.021425534039735794], [-0.03462936729192734, -0.025198571383953094, -0.017322326079010963], [-0.016921017318964005, -0.012419789098203182, -0.0154880927875638], [-0.08035065978765488, -0.08451078832149506, -0.09623870998620987], [-0.03870908170938492, -0.04211008921265602, -0.04383759945631027]] + latent_rgb_factors_bias = [-0.6957847476005554, -0.7276281118392944, -0.7405748963356018] + return latent_rgb_factors, latent_rgb_factors_bias + + +def prepare_callback(model, steps, x0_output_dict=None, shape=None, latent_upscale_model=None, vae=None, rate=8, taeltx=False, num_keyframes=0, is_23=False): + latent_rgb_factors, latent_rgb_factors_bias = get_ltx_rgb_factors(is_23) + preview_format = "JPEG" + if preview_format not in ["JPEG", "PNG"]: + preview_format = "JPEG" + + previewer = WrappedPreviewer(latent_rgb_factors, latent_rgb_factors_bias, rate=rate, taeltx=vae if taeltx else None) + + pbar = comfy.utils.ProgressBar(steps) + def callback(step, x0, x, total_steps): + if x0 is not None and shape is not None: + cut = math.prod(shape[1:]) + x0 = x0[:, :, :cut].reshape([x0.shape[0]] + list(shape)[1:]) + + if num_keyframes > 0: + x0 = x0[:, :, :-num_keyframes] + + if latent_upscale_model is not None: + x0 = vae.first_stage_model.per_channel_statistics.un_normalize(x0) + x0 = latent_upscale_model(x0.to(torch.bfloat16)) + x0 = vae.first_stage_model.per_channel_statistics.normalize(x0) + + preview_bytes = None + if previewer: + preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0) + pbar.update_absolute(step + 1, total_steps, preview_bytes) + return callback + +class OuterSampleCallbackWrapper: + def __init__(self, latent_upscale_model=None, vae=None, preview_rate=8, taeltx=False): + self.latent_upscale_model = latent_upscale_model + self.vae = vae + self.preview_rate = preview_rate + self.taeltx = taeltx + self.x0_output = {} + + def __call__(self, executor, noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes): + guider = executor.class_obj + diffusion_model = guider.model_patcher.model.diffusion_model + is_23 = not diffusion_model.caption_projection_first_linear + + original_callback = callback + if self.latent_upscale_model is not None: + self.latent_upscale_model.to(device) + if self.vae is not None and self.taeltx: + self.vae.first_stage_model.to(device) + + num_keyframes = 0 + if 'positive' in guider.conds and len(guider.conds['positive']) > 0: + keyframe_idxs = guider.conds['positive'][0].get('keyframe_idxs') + if keyframe_idxs is not None: + num_keyframes = len(torch.unique(keyframe_idxs[0, 0, :, 0])) + + new_callback = prepare_callback(guider.model_patcher, len(sigmas) -1, shape=latent_shapes[0] if len(latent_shapes) > 1 else None, + x0_output_dict=self.x0_output, latent_upscale_model=self.latent_upscale_model, vae=self.vae, rate=self.preview_rate, taeltx=self.taeltx, num_keyframes=num_keyframes, is_23=is_23) + # Wrapper that calls both callbacks + def combined_callback(step, x0, x, total_steps): + new_callback(step, x0, x, total_steps) + if original_callback is not None: + original_callback(step, x0, x, total_steps) + out = executor(noise, latent_image, sampler, sigmas, denoise_mask, combined_callback, disable_pbar, seed, latent_shapes=latent_shapes) + if self.latent_upscale_model is not None: + self.latent_upscale_model.to(mm.unet_offload_device()) + return out + +class LTX2SamplingPreviewOverride(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LTX2SamplingPreviewOverride", + display_name="LTX2 Sampling Preview Override", + description="Overrides the LTX2 preview sampling preview function, temporary measure until previews are in comfy core", + category="KJNodes/ltxv", + is_experimental=True, + inputs=[ + io.Model.Input("model", tooltip="The model to add preview override to."), + io.Int.Input("preview_rate", default=8, min=1, max=60, step=1, tooltip="Preview frame rate."), + io.LatentUpscaleModel.Input("latent_upscale_model", optional=True, tooltip="Optional upscale model to use for higher resolution previews."), + io.Vae.Input("vae", optional=True, tooltip="VAE model to use normalizing the latents for the upscale model."), + ], + outputs=[ + io.Model.Output(tooltip="The model with Sampling Preview Override."), + ], + ) + + @classmethod + def execute(cls, model, preview_rate, latent_upscale_model=None, vae=None) -> io.NodeOutput: + model = model.clone() + taeltx = False + if vae is not None: + if vae.first_stage_model.__class__.__name__ == "TAEHV": + taeltx = True + latent_upscale_model=None + model.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, "sampling_preview", OuterSampleCallbackWrapper(latent_upscale_model, vae, preview_rate, taeltx)) + return io.NodeOutput(model) + + +# based on https://github.com/Lightricks/ComfyUI-LTXVideo/blob/cd5d371518afb07d6b3641be8012f644f25269fc/easy_samplers.py#L916 +class OuterSampleAudioNormalizationWrapper: + def __init__(self, audio_normalization_factors): + self.audio_normalization_factors = audio_normalization_factors + + def __call__(self, executor, noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes): + guider = executor.class_obj + ltxav = guider.model_patcher.model.diffusion_model + + x0_output = {} + self.total_steps = sigmas.shape[-1] - 1 + pbar = comfy.utils.ProgressBar(self.total_steps) + self.full_step = 0 + + previewer = latent_preview.get_previewer(guider.model_patcher.load_device, guider.model_patcher.model.latent_format) + def custom_callback(step, x0, x, total_steps): + if x0_output is not None: + x0_output["x0"] = x0 + + preview_bytes = None + if previewer: + preview_bytes = previewer.decode_latent_to_preview_image("JPEG", x0) + self.full_step += 1 + pbar.update_absolute(self.full_step, self.total_steps, preview_bytes) + + callback = custom_callback + + audio_normalization_factors = self.audio_normalization_factors.strip().split(",") + audio_normalization_factors = [float(factor) for factor in audio_normalization_factors] + + # Extend normalization factors to match the length of sigmas + sigmas_len = self.total_steps + if len(audio_normalization_factors) < sigmas_len and len(audio_normalization_factors) > 0: + audio_normalization_factors.extend([audio_normalization_factors[-1]] * (sigmas_len - len(audio_normalization_factors))) + + # Calculate indices where both normalization factors are not 1.0 + sampling_split_indices = [i + 1 for i, a in enumerate(audio_normalization_factors) if a != 1.0] + + # Split sigmas according to sampling_split_indices + def split_by_indices(arr, indices): + """ + Splits arr into chunks according to indices (split points). + Indices are treated as starting a new chunk at each index in the list. + """ + if not indices: + return [arr] + split_points = sorted(set(indices)) + chunks = [] + prev = 0 + for idx in split_points: + if prev < idx: + chunks.append(arr[prev : idx + 1]) + prev = idx + if prev < len(arr): + chunks.append(arr[prev:]) + return chunks + + sigmas_chunks = split_by_indices(sigmas, sampling_split_indices) + + i = 0 + for sigmas_chunk in sigmas_chunks: + i += len(sigmas_chunk) - 1 + latent_image = executor(noise, latent_image, sampler, sigmas_chunk, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes) + + if "x0" in x0_output: + latent_image = guider.model_patcher.model.process_latent_out(x0_output["x0"]) + + if i - 1 < len(audio_normalization_factors): + vx, ax = ltxav.separate_audio_and_video_latents(comfy.utils.unpack_latents(latent_image, latent_shapes), None) + if denoise_mask is not None: + audio_mask = ltxav.separate_audio_and_video_latents(comfy.utils.unpack_latents(denoise_mask, latent_shapes), None)[1] + ax = ax * audio_mask * audio_normalization_factors[i - 1] + ax * (1 - audio_mask) + else: + ax = ax * audio_normalization_factors[i - 1] + latent_image = comfy.utils.pack_latents(ltxav.recombine_audio_and_video_latents(vx, ax))[0] + + print("After %d steps, the audio latent was normalized by %f" % (i, audio_normalization_factors[i - 1])) + + return latent_image + + +class LTX2AudioLatentNormalizingSampling(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="LTX2AudioLatentNormalizingSampling", + display_name="LTX2 Audio Latent Normalizing Sampling", + description="Improves LTX2 generated audio quality by normalizing audio latents at specified sampling steps.", + category="KJNodes/ltxv", + is_experimental=True, + inputs=[ + io.Model.Input("model", tooltip="The model to add preview override to."), + io.String.Input("audio_normalization_factors", default="1,1,0.25,1,1,0.25,1,1", tooltip="Comma-separated list of audio normalization factors to apply at each sampling step. For example, '1,1,0.25,1,1,0.25,1,1' will apply a factor of 0.25 at the 3rd and 6th steps."), + ], + outputs=[ + io.Model.Output(tooltip="The model with Audio Latent Normalizing Sampling."), + ], + ) + + @classmethod + def execute(cls, model, audio_normalization_factors) -> io.NodeOutput: + model = model.clone() + model.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, "ltx2_audio_normalization", OuterSampleAudioNormalizationWrapper(audio_normalization_factors)) + return io.NodeOutput(model) + + +class LTXVImgToVideoInplaceKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + options = [] + for num_images in range(1, 21): # 1 to 20 images + image_inputs = [] + for i in range(1, num_images + 1): + image_inputs.extend([ + io.Image.Input(f"image_{i}", optional=True, tooltip=f"Image {i} to insert into the video latent."), + io.Int.Input( + f"index_{i}", + default=0, + min=-9999, + max=9999, + step=1, + tooltip=f"Frame index for image {i} (in pixel space).", + optional=True, + ), + io.Float.Input(f"strength_{i}", default=1.0, min=0.0, max=1.0, step=0.01, tooltip=f"Strength for image {i}."), + ]) + options.append(io.DynamicCombo.Option( + key=str(num_images), + inputs=image_inputs + )) + + return io.Schema( + node_id="LTXVImgToVideoInplaceKJ", + category="KJNodes/ltxv", + description="Replaces video latent frames with the encoded input images, uses DynamicCombo which requires ComfyUI 0.8.1 and frontend 1.33.4 or later.", + inputs=[ + io.Vae.Input("vae", tooltip="Video VAE used to encode the images"), + io.Latent.Input("latent", tooltip="Video latent to insert images into"), + io.DynamicCombo.Input( + "num_images", + options=options, + display_name="Number of Images", + tooltip="Select how many images to insert", + ), + ], + outputs=[ + io.Latent.Output(display_name="latent", tooltip="The video latent with the images inserted and latent noise mask updated."), + ], + ) + + @classmethod + def execute(cls, vae, latent, num_images) -> io.NodeOutput: + + samples = latent["samples"].clone() + scale_factors = vae.downscale_index_formula + _, height_scale_factor, width_scale_factor = scale_factors + + batch, _, latent_frames, latent_height, latent_width = samples.shape + width = latent_width * width_scale_factor + height = latent_height * height_scale_factor + + # Get existing noise mask if present, otherwise create new one + if "noise_mask" in latent: + conditioning_latent_frames_mask = latent["noise_mask"].clone() + else: + conditioning_latent_frames_mask = torch.ones( + (batch, 1, latent_frames, 1, 1), + dtype=torch.float32, + device=samples.device, + ) + + # num_images is a dict containing the inputs from the selected option + # e.g., {'image_1': tensor, 'frame_idx_1': 0, 'strength_1': 1.0, 'image_2': tensor, 'frame_idx_2': 20, 'strength_2': 0.8, ...} + + image_keys = sorted([k for k in num_images.keys() if k.startswith('image_')]) + + for img_key in image_keys: + i = img_key.split('_')[1] + + image = num_images[f"image_{i}"] + if image is None: + continue + index = num_images.get(f"index_{i}") + if index is None: + continue + strength = num_images[f"strength_{i}"] + + if image.shape[1] != height or image.shape[2] != width: + pixels = comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + else: + pixels = image + encode_pixels = pixels[:, :, :, :3] + t = vae.encode(encode_pixels) + + # Convert pixel frame index to latent index + time_scale_factor = scale_factors[0] + + # Handle negative indexing in pixel space + pixel_frame_count = (latent_frames - 1) * time_scale_factor + 1 + if index < 0: + index = pixel_frame_count + index + + # Convert to latent index + latent_idx = index // time_scale_factor + + # Clamp to valid range + latent_idx = max(0, min(latent_idx, latent_frames - 1)) + + # Calculate end index, ensuring we don't exceed latent_frames + end_index = min(latent_idx + t.shape[2], latent_frames) + + # Replace samples at the specified index range + samples[:, :, latent_idx:end_index] = t[:, :, :end_index - latent_idx] + + # Update mask at the specified index range + conditioning_latent_frames_mask[:, :, latent_idx:end_index] = 1.0 - strength + + return io.NodeOutput({"samples": samples, "noise_mask": conditioning_latent_frames_mask}) + +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + +if HAS_TRITON: + @triton.jit + def _rms_norm_scale_shift_kernel( + X_ptr, Out_ptr, Scale_ptr, Shift_ptr, + N, scale_n_rows, # scale_n_rows: actual rows in scale/shift (handles broadcast) + eps: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_BF16: tl.constexpr, + HAS_BROADCAST: tl.constexpr, # True when scale_n_rows < n_rows + ): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_N) + mask = cols < N + off = row * N + cols + + x = tl.load(X_ptr + off, mask=mask, other=0.0).to(tl.float32) + + mean_sq = tl.sum(x * x, axis=0) / N + rrms = tl.math.rsqrt(mean_sq + eps) + x_norm = x * rrms + + # Fast path: no broadcast, scale row == x row + if HAS_BROADCAST: + scale_row = row % scale_n_rows + scale_off = scale_row * N + cols + else: + scale_off = off + scale = tl.load(Scale_ptr + scale_off, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(Shift_ptr + scale_off, mask=mask, other=0.0).to(tl.float32) + out = x_norm * (1.0 + scale) + shift + + if IS_BF16: + tl.store(Out_ptr + off, out.to(tl.bfloat16), mask=mask) + else: + tl.store(Out_ptr + off, out.to(tl.float16), mask=mask) + + @triton.jit + def _rope_qk_split_kernel( + Q_ptr, K_ptr, Cos_ptr, Sin_ptr, + H, T, D, # D = dim_head (full, even) + TH, # T * H, precomputed to avoid runtime multiply in index decomp + # q/k layout: [B, T, H*D] contiguous → q[b,t,h*D + d] + # cos/sin layout: [B, H, T, D//2] contiguous + IS_BF16: tl.constexpr, + BLOCK_HD: tl.constexpr, # >= D//2 + ): + # grid: (B*H*T,) + bht = tl.program_id(0) + t = bht % T + h = (bht // T) % H + b = bht // TH + + D_half = D // 2 + cols = tl.arange(0, BLOCK_HD) + mask = cols < D_half + + # offset into q/k tensor: [B, T, H*D] + qk_base = (b * T + t) * (H * D) + h * D + # offset into cos/sin: [B, H, T, D//2] + cs_base = (b * H * T + h * T + t) * D_half + + # load both halves of q and k + q_x = tl.load(Q_ptr + qk_base + cols, mask=mask, other=0.0).to(tl.float32) + q_y = tl.load(Q_ptr + qk_base + D_half + cols, mask=mask, other=0.0).to(tl.float32) + k_x = tl.load(K_ptr + qk_base + cols, mask=mask, other=0.0).to(tl.float32) + k_y = tl.load(K_ptr + qk_base + D_half + cols, mask=mask, other=0.0).to(tl.float32) + cos = tl.load(Cos_ptr + cs_base + cols, mask=mask, other=1.0).to(tl.float32) + sin = tl.load(Sin_ptr + cs_base + cols, mask=mask, other=0.0).to(tl.float32) + + # split RoPE: out_x = x*cos - y*sin, out_y = y*cos + x*sin + q_ox = q_x * cos - q_y * sin + q_oy = q_y * cos + q_x * sin + k_ox = k_x * cos - k_y * sin + k_oy = k_y * cos + k_x * sin + + if IS_BF16: + tl.store(Q_ptr + qk_base + cols, q_ox.to(tl.bfloat16), mask=mask) + tl.store(Q_ptr + qk_base + D_half + cols, q_oy.to(tl.bfloat16), mask=mask) + tl.store(K_ptr + qk_base + cols, k_ox.to(tl.bfloat16), mask=mask) + tl.store(K_ptr + qk_base + D_half + cols, k_oy.to(tl.bfloat16), mask=mask) + else: + tl.store(Q_ptr + qk_base + cols, q_ox.to(tl.float16), mask=mask) + tl.store(Q_ptr + qk_base + D_half + cols, q_oy.to(tl.float16), mask=mask) + tl.store(K_ptr + qk_base + cols, k_ox.to(tl.float16), mask=mask) + tl.store(K_ptr + qk_base + D_half + cols, k_oy.to(tl.float16), mask=mask) + + @triton.jit + def _rms_norm_dual_scale_shift_kernel( + X_ptr, Out1_ptr, Out2_ptr, + Scale1_ptr, Shift1_ptr, Scale2_ptr, Shift2_ptr, + N, scale_n_rows, + eps: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_BF16: tl.constexpr, + HAS_BROADCAST: tl.constexpr, + ): + """Compute rms_norm(x) once, write two differently-scaled outputs.""" + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_N) + mask = cols < N + off = row * N + cols + + x = tl.load(X_ptr + off, mask=mask, other=0.0).to(tl.float32) + mean_sq = tl.sum(x * x, axis=0) / N + x_norm = x * tl.math.rsqrt(mean_sq + eps) + + if HAS_BROADCAST: + scale_row = row % scale_n_rows + s_off = scale_row * N + cols + else: + s_off = off + + scale1 = tl.load(Scale1_ptr + s_off, mask=mask, other=0.0).to(tl.float32) + shift1 = tl.load(Shift1_ptr + s_off, mask=mask, other=0.0).to(tl.float32) + scale2 = tl.load(Scale2_ptr + s_off, mask=mask, other=0.0).to(tl.float32) + shift2 = tl.load(Shift2_ptr + s_off, mask=mask, other=0.0).to(tl.float32) + + out1 = x_norm * (1.0 + scale1) + shift1 + out2 = x_norm * (1.0 + scale2) + shift2 + + if IS_BF16: + tl.store(Out1_ptr + off, out1.to(tl.bfloat16), mask=mask) + tl.store(Out2_ptr + off, out2.to(tl.bfloat16), mask=mask) + else: + tl.store(Out1_ptr + off, out1.to(tl.float16), mask=mask) + tl.store(Out2_ptr + off, out2.to(tl.float16), mask=mask) + + +def fused_norm_scale_shift(x, scale, shift, eps=1e-6, use_triton=True): + if use_triton and HAS_TRITON and x.is_cuda: + orig_shape = x.shape + hidden = x.shape[-1] + x_2d = x.contiguous().reshape(-1, hidden) + n_rows = x_2d.shape[0] + # Flatten scale/shift without expanding — kernel handles broadcast via modulo + scale_2d = scale.contiguous().reshape(-1, hidden) + shift_2d = shift.contiguous().reshape(-1, hidden) + scale_n_rows = scale_2d.shape[0] + out = torch.empty_like(x_2d) + BLOCK_N = triton.next_power_of_2(hidden) + num_warps = min(max(BLOCK_N // 256, 1), 16) + _rms_norm_scale_shift_kernel[(n_rows,)]( + x_2d, out, scale_2d, shift_2d, + hidden, scale_n_rows, + eps=eps, BLOCK_N=BLOCK_N, + IS_BF16=(x.dtype == torch.bfloat16), + HAS_BROADCAST=(scale_n_rows < n_rows), + num_warps=num_warps, + ) + return out.view(orig_shape) + else: + return comfy.ldm.common_dit.rms_norm(x, eps=eps) * (1 + scale) + shift + + +def fused_rope_qk(q, k, freqs_cis, use_triton=True): + """Apply split RoPE to q and k in one fused kernel pass. + q, k: [B, T, H*D] contiguous + freqs_cis: (cos, sin, split_pe) where cos/sin: [B, H, T, D//2] + Falls back to original apply_rotary_emb if preconditions not met. + """ + + if not (use_triton and HAS_TRITON and q.is_cuda): + return _apply_rope(q, freqs_cis), _apply_rope(k, freqs_cis) + + cos, sin = freqs_cis[0], freqs_cis[1] + split_pe = freqs_cis[2] if len(freqs_cis) > 2 else False + + if not split_pe or cos.ndim != 4 or q.ndim != 3: + return _apply_rope(q, freqs_cis), _apply_rope(k, freqs_cis) + + B_cos, H, T_cos, D_half = cos.shape + D = D_half * 2 + if q.shape != (B_cos, T_cos, H * D) or k.shape != (B_cos, T_cos, H * D): + return _apply_rope(q, freqs_cis), _apply_rope(k, freqs_cis) + + q = q.contiguous() + k = k.contiguous() + cos_c = cos.contiguous() + sin_c = sin.contiguous() + + BLOCK_HD = triton.next_power_of_2(D_half) + num_warps = min(max(BLOCK_HD // 32, 1), 8) + _rope_qk_split_kernel[(B_cos * H * T_cos,)]( + q, k, cos_c, sin_c, + H, T_cos, D, + T_cos * H, + IS_BF16=(q.dtype == torch.bfloat16), + BLOCK_HD=BLOCK_HD, + num_warps=num_warps, + ) + return q, k + + +def fused_norm_dual_scale_shift(x, scale1, shift1, scale2, shift2, eps=1e-6, use_triton=True): + """RMS-norm x once, return two scaled outputs: (x_norm*(1+s1)+b1, x_norm*(1+s2)+b2).""" + if use_triton and HAS_TRITON and x.is_cuda: + orig_shape = x.shape + hidden = x.shape[-1] + x_2d = x.contiguous().reshape(-1, hidden) + n_rows = x_2d.shape[0] + scale1_2d = scale1.contiguous().reshape(-1, hidden) + shift1_2d = shift1.contiguous().reshape(-1, hidden) + scale2_2d = scale2.contiguous().reshape(-1, hidden) + shift2_2d = shift2.contiguous().reshape(-1, hidden) + scale_n_rows = scale1_2d.shape[0] + out1 = torch.empty_like(x_2d) + out2 = torch.empty_like(x_2d) + BLOCK_N = triton.next_power_of_2(hidden) + num_warps = min(max(BLOCK_N // 256, 1), 16) + _rms_norm_dual_scale_shift_kernel[(n_rows,)]( + x_2d, out1, out2, + scale1_2d, shift1_2d, scale2_2d, shift2_2d, + hidden, scale_n_rows, + eps=eps, BLOCK_N=BLOCK_N, + IS_BF16=(x.dtype == torch.bfloat16), + HAS_BROADCAST=(scale_n_rows < n_rows), + num_warps=num_warps, + ) + return out1.view(orig_shape), out2.view(orig_shape) + else: + x_norm = comfy.ldm.common_dit.rms_norm(x, eps=eps) + return x_norm * (1 + scale1) + shift1, x_norm * (1 + scale2) + shift2 + + +def _apply_text_cross_attention_patched( + self, x, context, attn, scale_shift_table, prompt_scale_shift_table, + timestep, prompt_timestep, attention_mask, transformer_options, +): + """Drop-in replacement for _apply_text_cross_attention with fused norm+scale+shift. + Patched onto the block instance so self._apply_text_cross_attention resolves here.""" + if self.cross_attention_adaln: + shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9)) + batch_size = x.shape[0] + shift_kv, scale_kv = ( + prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1) + ).unbind(dim=2) + attn_input = fused_norm_scale_shift(x, scale_q, shift_q, + use_triton=getattr(self, 'use_triton_kernels', True)) + del shift_q, scale_q + encoder_hidden_states = context * (1 + scale_kv) + shift_kv + del scale_kv, shift_kv + return attn(attn_input, context=encoder_hidden_states, + mask=attention_mask, transformer_options=transformer_options) * gate + return attn( + comfy.ldm.common_dit.rms_norm(x), context=context, + mask=attention_mask, transformer_options=transformer_options, + ) + + +def ltx2_forward( + self, x: Tuple[torch.Tensor, torch.Tensor], v_context=None, a_context=None, attention_mask=None, v_timestep=None, a_timestep=None, + v_pe=None, a_pe=None, v_cross_pe=None, a_cross_pe=None, v_cross_scale_shift_timestep=None, a_cross_scale_shift_timestep=None, + v_cross_gate_timestep=None, a_cross_gate_timestep=None, transformer_options=None, self_attention_mask=None, + v_prompt_timestep=None, a_prompt_timestep=None, **kwargs + ) -> Tuple[torch.Tensor, torch.Tensor]: + run_vx = transformer_options.get("run_vx", True) + run_ax = transformer_options.get("run_ax", True) + use_triton = getattr(self, 'use_triton_kernels', True) + video_scale = getattr(self, 'video_scale', 1.0) + audio_scale = getattr(self, 'audio_scale', 1.0) + audio_to_video_scale = getattr(self, 'audio_to_video_scale', 1.0) + video_to_audio_scale = getattr(self, 'video_to_audio_scale', 1.0) + + vx, ax = x + run_ax = run_ax and ax.numel() > 0 and audio_scale != 0.0 + run_a2v = run_vx and transformer_options.get("a2v_cross_attn", True) and ax.numel() > 0 and audio_to_video_scale != 0.0 + run_v2a = run_ax and transformer_options.get("v2a_cross_attn", True) and video_to_audio_scale != 0.0 + + if run_vx: + # video self-attention + vshift_msa, vscale_msa = (self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(0, 2))) + norm_vx = fused_norm_scale_shift(vx, vscale_msa, vshift_msa, use_triton=use_triton) + del vshift_msa, vscale_msa + # inline attn1 only when triton is active (needed to intercept q/k for fused RoPE) + if use_triton: + _a1 = self.attn1 + q = _a1.to_q(norm_vx) + k = _a1.to_k(norm_vx) + v = _a1.to_v(norm_vx) + + q = _a1.q_norm(q) + k = _a1.k_norm(k) + + if v_pe is not None: + q, k = fused_rope_qk(q, k, v_pe, use_triton=True) + + if self_attention_mask is None: + _sa_out = _comfy_attn.optimized_attention(q, k, v, _a1.heads, attn_precision=_a1.attn_precision, transformer_options=transformer_options) + elif _GuideAttentionMask is not None and isinstance(self_attention_mask, _GuideAttentionMask): + _sa_out = _ltx_attn_with_guide_mask(q, k, v, _a1.heads, self_attention_mask, attn_precision=_a1.attn_precision, transformer_options=transformer_options) + else: + _sa_out = _comfy_attn.optimized_attention_masked(q, k, v, _a1.heads, self_attention_mask, attn_precision=_a1.attn_precision, transformer_options=transformer_options) + del q, k, v + if _a1.to_gate_logits is not None: + _gate = _a1.to_gate_logits(norm_vx) + _b, _t, _ = _sa_out.shape + _sa_out = _sa_out.view(_b, _t, _a1.heads, _a1.dim_head) + _sa_out.mul_((2.0 * torch.sigmoid(_gate)).unsqueeze(-1)) + _sa_out = _sa_out.view(_b, _t, _a1.heads * _a1.dim_head) + del _gate + + attn1_out = _a1.to_out(_sa_out) + del _sa_out, norm_vx + else: + attn1_out = self.attn1(norm_vx, pe=v_pe, mask=self_attention_mask, transformer_options=transformer_options) + del norm_vx + # video cross-attention + vgate_msa = self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(2, 3))[0] + vx.addcmul_(attn1_out, vgate_msa, value=video_scale) + del vgate_msa, attn1_out + vx.add_(self._apply_text_cross_attention( + vx, v_context, self.attn2, self.scale_shift_table, + getattr(self, 'prompt_scale_shift_table', None), + v_timestep, v_prompt_timestep, attention_mask, transformer_options,), + alpha=video_scale + ) + # audio + if run_ax: + # audio self-attention + ashift_msa, ascale_msa = (self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], a_timestep, slice(0, 2))) + norm_ax = fused_norm_scale_shift(ax, ascale_msa, ashift_msa, use_triton=use_triton) + del ashift_msa, ascale_msa + attn1_out = self.audio_attn1(norm_ax, pe=a_pe, transformer_options=transformer_options) + del norm_ax + # audio cross-attention + agate_msa = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], a_timestep, slice(2, 3))[0] + ax.addcmul_(attn1_out, agate_msa, value=audio_scale) + del agate_msa, attn1_out + ax.add_(self._apply_text_cross_attention( + ax, a_context, self.audio_attn2, self.audio_scale_shift_table, + getattr(self, 'audio_prompt_scale_shift_table', None), + a_timestep, a_prompt_timestep, attention_mask, transformer_options,), + alpha=audio_scale + ) + + # video - audio cross attention. + if run_a2v or run_v2a: + if run_a2v and run_v2a: + # Fetch all 4 scale/shift values per table at once (avoids two get_ada_values calls each) + # and compute rms_norm once per tensor via dual kernel. + v_ca_vals = self.get_ada_values(self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep) + a_ca_vals = self.get_ada_values(self.scale_shift_table_a2v_ca_audio[:4, :], ax.shape[0], a_cross_scale_shift_timestep) + vx_a2v, vx_v2a = fused_norm_dual_scale_shift(vx, v_ca_vals[0], v_ca_vals[1], v_ca_vals[2], v_ca_vals[3], use_triton=use_triton) + ax_a2v, ax_v2a = fused_norm_dual_scale_shift(ax, a_ca_vals[0], a_ca_vals[1], a_ca_vals[2], a_ca_vals[3], use_triton=use_triton) + del v_ca_vals, a_ca_vals + elif run_a2v: + sv, bv = self.get_ada_values(self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[:2] + sa, ba = self.get_ada_values(self.scale_shift_table_a2v_ca_audio[:4, :], ax.shape[0], a_cross_scale_shift_timestep)[:2] + vx_a2v = fused_norm_scale_shift(vx, sv, bv, use_triton=use_triton) + ax_a2v = fused_norm_scale_shift(ax, sa, ba, use_triton=use_triton) + del sv, bv, sa, ba + else: # only v2a + sv, bv = self.get_ada_values(self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[2:4] + sa, ba = self.get_ada_values(self.scale_shift_table_a2v_ca_audio[:4, :], ax.shape[0], a_cross_scale_shift_timestep)[2:4] + vx_v2a = fused_norm_scale_shift(vx, sv, bv, use_triton=use_triton) + ax_v2a = fused_norm_scale_shift(ax, sa, ba, use_triton=use_triton) + del sv, bv, sa, ba + + # audio to video cross attention + if run_a2v: + a2v_out = self.audio_to_video_attn(vx_a2v, context=ax_a2v, pe=v_cross_pe, k_pe=a_cross_pe, transformer_options=transformer_options) + del vx_a2v, ax_a2v + gate_out_a2v = self.get_ada_values(self.scale_shift_table_a2v_ca_video[4:, :], vx.shape[0], v_cross_gate_timestep)[0] + vx.addcmul_(a2v_out, gate_out_a2v, value=audio_to_video_scale) + del gate_out_a2v, a2v_out + + # video to audio cross attention + if run_v2a: + v2a_out = self.video_to_audio_attn(ax_v2a, context=vx_v2a, pe=a_cross_pe, k_pe=v_cross_pe, transformer_options=transformer_options) + del ax_v2a, vx_v2a + gate_out_v2a = self.get_ada_values(self.scale_shift_table_a2v_ca_audio[4:, :], ax.shape[0], a_cross_gate_timestep)[0] + ax.addcmul_(v2a_out, gate_out_v2a, value=video_to_audio_scale) + del gate_out_v2a, v2a_out + + # video feedforward + if run_vx: + vshift_mlp, vscale_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(3, 5)) + vx_scaled = fused_norm_scale_shift(vx, vscale_mlp, vshift_mlp, use_triton=use_triton) + del vshift_mlp, vscale_mlp + + ff_out = self.ff(vx_scaled) + del vx_scaled + + vgate_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(5, 6))[0] + vx.addcmul_(ff_out, vgate_mlp, value=video_scale) + del vgate_mlp, ff_out + + # audio feedforward + if run_ax: + ashift_mlp, ascale_mlp = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], a_timestep, slice(3, 5)) + ax_scaled = fused_norm_scale_shift(ax, ascale_mlp, ashift_mlp, use_triton=use_triton) + del ashift_mlp, ascale_mlp + + ff_out = self.audio_ff(ax_scaled) + del ax_scaled + + agate_mlp = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], a_timestep, slice(5, 6))[0] + ax.addcmul_(ff_out, agate_mlp, value=audio_scale) + del agate_mlp, ff_out + + return vx, ax + +class LTX2ForwardPatch: + def __init__(self, video, audio, audio_to_video, video_to_audio, use_triton=True): + self.video_scale = video + self.audio_scale = audio + self.video_to_audio_scale = video_to_audio + self.audio_to_video_scale = audio_to_video + self.use_triton_kernels = use_triton + def __get__(self, obj, objtype=None): + def wrapped_forward(self_module, *args, **kwargs): + self_module.video_scale = self.video_scale + self_module.audio_scale = self.audio_scale + self_module.video_to_audio_scale = self.video_to_audio_scale + self_module.audio_to_video_scale = self.audio_to_video_scale + self_module.use_triton_kernels = self.use_triton_kernels + self_module._apply_text_cross_attention = types.MethodType( + _apply_text_cross_attention_patched, self_module) + return ltx2_forward(self_module, *args, **kwargs) + return types.MethodType(wrapped_forward, obj) + +class LTX2AttentionTunerPatch(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTX2AttentionTunerPatch", + display_name="LTX2 Attention Tuner Patch", + category="KJNodes/ltxv", + description="EXPERIMENTAL! Custom LTX2 forward pass with attention scaling factors per modality, also reduces peak VRAM usage.", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + io.String.Input("blocks", default="", tooltip="Comma separated list of transformer block indices to apply the patch to. Leave empty to apply to all blocks."), + io.Float.Input("video_scale", default=1.0, min=0.0, max=100, step=0.01, tooltip="Scaling factor for video attention."), + io.Float.Input("audio_scale", default=1.0, min=0.0, max=100, step=0.01, tooltip="Scaling factor for audio attention."), + io.Float.Input("audio_to_video_scale", default=1.0, min=0.0, max=100, step=0.01, tooltip="Scaling factor for video attention."), + io.Float.Input("video_to_audio_scale", default=1.0, min=0.0, max=100, step=0.01, tooltip="Scaling factor for audio attention."), + io.Boolean.Input("triton_kernels", default=True, tooltip="Use Triton fused kernels for norm+scale+shift and rope application operations, can be very slightly faster."), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, blocks, video_scale, audio_scale, audio_to_video_scale, video_to_audio_scale, triton_kernels) -> io.NodeOutput: + model_clone = model.clone() + diffusion_model = model_clone.get_model_object("diffusion_model") + + # Parse selected block indices + if blocks.strip() == "": + selected_blocks = set(range(len(diffusion_model.transformer_blocks))) + else: + selected_blocks = set(int(idx) for idx in blocks.strip().split(",")) + + logging.info(f"Applying LTX2 Attention Tuner Patch with custom scales to blocks: {sorted(selected_blocks)}, triton_kernels={triton_kernels}") + + # Apply patch to all blocks, but use 1.0 scales for non-selected blocks + for idx in range(len(diffusion_model.transformer_blocks)): + block = diffusion_model.transformer_blocks[idx] + if idx in selected_blocks: + patched_forward = LTX2ForwardPatch(video_scale, audio_scale, audio_to_video_scale, video_to_audio_scale, use_triton=triton_kernels).__get__(block, block.__class__) + else: + patched_forward = LTX2ForwardPatch(1.0, 1.0, 1.0, 1.0, use_triton=triton_kernels).__get__(block, block.__class__) + model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.forward", patched_forward) + + return io.NodeOutput(model_clone) + +class LTX2MemoryEfficientSageAttentionPatch(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTX2MemoryEfficientSageAttentionPatch", + display_name="LTX2 Mem Eff Sage Attention Patch", + category="KJNodes/ltxv", + description="EXPERIMENTAL! Activates custom sageattention to reduce peak VRAM usage, overrides the attention mode. Requires latest sageattention version.", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + io.Boolean.Input("triton_kernels", default=True, tooltip="Use Triton fused RoPE kernel on the self-attention Q/K. Requires Triton."), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, triton_kernels) -> io.NodeOutput: + if _cuda_archs is None: + raise RuntimeError("sageattention is not new enough version or could not determine CUDA architecture, cannot apply LTX2 Memory Efficient Sage Attention Patch.") + model_clone = model.clone() + diffusion_model = model_clone.get_model_object("diffusion_model") + + logging.info(f"Applying LTX2 Memory Efficient Sage Attention Patch to all transformer blocks, triton_kernels={triton_kernels}") + + for idx, block in enumerate(diffusion_model.transformer_blocks): + block.attn1.use_triton_kernels = triton_kernels + model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.attn1.forward", ltx2_sageattn_forward.__get__(block.attn1, block.attn1.__class__)) + + return io.NodeOutput(model_clone) + + +def get_cuda_version(): + try: + version = torch.version.cuda + if version is not None: + major, minor = version.split('.') + return int(major), int(minor) + else: + return 0, 0 + except Exception: + return 0, 0 + +sageplus_sm89_available = False +_cuda_archs = None +try: + from sageattention.core import per_thread_int8_triton, per_warp_int8_cuda, per_block_int8_triton, per_channel_fp8, get_cuda_arch_versions, attn_false + _cuda_archs = get_cuda_arch_versions() +except Exception: + pass +try: + from sageattention.core import _qattn_sm89 + cuda_version = get_cuda_version() + sageplus_sm89_available = hasattr(_qattn_sm89, 'qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf') and cuda_version >= (12, 8) +except ImportError: + try: + from sageattention.core import sm89_compile as _qattn_sm89 + except ImportError: + _qattn_sm89 = None +try: + from sageattention.core import _qattn_sm80 +except ImportError: + try: + from sageattention.core import sm80_compile as _qattn_sm80 + except ImportError: + _qattn_sm80 = None +try: + from sageattention.core import _qattn_sm90 +except ImportError: + try: + from sageattention.core import sm90_compile as _qattn_sm90 + except ImportError: + _qattn_sm90 = None + +from comfy.ldm.lightricks.model import apply_rotary_emb +try: + from comfy.ldm.flux.math import apply_rope as _wan_apply_rope +except ImportError: + _wan_apply_rope = None +try: + from comfy.ldm.wan.model import WanT2VCrossAttention as _WanT2VCrossAttention, WanI2VCrossAttention as _WanI2VCrossAttention +except ImportError: + _WanT2VCrossAttention = _WanI2VCrossAttention = None + + +def _sageattn_int8_fp8_nhd(qkv, dtype): + # qkv: [q, k, v], each [batch, seq_len, num_heads, head_dim] in NHD layout. Returns o in the same layout. + # The list is consumed so the only references to the float q/k/v are the locals below, letting `del` + # actually free them before the kernel runs — attention is the VRAM peak in these models. + q, k, v = qkv + qkv.clear() + head_dim_og = q.shape[-1] + + tensor_layout="NHD" + _tensor_layout = 0 # NHD + _is_caual = 0 + _qk_quant_gran = 3 + _return_lse = 0 + sm_scale = head_dim_og**-0.5 + quant_v_scale_max = 448.0 + + if _cuda_archs[0] in {"sm80", "sm86"}: + # mean-sub in-place: passing km= makes the triton quant materialize k - km as a full float copy + k.sub_(k.mean(dim=1, keepdim=True)) + q_int8, q_scale, k_int8, k_scale = per_thread_int8_triton(q, k, tensor_layout=tensor_layout, BLKQ=128, WARPQ=32, BLKK=64, WARPK=64) + del q, k + o = torch.empty(q_int8.size(), dtype=dtype, device=q_int8.device) + v_fp16 = v.to(torch.float16) + del v + _qattn_sm80.qk_int8_sv_f16_accum_f32_attn(q_int8, k_int8, v_fp16, o, q_scale, k_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + elif _cuda_archs[0] == "sm75": + k.sub_(k.mean(dim=1, keepdim=True)) + q_int8, q_scale, k_int8, k_scale = per_block_int8_triton(q, k, sm_scale=sm_scale, tensor_layout=tensor_layout) + del q, k + o, _ = attn_false(q_int8, k_int8, v, q_scale, k_scale, tensor_layout=tensor_layout, output_dtype=dtype, attn_mask=None, return_lse=False) + del v + elif _cuda_archs[0] == "sm89": + if not sageplus_sm89_available: + pv_accum_dtype = "fp32+fp32" + else: + pv_accum_dtype = "fp32+fp16" + quant_v_scale_max = 2.25 + k.sub_(k.mean(dim=1, keepdim=True)) + q_int8, q_scale, k_int8, k_scale = per_thread_int8_triton(q, k, tensor_layout=tensor_layout, BLKQ=128, WARPQ=32, BLKK=64, WARPK=64) + del q, k + v_fp8, v_scale, _ = per_channel_fp8(v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=False) + del v + o = torch.empty(q_int8.size(), dtype=dtype, device=q_int8.device) + if pv_accum_dtype == "fp32+fp16": + _qattn_sm89.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(q_int8, k_int8, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + elif pv_accum_dtype == "fp32+fp32": + _qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(q_int8, k_int8, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + del v_fp8, v_scale + elif _cuda_archs[0] == "sm90": + k.sub_(k.mean(dim=1, keepdim=True)) + q_int8, q_scale, k_int8, k_scale = per_thread_int8_triton(q, k, tensor_layout=tensor_layout, BLKQ=64, WARPQ=16, BLKK=128, WARPK=128) + del q, k + v_fp8, v_scale, _ = per_channel_fp8(v, tensor_layout=tensor_layout, smooth_v=False) + del v + o = torch.empty(q_int8.size(), dtype=dtype, device=q_int8.device) + _qattn_sm90.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(q_int8, k_int8, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + del v_fp8, v_scale + elif _cuda_archs[0] == "sm120": + if not sageplus_sm89_available: + pv_accum_dtype = "fp32" + else: + pv_accum_dtype = "fp32+fp16" + quant_v_scale_max = 2.25 + _qk_quant_gran = 2 # per warp + # km kept here: the CUDA quant fuses the mean-sub with no temp copy + q_int8, q_scale, k_int8, k_scale = per_warp_int8_cuda(q, k, km=k.mean(dim=1, keepdim=True), tensor_layout=tensor_layout, BLKQ=128, WARPQ=32, BLKK=64) + del q, k + v_fp8, v_scale, _ = per_channel_fp8(v, tensor_layout=tensor_layout, scale_max=quant_v_scale_max, smooth_v=False) + del v + o = torch.empty(q_int8.size(), dtype=dtype, device=q_int8.device) + if pv_accum_dtype == "fp32": + _qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(q_int8, k_int8, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + elif pv_accum_dtype == "fp32+fp16": + _qattn_sm89.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(q_int8, k_int8, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + del v_fp8, v_scale + + del q_int8, q_scale, k_int8, k_scale + return o + + +def ltx2_sageattn_forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}): + dtype = x.dtype + context = x if context is None else context + + # query + q = self.to_q(x) + q = self.q_norm(q) + # key + k = self.to_k(context) + k = self.k_norm(k) + # apply RoPE — fuse q+k into one kernel when both share the same pe + if pe is not None: + use_triton = getattr(self, 'use_triton_kernels', False) + if k_pe is None: + q, k = fused_rope_qk(q, k, pe, use_triton=use_triton) + else: + q = apply_rotary_emb(q, pe) + k = apply_rotary_emb(k, k_pe) + # value + v = self.to_v(context) + + # Sage kernels don't support masking — fall back to the upstream masked paths. + if mask is not None: + if _GuideAttentionMask is not None and isinstance(mask, _GuideAttentionMask): + o = _ltx_attn_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options) + else: + o = _comfy_attn.optimized_attention_masked(q, k, v, self.heads, 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, _ = o.shape + o = o.view(_b, _t, self.heads, self.dim_head) + o.mul_((2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)) + o = o.view(_b, _t, self.heads * self.dim_head) + del gate_logits + return self.to_out(o) + + # Reshape from [batch, seq_len, total_dim] to [batch, seq_len, num_heads, head_dim] + batch_size, seq_len, _ = q.shape + head_dim_og = self.dim_head + + q = q.view(batch_size, seq_len, self.heads, head_dim_og) + k = k.view(batch_size, k.shape[1], self.heads, head_dim_og) + v = v.view(batch_size, v.shape[1], self.heads, head_dim_og) + + qkv = [q, k, v] + del q, k, v + o = _sageattn_int8_fp8_nhd(qkv, dtype) + + # o is [B, T, H, D] from sage kernel (NHD layout) + if self.to_gate_logits is not None: + gate_logits = self.to_gate_logits(x) # (B, T, H) + o.mul_((2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)) + del gate_logits + + return self.to_out(o.view(batch_size, seq_len, -1)) + + +def wan_sageattn_forward(self, x, freqs, transformer_options={}): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + dtype = x.dtype + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + q, k = _wan_apply_rope(q, k, freqs) + v = self.v(x).view(b, s, n, d) + + qkv = [q, k, v] + del q, k, v + o = _sageattn_int8_fp8_nhd(qkv, dtype) + + return self.o(o.view(b, s, n * d)) + + +def wan_t2v_cross_sageattn_forward(self, x, context, transformer_options={}, **kwargs): + dtype = x.dtype + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + qkv = [q, k, v] + del q, k, v + o = _sageattn_int8_fp8_nhd(qkv, dtype) + + return self.o(o.view(b, s, n * d)) + + +def wan_i2v_cross_sageattn_forward(self, x, context, context_img_len, transformer_options={}): + dtype = x.dtype + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + context_img = context[:, :context_img_len] + context = context[:, context_img_len:] + + q = self.norm_q(self.q(x)).view(b, s, n, d) + k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) + v_img = self.v_img(context_img).view(b, -1, n, d) + # local q reference survives the helper's consume — still needed for the text attention below + qkv_img = [q, k_img, v_img] + del k_img, v_img + img_o = _sageattn_int8_fp8_nhd(qkv_img, dtype) + + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + qkv = [q, k, v] + del q, k, v + o = _sageattn_int8_fp8_nhd(qkv, dtype) + + o.add_(img_o) + del img_o + return self.o(o.view(b, s, n * d)) + + +class WanVideoMemoryEfficientSageAttentionPatch(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WanVideoMemoryEfficientSageAttentionPatch", + display_name="WanVideo Mem Eff Sage Attention Patch", + category="KJNodes/wan", + description="EXPERIMENTAL! Activates custom sageattention on the WanVideo self-attention to reduce peak VRAM usage, overrides the attention mode. Requires latest sageattention version.", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model) -> io.NodeOutput: + if _cuda_archs is None: + raise RuntimeError("sageattention is not new enough version or could not determine CUDA architecture, cannot apply WanVideo Memory Efficient Sage Attention Patch.") + if _wan_apply_rope is None: + raise RuntimeError("Could not import apply_rope from comfy.ldm.flux.math, cannot apply WanVideo Memory Efficient Sage Attention Patch.") + model_clone = model.clone() + diffusion_model = model_clone.get_model_object("diffusion_model") + + logging.info("Applying WanVideo Memory Efficient Sage Attention Patch to all transformer blocks") + + for idx, block in enumerate(diffusion_model.blocks): + model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.self_attn.forward", wan_sageattn_forward.__get__(block.self_attn, block.self_attn.__class__)) + cross_attn = getattr(block, "cross_attn", None) + # exact type match on purpose: subclasses like WanT2VCrossAttentionGather have different semantics + if cross_attn is not None and type(cross_attn) is _WanI2VCrossAttention: + model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.cross_attn.forward", wan_i2v_cross_sageattn_forward.__get__(cross_attn, cross_attn.__class__)) + elif cross_attn is not None and type(cross_attn) is _WanT2VCrossAttention: + model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.cross_attn.forward", wan_t2v_cross_sageattn_forward.__get__(cross_attn, cross_attn.__class__)) + + return io.NodeOutput(model_clone) + + +import folder_paths + +class LTX2LoraLoaderAdvanced(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTX2LoraLoaderAdvanced", + display_name="LTX2 LoRA Loader Advanced", + category="KJNodes/ltxv", + description="Advanced LoRA loader with per-block strength control for LTX2 models", + is_experimental=True, + inputs=[ + io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras"), tooltip="The name of the LoRA."), + io.Model.Input("model", tooltip="The diffusion model the LoRA will be applied to."), + io.Float.Input("strength_model", default=1.0, min=-100.0, max=100.0, step=0.01, tooltip="How strongly to modify the diffusion model. This value can be negative."), + io.String.Input("opt_lora_path", optional=True, force_input=True,tooltip="Absolute path of the LoRA."), + io.Custom("SELECTEDDITBLOCKS").Input("blocks", optional=True, tooltip="Selected DiT blocks configuration"), + io.Float.Input("video", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Strength for video attention layers."), + io.Float.Input("video_to_audio", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Strength for video to audio cross-attention layers."), + io.Float.Input("audio", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Strength for audio attention layers."), + io.Float.Input("audio_to_video", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Strength for audio to video cross-attention layers."), + io.Float.Input("other", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Strength for layers not caught by other layer filters."), + ], + outputs=[ + io.Model.Output(display_name="model", tooltip="The modified diffusion model."), + io.String.Output(display_name="rank", tooltip="Possible rank of the LoRA."), + io.String.Output(display_name="loaded_keys_info", tooltip="List of loaded keys and their alpha values."), + ], + ) + + @classmethod + def execute(cls, model, lora_name, strength_model, video, video_to_audio, audio, audio_to_video, other, opt_lora_path=None, blocks=None) -> io.NodeOutput: + from comfy.utils import load_torch_file + import comfy.lora + + if opt_lora_path: + lora_path = opt_lora_path + else: + lora_path = folder_paths.get_full_path("loras", lora_name) + + lora = load_torch_file(lora_path, safe_load=True) + + # Find the first key that ends with "weight" + rank = "unknown" + weight_key = next((key for key in lora.keys() if key.endswith('weight')), None) + # Print the shape of the value corresponding to the key + if weight_key: + print(f"Shape of the first 'weight' key ({weight_key}): {lora[weight_key].shape}") + rank = str(lora[weight_key].shape[0]) + else: + print("No key ending with 'weight' found.") + rank = "Couldn't find rank" + + key_map = {} + if model is not None: + key_map = comfy.lora.model_lora_keys_unet(model.model, key_map) + + loaded = comfy.lora.load_lora(lora, key_map) + + keys_to_delete = [] + + # First apply blocks filtering if provided + if blocks is not None: + for block in blocks: + for key in list(loaded.keys()): + match = False + if isinstance(key, str) and block in key: + match = True + elif isinstance(key, tuple): + for k in key: + if block in k: + match = True + break + + if match: + ratio = blocks[block] + if ratio == 0: + keys_to_delete.append(key) + else: + # Only modify LoRA adapters, skip diff tuples + value = loaded[key] + if hasattr(value, 'weights'): + weights_list = list(value.weights) + weights_list[2] = ratio + loaded[key].weights = tuple(weights_list) + + # Then apply layer-based attention strength filtering (takes priority) + for key in list(loaded.keys()): + if key in keys_to_delete: + continue + + key_str = key if isinstance(key, str) else (key[0] if isinstance(key, tuple) else str(key)) + + # Determine the strength multiplier based on layer name + # Check more specific patterns first + strength_multiplier = None + + # Video to audio cross-attention (check first - most specific) + if "video_to_audio_attn" in key_str: + strength_multiplier = video_to_audio + # Audio to video cross-attention + elif "audio_to_video_attn" in key_str: + strength_multiplier = audio_to_video + # Audio layers + elif "audio_attn" in key_str or "audio_ff.net" in key_str: + strength_multiplier = audio + # Video layers (check last - most general) + elif "attn" in key_str or "ff.net" in key_str: + strength_multiplier = video + # Everything else not caught by above filters + else: + strength_multiplier = other + + # Apply strength or mark for deletion + if strength_multiplier is not None: + if strength_multiplier == 0: + keys_to_delete.append(key) + elif strength_multiplier != 1.0: + value = loaded[key] + if hasattr(value, 'weights'): + weights_list = list(value.weights) + # Handle case where alpha (weights[2]) might be None + current_alpha = weights_list[2] if weights_list[2] is not None else 1.0 + weights_list[2] = current_alpha * strength_multiplier + loaded[key].weights = tuple(weights_list) + + for key in keys_to_delete: + if key in loaded: + del loaded[key] + + # Build list of loaded keys and their alphas + loaded_keys_list = [] + for key, value in loaded.items(): + if hasattr(value, 'weights'): + key_str = key if isinstance(key, str) else str(key) + alpha = value.weights[2] if value.weights[2] is not None else "None" + loaded_keys_list.append(f"{key_str}: alpha={alpha}") + else: + key_str = key if isinstance(key, str) else str(key) + loaded_keys_list.append(f"{key_str}: type={type(value).__name__}") + + if model is not None: + new_modelpatcher = model.clone() + k = new_modelpatcher.add_patches(loaded, strength_model) + + # Add not loaded keys to the info + k = set(k) + not_loaded = [] + for x in loaded: + if x not in k: + key_str = x if isinstance(x, str) else str(x) + not_loaded.append(f"NOT LOADED: {key_str}") + + if not_loaded: + loaded_keys_list.extend(not_loaded) + + loaded_keys_info = "\n".join(loaded_keys_list) + + return io.NodeOutput(new_modelpatcher, rank, loaded_keys_info) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/mask_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/mask_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..ca4dc69ce972751f85d2b80e1c3e5e2de1d82099 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/mask_nodes.py @@ -0,0 +1,1689 @@ +import torch +import torch.nn.functional as F +from torchvision.transforms import functional as TF +from PIL import Image, ImageDraw, ImageFilter, ImageFont +import scipy.ndimage +import numpy as np +from contextlib import nullcontext +import os +from tqdm import tqdm +import logging + +from comfy import model_management +from comfy.utils import ProgressBar +from comfy.utils import common_upscale +from nodes import MAX_RESOLUTION + +import folder_paths + +from ..utility.utility import tensor2pil, pil2tensor + +script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +main_device = model_management.get_torch_device() +offload_device = model_management.unet_offload_device() + +class BatchCLIPSeg: + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + + return {"required": + { + "images": ("IMAGE",), + "text": ("STRING", {"multiline": False}), + "threshold": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 10.0, "step": 0.001}), + "binary_mask": ("BOOLEAN", {"default": True}), + "combine_mask": ("BOOLEAN", {"default": False}), + "use_cuda": ("BOOLEAN", {"default": True}), + }, + "optional": + { + "blur_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), + "opt_model": ("CLIPSEGMODEL", ), + "prev_mask": ("MASK", {"default": None}), + "image_bg_level": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "invert": ("BOOLEAN", {"default": False}), + } + } + + CATEGORY = "KJNodes/masking" + RETURN_TYPES = ("MASK", "IMAGE", ) + RETURN_NAMES = ("Mask", "Image", ) + FUNCTION = "segment_image" + DESCRIPTION = """ +Segments an image or batch of images using CLIPSeg. +""" + + def segment_image(self, images, text, threshold, binary_mask, combine_mask, use_cuda, blur_sigma=0.0, opt_model=None, prev_mask=None, invert= False, image_bg_level=0.5): + from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation + import torchvision.transforms as transforms + offload_device = model_management.unet_offload_device() + device = model_management.get_torch_device() + if not use_cuda: + device = torch.device("cpu") + dtype = model_management.unet_dtype() + + if opt_model is None: + checkpoint_path = os.path.join(folder_paths.models_dir,'clip_seg', 'clipseg-rd64-refined-fp16') + if not hasattr(self, "model"): + try: + if not os.path.exists(checkpoint_path): + from huggingface_hub import snapshot_download + snapshot_download(repo_id="Kijai/clipseg-rd64-refined-fp16", local_dir=checkpoint_path, local_dir_use_symlinks=False) + self.model = CLIPSegForImageSegmentation.from_pretrained(checkpoint_path) + except Exception: + checkpoint_path = "CIDAS/clipseg-rd64-refined" + self.model = CLIPSegForImageSegmentation.from_pretrained(checkpoint_path) + processor = CLIPSegProcessor.from_pretrained(checkpoint_path) + + else: + self.model = opt_model['model'] + processor = opt_model['processor'] + + self.model.to(dtype).to(device) + + B, H, W, C = images.shape + images = images.to(device) + + autocast_condition = (dtype != torch.float32) and not model_management.is_device_mps(device) + with torch.autocast(model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): + + PIL_images = [Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) for image in images ] + prompt = [text] * len(images) + input_prc = processor(text=prompt, images=PIL_images, return_tensors="pt") + + for key in input_prc: + input_prc[key] = input_prc[key].to(device) + outputs = self.model(**input_prc) + + mask_tensor = torch.sigmoid(outputs.logits) + mask_tensor = (mask_tensor - mask_tensor.min()) / (mask_tensor.max() - mask_tensor.min()) + mask_tensor = torch.where(mask_tensor > (threshold), mask_tensor, torch.tensor(0, dtype=torch.float)) + + if len(mask_tensor.shape) == 2: + mask_tensor = mask_tensor.unsqueeze(0) + mask_tensor = F.interpolate(mask_tensor.unsqueeze(1), size=(H, W), mode='nearest') + mask_tensor = mask_tensor.squeeze(1) + + self.model.to(offload_device) + + if binary_mask: + mask_tensor = (mask_tensor > 0).float() + if blur_sigma > 0: + kernel_size = int(6 * int(blur_sigma) + 1) + blur = transforms.GaussianBlur(kernel_size=(kernel_size, kernel_size), sigma=(blur_sigma, blur_sigma)) + mask_tensor = blur(mask_tensor) + + if combine_mask: + mask_tensor = torch.max(mask_tensor, dim=0)[0] + mask_tensor = mask_tensor.unsqueeze(0).repeat(len(images),1,1) + + del outputs + model_management.soft_empty_cache() + + if prev_mask is not None: + if prev_mask.shape != mask_tensor.shape: + prev_mask = F.interpolate(prev_mask.unsqueeze(1), size=(H, W), mode='nearest') + mask_tensor = mask_tensor + prev_mask.to(device) + torch.clamp(mask_tensor, min=0.0, max=1.0) + + if invert: + mask_tensor = 1 - mask_tensor + + image_tensor = images * mask_tensor.unsqueeze(-1) + (1 - mask_tensor.unsqueeze(-1)) * image_bg_level + image_tensor = torch.clamp(image_tensor, min=0.0, max=1.0).cpu().float() + + mask_tensor = mask_tensor.cpu().float() + + return mask_tensor, image_tensor, + +class DownloadAndLoadCLIPSeg: + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + + return {"required": + { + "model": ( + [ 'Kijai/clipseg-rd64-refined-fp16', + 'CIDAS/clipseg-rd64-refined', + ], + ), + }, + } + + CATEGORY = "KJNodes/masking" + RETURN_TYPES = ("CLIPSEGMODEL",) + RETURN_NAMES = ("clipseg_model",) + FUNCTION = "segment_image" + DESCRIPTION = """ +Downloads and loads CLIPSeg model with huggingface_hub, +to ComfyUI/models/clip_seg +""" + + def segment_image(self, model): + from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation + checkpoint_path = os.path.join(folder_paths.models_dir,'clip_seg', os.path.basename(model)) + if not hasattr(self, "model"): + if not os.path.exists(checkpoint_path): + from huggingface_hub import snapshot_download + snapshot_download(repo_id=model, local_dir=checkpoint_path, local_dir_use_symlinks=False) + self.model = CLIPSegForImageSegmentation.from_pretrained(checkpoint_path) + + processor = CLIPSegProcessor.from_pretrained(checkpoint_path) + + clipseg_model = {} + clipseg_model['model'] = self.model + clipseg_model['processor'] = processor + + return clipseg_model, + +class CreateTextMask: + + RETURN_TYPES = ("IMAGE", "MASK",) + FUNCTION = "createtextmask" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +Creates a text image and mask. +Looks for fonts from this folder: +ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts + +If start_rotation and/or end_rotation are different values, +creates animation between them. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), + "text_x": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "text_y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "font_size": ("INT", {"default": 32,"min": 8, "max": 4096, "step": 1}), + "font_color": ("STRING", {"default": "white"}), + "text": ("STRING", {"default": "HELLO!", "multiline": True}), + "font": (folder_paths.get_filename_list("kjnodes_fonts"), ), + "width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "start_rotation": ("INT", {"default": 0,"min": 0, "max": 359, "step": 1}), + "end_rotation": ("INT", {"default": 0,"min": -359, "max": 359, "step": 1}), + }, + } + + def createtextmask(self, frames, width, height, invert, text_x, text_y, text, font_size, font_color, font, start_rotation, end_rotation): + # Define the number of images in the batch + batch_size = frames + out = [] + masks = [] + rotation = start_rotation + if start_rotation != end_rotation: + rotation_increment = (end_rotation - start_rotation) / (batch_size - 1) + + font_path = folder_paths.get_full_path("kjnodes_fonts", font) + # Generate the text + for i in range(batch_size): + image = Image.new("RGB", (width, height), "black") + draw = ImageDraw.Draw(image) + font = ImageFont.truetype(font_path, font_size) + + # Split the text into lines and wrap words to fit width + text_lines = text.split('\n') + lines = [] + for text_line in text_lines: + if text_line.strip() == "": + # Preserve empty lines for multiple newlines + lines.append("") + continue + words = text_line.split() + current_line = [] + for word in words: + if current_line: + test_line = " ".join(current_line + [word]) + else: + test_line = word + try: + test_line_width = font.getbbox(test_line)[2] + except Exception: + test_line_width = font.getsize(test_line)[0] + if test_line_width <= width - 2 * text_x: + current_line.append(word) + else: + lines.append(" ".join(current_line)) + current_line = [word] + if current_line: + lines.append(" ".join(current_line)) + + # Draw each line of text separately + y_offset = text_y + for line in lines: + text_width = font.getlength(line) + text_height = font_size + text_center_x = text_x + text_width / 2 + text_center_y = y_offset + text_height / 2 + try: + draw.text((text_x, y_offset), line, font=font, fill=font_color, features=['-liga']) + except Exception: + draw.text((text_x, y_offset), line, font=font, fill=font_color) + y_offset += text_height # Move to the next line + + if start_rotation != end_rotation: + image = image.rotate(rotation, center=(text_center_x, text_center_y)) + rotation += rotation_increment + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0), 1.0 - torch.cat(masks, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + +class ColorToMask: + + RETURN_TYPES = ("MASK",) + FUNCTION = "clip" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Converts chosen RGB value to a mask. +With batch inputs, the **per_batch** +controls the number of images processed at once. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "invert": ("BOOLEAN", {"default": False}), + "red": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "green": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "blue": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "threshold": ("INT", {"default": 10,"min": 0, "max": 255, "step": 1}), + "per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}), + }, + } + + def clip(self, images, red, green, blue, threshold, invert, per_batch): + + color = torch.tensor([red, green, blue], dtype=torch.uint8) + black = torch.tensor([0, 0, 0], dtype=torch.uint8) + white = torch.tensor([255, 255, 255], dtype=torch.uint8) + + if invert: + black, white = white, black + + steps = images.shape[0] + pbar = ProgressBar(steps) + tensors_out = [] + + for start_idx in range(0, images.shape[0], per_batch): + + # Calculate color distances + color_distances = torch.norm(images[start_idx:start_idx+per_batch] * 255 - color, dim=-1) + + # Create a mask based on the threshold + mask = color_distances <= threshold + + # Apply the mask to create new images + mask_out = torch.where(mask.unsqueeze(-1), white, black).float() + mask_out = mask_out.mean(dim=-1) + + tensors_out.append(mask_out.cpu()) + batch_count = mask_out.shape[0] + pbar.update(batch_count) + + tensors_out = torch.cat(tensors_out, dim=0) + tensors_out = torch.clamp(tensors_out, min=0.0, max=1.0) + return tensors_out, + +class CreateFluidMask: + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "createfluidmask" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "inflow_count": ("INT", {"default": 3,"min": 0, "max": 255, "step": 1}), + "inflow_velocity": ("INT", {"default": 1,"min": 0, "max": 255, "step": 1}), + "inflow_radius": ("INT", {"default": 8,"min": 0, "max": 255, "step": 1}), + "inflow_padding": ("INT", {"default": 50,"min": 0, "max": 255, "step": 1}), + "inflow_duration": ("INT", {"default": 60,"min": 0, "max": 255, "step": 1}), + }, + } + #using code from https://github.com/GregTJ/stable-fluids + def createfluidmask(self, frames, width, height, invert, inflow_count, inflow_velocity, inflow_radius, inflow_padding, inflow_duration): + from ..utility.fluid import Fluid + try: + from scipy.special import erf + except ImportError: + from scipy.spatial import erf + out = [] + masks = [] + RESOLUTION = width, height + DURATION = frames + + INFLOW_PADDING = inflow_padding + INFLOW_DURATION = inflow_duration + INFLOW_RADIUS = inflow_radius + INFLOW_VELOCITY = inflow_velocity + INFLOW_COUNT = inflow_count + + logging.info('Generating fluid solver, this may take some time.') + fluid = Fluid(RESOLUTION, 'dye') + + center = np.floor_divide(RESOLUTION, 2) + r = np.min(center) - INFLOW_PADDING + + points = np.linspace(-np.pi, np.pi, INFLOW_COUNT, endpoint=False) + points = tuple(np.array((np.cos(p), np.sin(p))) for p in points) + normals = tuple(-p for p in points) + points = tuple(r * p + center for p in points) + + inflow_velocity = np.zeros_like(fluid.velocity) + inflow_dye = np.zeros(fluid.shape) + for p, n in zip(points, normals): + mask = np.linalg.norm(fluid.indices - p[:, None, None], axis=0) <= INFLOW_RADIUS + inflow_velocity[:, mask] += n[:, None] * INFLOW_VELOCITY + inflow_dye[mask] = 1 + + + for f in range(DURATION): + logging.info(f'Computing frame {f + 1} of {DURATION}.') + if f <= INFLOW_DURATION: + fluid.velocity += inflow_velocity + fluid.dye += inflow_dye + + curl = fluid.step()[1] + # Using the error function to make the contrast a bit higher. + # Any other sigmoid function e.g. smoothstep would work. + curl = (erf(curl * 2) + 1) / 4 + + color = np.dstack((curl, np.ones(fluid.shape), fluid.dye)) + color = (np.clip(color, 0, 1) * 255).astype('uint8') + image = np.array(color).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0),1.0 - torch.cat(masks, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + +class CreateAudioMask: + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "createaudiomask" + CATEGORY = "KJNodes/deprecated" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 16,"min": 1, "max": 255, "step": 1}), + "scale": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 2.0, "step": 0.01}), + "audio_path": ("STRING", {"default": "audio.wav"}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + }, + } + + def createaudiomask(self, frames, width, height, invert, audio_path, scale): + try: + import librosa + except ImportError as e: + raise ImportError("Can not import librosa. Install it with 'pip install librosa'") from e + batch_size = frames + out = [] + masks = [] + if audio_path == "audio.wav": #I don't know why relative path won't work otherwise... + audio_path = os.path.join(script_directory, audio_path) + audio, sr = librosa.load(audio_path) + spectrogram = np.abs(librosa.stft(audio)) + + for i in range(batch_size): + image = Image.new("RGB", (width, height), "black") + draw = ImageDraw.Draw(image) + frame = spectrogram[:, i] + circle_radius = int(height * np.mean(frame)) + circle_radius *= scale + circle_center = (width // 2, height // 2) # Calculate the center of the image + + draw.ellipse([(circle_center[0] - circle_radius, circle_center[1] - circle_radius), + (circle_center[0] + circle_radius, circle_center[1] + circle_radius)], + fill='white') + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + +class CreateGradientMask: + + RETURN_TYPES = ("MASK",) + FUNCTION = "createmask" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + }, + } + def createmask(self, frames, width, height, invert): + # Define the number of images in the batch + batch_size = frames + out = [] + # Create an empty array to store the image batch + image_batch = np.zeros((batch_size, height, width), dtype=np.float32) + # Generate the black to white gradient for each image + for i in range(batch_size): + gradient = np.linspace(1.0, 0.0, width, dtype=np.float32) + time = i / frames # Calculate the time variable + offset_gradient = gradient - time # Offset the gradient values based on time + image_batch[i] = offset_gradient.reshape(1, -1) + output = torch.from_numpy(image_batch) + mask = output + out.append(mask) + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),) + +class CreateFadeMask: + + RETURN_TYPES = ("MASK",) + FUNCTION = "createfademask" + CATEGORY = "KJNodes/deprecated" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 2,"min": 2, "max": 10000, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), + "start_level": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 1.0, "step": 0.01}), + "midpoint_level": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}), + "end_level": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 1.0, "step": 0.01}), + "midpoint_frame": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + }, + } + + def createfademask(self, frames, width, height, invert, interpolation, start_level, midpoint_level, end_level, midpoint_frame): + def ease_in(t): + return t * t + + def ease_out(t): + return 1 - (1 - t) * (1 - t) + + def ease_in_out(t): + return 3 * t * t - 2 * t * t * t + + batch_size = frames + out = [] + image_batch = np.zeros((batch_size, height, width), dtype=np.float32) + + if midpoint_frame == 0: + midpoint_frame = batch_size // 2 + + for i in range(batch_size): + if i <= midpoint_frame: + t = i / midpoint_frame + if interpolation == "ease_in": + t = ease_in(t) + elif interpolation == "ease_out": + t = ease_out(t) + elif interpolation == "ease_in_out": + t = ease_in_out(t) + color = start_level - t * (start_level - midpoint_level) + else: + t = (i - midpoint_frame) / (batch_size - midpoint_frame) + if interpolation == "ease_in": + t = ease_in(t) + elif interpolation == "ease_out": + t = ease_out(t) + elif interpolation == "ease_in_out": + t = ease_in_out(t) + color = midpoint_level - t * (midpoint_level - end_level) + + color = np.clip(color, 0, 255) + image = np.full((height, width), color, dtype=np.float32) + image_batch[i] = image + + output = torch.from_numpy(image_batch) + mask = output + out.append(mask) + + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),) + +class CreateFadeMaskAdvanced: + + RETURN_TYPES = ("MASK",) + FUNCTION = "createfademask" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Create a batch of masks interpolated between given frames and values. +Uses same syntax as Fizz' BatchValueSchedule. +First value is the frame index (not that this starts from 0, not 1) +and the second value inside the brackets is the float value of the mask in range 0.0 - 1.0 + +For example the default values: +0:(0.0) +7:(1.0) +15:(0.0) + +Would create a mask batch fo 16 frames, starting from black, +interpolating with the chosen curve to fully white at the 8th frame, +and interpolating from that to fully black at the 16th frame. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "points_string": ("STRING", {"default": "0:(0.0),\n7:(1.0),\n15:(0.0)\n", "multiline": True}), + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 16,"min": 2, "max": 10000, "step": 1}), + "width": ("INT", {"default": 512,"min": 1, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 1, "max": 4096, "step": 1}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out", "none", "default_to_black"],), + }, + } + + def createfademask(self, frames, width, height, invert, points_string, interpolation): + def ease_in(t): + return t * t + + def ease_out(t): + return 1 - (1 - t) * (1 - t) + + def ease_in_out(t): + return 3 * t * t - 2 * t * t * t + + # Parse the input string into a list of tuples + points = [] + points_string = points_string.rstrip(',\n') + for point_str in points_string.split(','): + frame_str, color_str = point_str.split(':') + frame = int(frame_str.strip()) + color = float(color_str.strip()[1:-1]) # Remove parentheses around color + points.append((frame, color)) + + # Check if the last frame is already in the points + if (interpolation != "default_to_black") and (len(points) == 0 or points[-1][0] != frames - 1): + # If not, add it with the color of the last specified frame + points.append((frames - 1, points[-1][1] if points else 0)) + + # Sort the points by frame number + points.sort(key=lambda x: x[0]) + + batch_size = frames + out = [] + image_batch = np.zeros((batch_size, height, width), dtype=np.float32) + + # Index of the next point to interpolate towards + next_point = 1 + + for i in range(batch_size): + while next_point < len(points) and i > points[next_point][0]: + next_point += 1 + + # Interpolate between the previous point and the next point + prev_point = next_point - 1 + + if interpolation == "none": + exact_match = False + for p in points: + if p[0] == i: # Exact frame match + color = p[1] + exact_match = True + break + if not exact_match: + color = points[prev_point][1] + + elif interpolation == "default_to_black": + exact_match = False + for p in points: + if p[0] == i: # Exact frame match + color = p[1] + exact_match = True + break + if not exact_match: + color = 0 + else: + t = (i - points[prev_point][0]) / (points[next_point][0] - points[prev_point][0]) + if interpolation == "ease_in": + t = ease_in(t) + elif interpolation == "ease_out": + t = ease_out(t) + elif interpolation == "ease_in_out": + t = ease_in_out(t) + elif interpolation == "linear": + pass # No need to modify `t` for linear interpolation + + color = points[prev_point][1] - t * (points[prev_point][1] - points[next_point][1]) + + color = np.clip(color, 0, 255) + image = np.full((height, width), color, dtype=np.float32) + image_batch[i] = image + + output = torch.from_numpy(image_batch) + mask = output + out.append(mask) + + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),) + +class CreateMagicMask: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createmagicmask" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}), + "depth": ("INT", {"default": 12,"min": 1, "max": 500, "step": 1}), + "distortion": ("FLOAT", {"default": 1.5,"min": 0.0, "max": 100.0, "step": 0.01}), + "seed": ("INT", {"default": 123,"min": 0, "max": 99999999, "step": 1}), + "transitions": ("INT", {"default": 1,"min": 1, "max": 20, "step": 1}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + }, + } + + def createmagicmask(self, frames, transitions, depth, distortion, seed, frame_width, frame_height): + from ..utility.magictex import coordinate_grid, random_transform, magic + import matplotlib.pyplot as plt + rng = np.random.default_rng(seed) + out = [] + coords = coordinate_grid((frame_width, frame_height)) + + # Calculate the number of frames for each transition + frames_per_transition = frames // transitions + + # Generate a base set of parameters + base_params = { + "coords": random_transform(coords, rng), + "depth": depth, + "distortion": distortion, + } + for t in range(transitions): + # Generate a second set of parameters that is at most max_diff away from the base parameters + params1 = base_params.copy() + params2 = base_params.copy() + + params1['coords'] = random_transform(coords, rng) + params2['coords'] = random_transform(coords, rng) + + for i in range(frames_per_transition): + # Compute the interpolation factor + alpha = i / frames_per_transition + + # Interpolate between the two sets of parameters + params = params1.copy() + params['coords'] = (1 - alpha) * params1['coords'] + alpha * params2['coords'] + + tex = magic(**params) + + dpi = frame_width / 10 + fig = plt.figure(figsize=(10, 10), dpi=dpi) + + ax = fig.add_subplot(111) + plt.subplots_adjust(left=0, right=1, bottom=0, top=1) + + ax.get_yaxis().set_ticks([]) + ax.get_xaxis().set_ticks([]) + ax.imshow(tex, aspect='auto') + + fig.canvas.draw() + img = np.array(fig.canvas.renderer._renderer) + + plt.close(fig) + + pil_img = Image.fromarray(img).convert("L") + mask = torch.tensor(np.array(pil_img)) / 255.0 + + out.append(mask) + + return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) + +class CreateShapeMask: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createshapemask" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Creates a mask or batch of masks with the specified shape. +Locations are center locations. +Grow value is the amount to grow the shape on each frame, creating animated masks. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "shape": ( + [ 'circle', + 'square', + 'triangle', + ], + { + "default": 'circle' + }), + "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), + "location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), + "location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), + "grow": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "shape_width": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + "shape_height": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + }, + } + + def createshapemask(self, frames, frame_width, frame_height, location_x, location_y, shape_width, shape_height, grow, shape): + # Define the number of images in the batch + batch_size = frames + out = [] + color = "white" + for i in range(batch_size): + image = Image.new("RGB", (frame_width, frame_height), "black") + draw = ImageDraw.Draw(image) + + # Calculate the size for this frame and ensure it's not less than 0 + current_width = max(0, shape_width + i*grow) + current_height = max(0, shape_height + i*grow) + + if shape == 'circle' or shape == 'square': + # Define the bounding box for the shape + left_up_point = (location_x - current_width // 2, location_y - current_height // 2) + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) + two_points = [left_up_point, right_down_point] + + if shape == 'circle': + draw.ellipse(two_points, fill=color) + elif shape == 'square': + draw.rectangle(two_points, fill=color) + + elif shape == 'triangle': + # Define the points for the triangle + left_up_point = (location_x - current_width // 2, location_y + current_height // 2) # bottom left + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) # bottom right + top_point = (location_x, location_y - current_height // 2) # top point + draw.polygon([top_point, left_up_point, right_down_point], fill=color) + + image = pil2tensor(image) + mask = image[:, :, :, 0] + out.append(mask) + outstack = torch.cat(out, dim=0) + return (outstack, 1.0 - outstack,) + +class CreateVoronoiMask: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createvoronoi" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}), + "num_points": ("INT", {"default": 15,"min": 1, "max": 4096, "step": 1}), + "line_width": ("INT", {"default": 4,"min": 1, "max": 4096, "step": 1}), + "speed": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + }, + } + + def createvoronoi(self, frames, num_points, line_width, speed, frame_width, frame_height): + from scipy.spatial import Voronoi + from matplotlib import pyplot as plt + # Define the number of images in the batch + batch_size = frames + out = [] + + # Calculate aspect ratio + aspect_ratio = frame_width / frame_height + + # Create start and end points for each point, considering the aspect ratio + start_points = np.random.rand(num_points, 2) + start_points[:, 0] *= aspect_ratio + + end_points = np.random.rand(num_points, 2) + end_points[:, 0] *= aspect_ratio + + for i in range(batch_size): + # Interpolate the points' positions based on the current frame + t = (i * speed) / (batch_size - 1) # normalize to [0, 1] over the frames + t = np.clip(t, 0, 1) # ensure t is in [0, 1] + points = (1 - t) * start_points + t * end_points # lerp + + # Adjust points for aspect ratio + points[:, 0] *= aspect_ratio + + vor = Voronoi(points) + + # Create a blank image with a white background + fig, ax = plt.subplots() + plt.subplots_adjust(left=0, right=1, bottom=0, top=1) + ax.set_xlim([0, aspect_ratio]); ax.set_ylim([0, 1]) # adjust x limits + ax.axis('off') + ax.margins(0, 0) + fig.set_size_inches(aspect_ratio * frame_height/100, frame_height/100) # adjust figure size + ax.fill_between([0, 1], [0, 1], color='white') + + # Plot each Voronoi ridge + for simplex in vor.ridge_vertices: + simplex = np.asarray(simplex) + if np.all(simplex >= 0): + plt.plot(vor.vertices[simplex, 0], vor.vertices[simplex, 1], 'k-', linewidth=line_width) + + fig.canvas.draw() + img = np.array(fig.canvas.renderer._renderer) + + plt.close(fig) + + pil_img = Image.fromarray(img).convert("L") + mask = torch.tensor(np.array(pil_img)) / 255.0 + + out.append(mask) + + return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) + +class GetMaskSizeAndCount: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mask": ("MASK",), + }} + + RETURN_TYPES = ("MASK","INT", "INT", "INT",) + RETURN_NAMES = ("mask", "width", "height", "count",) + FUNCTION = "getsize" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Returns the width, height and batch size of the mask, +and passes it through unchanged. + +""" + + def getsize(self, mask): + width = mask.shape[2] + height = mask.shape[1] + count = mask.shape[0] + return {"ui": { + "text": [f"{count}x{width}x{height}"]}, + "result": (mask, width, height, count) + } + +class GrowMaskWithBlur: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), + "incremental_expandrate": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), + "tapered_corners": ("BOOLEAN", {"default": True}), + "flip_input": ("BOOLEAN", {"default": False}), + "blur_radius": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 100, + "step": 0.1 + }), + "lerp_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "decay_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "fill_holes": ("BOOLEAN", {"default": False}), + }, + } + + CATEGORY = "KJNodes/masking" + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "expand_mask" + DESCRIPTION = """ +# GrowMaskWithBlur +- mask: Input mask or mask batch +- expand: Expand or contract mask or mask batch by a given amount +- incremental_expandrate: increase expand rate by a given amount per frame +- tapered_corners: use tapered corners +- flip_input: flip input mask +- blur_radius: value higher than 0 will blur the mask +- lerp_alpha: alpha value for interpolation between frames +- decay_factor: decay value for interpolation between frames +- fill_holes: fill holes in the mask (slow)""" + + def expand_mask(self, mask, expand, tapered_corners, flip_input, blur_radius, incremental_expandrate, lerp_alpha, decay_factor, fill_holes=False): + import kornia.morphology as morph + alpha = lerp_alpha + decay = decay_factor + if flip_input: + mask = 1.0 - mask + + growmask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) + out = [] + previous_output = None + current_expand = expand + for m in tqdm(growmask, desc="Expanding/Contracting Mask"): + output = m.unsqueeze(0).unsqueeze(0).to(main_device) # Add batch and channel dims for kornia + if abs(round(current_expand)) > 0 and output.max() > 0: + # Create kernel - kornia expects kernel on same device as input + if tapered_corners: + kernel = torch.tensor([[0, 1, 0], + [1, 1, 1], + [0, 1, 0]], dtype=torch.float32, device=output.device) + else: + kernel = torch.tensor([[1, 1, 1], + [1, 1, 1], + [1, 1, 1]], dtype=torch.float32, device=output.device) + + for _ in range(abs(round(current_expand))): + if current_expand < 0: + output = morph.erosion(output, kernel) + else: + output = morph.dilation(output, kernel) + + output = output.squeeze(0).squeeze(0) # Remove batch and channel dims + + if current_expand < 0: + current_expand -= abs(incremental_expandrate) + else: + current_expand += abs(incremental_expandrate) + + if fill_holes: + binary_mask = output > 0 + output_np = binary_mask.cpu().numpy() + filled = scipy.ndimage.binary_fill_holes(output_np) + output = torch.from_numpy(filled.astype(np.float32)).to(output.device) + + if alpha < 1.0 and previous_output is not None: + output = alpha * output + (1 - alpha) * previous_output + if decay < 1.0 and previous_output is not None: + output += decay * previous_output + output = output / output.max() + previous_output = output + out.append(output.cpu()) + + if blur_radius != 0: + # Convert the tensor list to PIL images, apply blur, and convert back + for idx, tensor in enumerate(out): + # Convert tensor to PIL image + pil_image = tensor2pil(tensor.cpu().detach())[0] + # Apply Gaussian blur + pil_image = pil_image.filter(ImageFilter.GaussianBlur(blur_radius)) + # Convert back to tensor + out[idx] = pil2tensor(pil_image) + blurred = torch.cat(out, dim=0) + return (blurred, 1.0 - blurred) + else: + return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) + +class MaskBatchMulti: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "mask_1": ("MASK", ), + "mask_2": ("MASK", ), + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("masks",) + FUNCTION = "combine" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Creates an image batch from multiple masks. +You can set how many inputs the node has, +with the **inputcount** and clicking update. +""" + + def combine(self, inputcount, **kwargs): + mask = kwargs["mask_1"] + for c in range(1, inputcount): + new_mask = kwargs[f"mask_{c + 1}"] + if mask.shape[1:] != new_mask.shape[1:]: + new_mask = F.interpolate(new_mask.unsqueeze(1), size=(mask.shape[1], mask.shape[2]), mode="bicubic").squeeze(1) + mask = torch.cat((mask, new_mask), dim=0) + return (mask,) + +class OffsetMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + "x": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "y": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "angle": ("INT", { "default": 0, "min": -360, "max": 360, "step": 1, "display": "number" }), + "duplication_factor": ("INT", { "default": 1, "min": 1, "max": 1000, "step": 1, "display": "number" }), + "roll": ("BOOLEAN", { "default": False }), + "incremental": ("BOOLEAN", { "default": False }), + "padding_mode": ( + [ + 'empty', + 'border', + 'reflection', + + ], { + "default": 'empty' + }), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "offset" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Offsets the mask by the specified amount. + - mask: Input mask or mask batch + - x: Horizontal offset + - y: Vertical offset + - angle: Angle in degrees + - roll: roll edge wrapping + - duplication_factor: Number of times to duplicate the mask to form a batch + - border padding_mode: Padding mode for the mask +""" + + def offset(self, mask, x, y, angle, roll=False, incremental=False, duplication_factor=1, padding_mode="empty"): + # Create duplicates of the mask batch + mask = mask.repeat(duplication_factor, 1, 1).clone() + + batch_size, height, width = mask.shape + + if angle != 0 and incremental: + for i in range(batch_size): + rotation_angle = angle * (i+1) + mask[i] = TF.rotate(mask[i].unsqueeze(0), rotation_angle).squeeze(0) + elif angle > 0: + for i in range(batch_size): + mask[i] = TF.rotate(mask[i].unsqueeze(0), angle).squeeze(0) + + if roll: + if incremental: + for i in range(batch_size): + shift_x = min(x*(i+1), width-1) + shift_y = min(y*(i+1), height-1) + if shift_x != 0: + mask[i] = torch.roll(mask[i], shifts=shift_x, dims=1) + if shift_y != 0: + mask[i] = torch.roll(mask[i], shifts=shift_y, dims=0) + else: + shift_x = min(x, width-1) + shift_y = min(y, height-1) + if shift_x != 0: + mask = torch.roll(mask, shifts=shift_x, dims=2) + if shift_y != 0: + mask = torch.roll(mask, shifts=shift_y, dims=1) + else: + + for i in range(batch_size): + if incremental: + temp_x = min(x * (i+1), width-1) + temp_y = min(y * (i+1), height-1) + else: + temp_x = min(x, width-1) + temp_y = min(y, height-1) + if temp_x > 0: + if padding_mode == 'empty': + mask[i] = torch.cat([torch.zeros((height, temp_x)), mask[i, :, :-temp_x]], dim=1) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, :, :-temp_x], (0, temp_x), mode=padding_mode) + elif temp_x < 0: + if padding_mode == 'empty': + mask[i] = torch.cat([mask[i, :, :temp_x], torch.zeros((height, -temp_x))], dim=1) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, :, -temp_x:], (temp_x, 0), mode=padding_mode) + + if temp_y > 0: + if padding_mode == 'empty': + mask[i] = torch.cat([torch.zeros((temp_y, width)), mask[i, :-temp_y, :]], dim=0) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, :-temp_y, :], (0, temp_y), mode=padding_mode) + elif temp_y < 0: + if padding_mode == 'empty': + mask[i] = torch.cat([mask[i, :temp_y, :], torch.zeros((-temp_y, width))], dim=0) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, -temp_y:, :], (temp_y, 0), mode=padding_mode) + + return mask, + +class RoundMask: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mask": ("MASK",), + }} + + RETURN_TYPES = ("MASK",) + FUNCTION = "round" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Rounds the mask or batch of masks to a binary mask. +RoundMask example + +""" + + def round(self, mask): + mask = mask.round() + return (mask,) + +class ResizeMask: + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "keep_proportions": ("BOOLEAN", { "default": False }), + "upscale_method": (s.upscale_methods,), + "crop": (["disabled","center"],), + } + } + + RETURN_TYPES = ("MASK", "INT", "INT",) + RETURN_NAMES = ("mask", "width", "height",) + FUNCTION = "resize" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Resizes the mask or batch of masks to the specified width and height. +""" + + def resize(self, mask, width, height, keep_proportions, upscale_method,crop): + if keep_proportions: + _, oh, ow = mask.shape + width = ow if width == 0 else width + height = oh if height == 0 else height + ratio = min(width / ow, height / oh) + width = round(ow*ratio) + height = round(oh*ratio) + + if upscale_method == "lanczos": + out_mask = common_upscale(mask.unsqueeze(1).repeat(1, 3, 1, 1), width, height, upscale_method, crop=crop).movedim(1,-1)[:, :, :, 0] + else: + out_mask = common_upscale(mask.unsqueeze(1), width, height, upscale_method, crop=crop).squeeze(1) + + return(out_mask, out_mask.shape[2], out_mask.shape[1],) + +class RemapMaskRange: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + "min": ("FLOAT", {"default": 0.0,"min": -10.0, "max": 1.0, "step": 0.01}), + "max": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 10.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "remap" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Sets new min and max values for the mask. +""" + + def remap(self, mask, min, max): + + # Find the maximum value in the mask + mask_max = torch.max(mask) + + # If the maximum mask value is zero, avoid division by zero by setting it to 1 + mask_max = mask_max if mask_max > 0 else 1 + + # Scale the mask values to the new range defined by min and max + # The highest pixel value in the mask will be scaled to max + scaled_mask = (mask / mask_max) * (max - min) + min + + # Clamp the values to ensure they are within [0.0, 1.0] + scaled_mask = torch.clamp(scaled_mask, min=0.0, max=1.0) + + return (scaled_mask, ) + + +def get_mask_polygon(self, mask_np): + import cv2 + """Helper function to get polygon points from mask""" + # Find contours + contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if not contours: + return None + + # Get the largest contour + largest_contour = max(contours, key=cv2.contourArea) + + # Approximate polygon + epsilon = 0.02 * cv2.arcLength(largest_contour, True) + polygon = cv2.approxPolyDP(largest_contour, epsilon, True) + + return polygon.squeeze() + +class SeparateMasks: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK", ), + "size_threshold_width" : ("INT", {"default": 256, "min": 0.0, "max": 4096, "step": 1}), + "size_threshold_height" : ("INT", {"default": 256, "min": 0.0, "max": 4096, "step": 1}), + "mode": (["convex_polygons", "area", "box"],), + "max_poly_points": ("INT", {"default": 8, "min": 3, "max": 32, "step": 1}), + + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "separate" + CATEGORY = "KJNodes/masking" + OUTPUT_NODE = True + DESCRIPTION = "Separates a mask into multiple masks based on the size of the connected components." + + def polygon_to_mask(self, polygon, shape): + import cv2 + mask = np.zeros((shape[0], shape[1]), dtype=np.uint8) # Fixed shape handling + + if len(polygon.shape) == 2: # Check if polygon points are valid + polygon = polygon.astype(np.int32) + cv2.fillPoly(mask, [polygon], 1) + return mask + + def get_mask_polygon(self, mask_np, max_points): + import cv2 + contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + return None + + largest_contour = max(contours, key=cv2.contourArea) + hull = cv2.convexHull(largest_contour) + + # Initialize with smaller epsilon for more points + perimeter = cv2.arcLength(hull, True) + + min_eps = perimeter * 0.001 # Much smaller minimum + max_eps = perimeter * 0.2 # Smaller maximum + + best_approx = None + best_diff = float('inf') + max_iterations = 20 + + #print(f"Target points: {max_points}, Perimeter: {perimeter}") + + for i in range(max_iterations): + curr_eps = (min_eps + max_eps) / 2 + approx = cv2.approxPolyDP(hull, curr_eps, True) + points_diff = len(approx) - max_points + + #print(f"Iteration {i}: points={len(approx)}, eps={curr_eps:.4f}") + + if abs(points_diff) < best_diff: + best_approx = approx + best_diff = abs(points_diff) + + if len(approx) > max_points: + min_eps = curr_eps * 1.1 # More gradual adjustment + elif len(approx) < max_points: + max_eps = curr_eps * 0.9 # More gradual adjustment + else: + return approx.squeeze() + + if abs(max_eps - min_eps) < perimeter * 0.0001: # Relative tolerance + break + + # If we didn't find exact match, return best approximation + return best_approx.squeeze() if best_approx is not None else hull.squeeze() + + def separate(self, mask: torch.Tensor, size_threshold_width: int, size_threshold_height: int, max_poly_points: int, mode: str): + B, H, W = mask.shape + separated = [] + + mask = mask.round() + + for b in range(B): + mask_np = mask[b].cpu().numpy().astype(np.uint8) + structure = np.ones((3, 3), dtype=np.int8) + labeled, ncomponents = scipy.ndimage.label(mask_np, structure=structure) + pbar = ProgressBar(ncomponents) + + for component in range(1, ncomponents + 1): + component_mask_np = (labeled == component).astype(np.uint8) + + rows = np.any(component_mask_np, axis=1) + cols = np.any(component_mask_np, axis=0) + y_min, y_max = np.where(rows)[0][[0, -1]] + x_min, x_max = np.where(cols)[0][[0, -1]] + + width = x_max - x_min + 1 + height = y_max - y_min + 1 + centroid_x = (x_min + x_max) / 2 # Calculate x centroid + logging.info(f"Component {component}: width={width}, height={height}, x_pos={centroid_x}") + + if width >= size_threshold_width and height >= size_threshold_height: + if mode == "convex_polygons": + polygon = self.get_mask_polygon(component_mask_np, max_poly_points) + if polygon is not None: + poly_mask = self.polygon_to_mask(polygon, (H, W)) + poly_mask = torch.tensor(poly_mask, device=mask.device) + separated.append((centroid_x, poly_mask)) + elif mode == "box": + # Create bounding box mask + box_mask = np.zeros((H, W), dtype=np.uint8) + box_mask[y_min:y_max+1, x_min:x_max+1] = 1 + box_mask = torch.tensor(box_mask, device=mask.device) + separated.append((centroid_x, box_mask)) + else: + area_mask = torch.tensor(component_mask_np, device=mask.device) + separated.append((centroid_x, area_mask)) + pbar.update(1) + + if len(separated) > 0: + # Sort by x position and extract only the masks + separated.sort(key=lambda x: x[0]) + separated = [x[1] for x in separated] + out_masks = torch.stack(separated, dim=0) + return out_masks, + else: + return torch.empty((1, 64, 64), device=mask.device), + + +class ConsolidateMasksKJ: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "masks": ("MASK",), + "width": ("INT", {"default": 512, "min": 0, "max": 4096, "step": 64}), + "height": ("INT", {"default": 512, "min": 0, "max": 4096, "step": 64}), + "padding": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + }, + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "consolidate" + + CATEGORY = "KJNodes/masking" + DESCRIPTION = "Consolidates a batch of separate masks by finding the largest group of masks that fit inside a tile of the given width and height (including the padding), and repeating until no more masks can be combined." + + def consolidate(self, masks, width=512, height=512, padding=0): + B, H, W = masks.shape + + def mask_fits(coords, candidate_coords): + x_min, y_min, x_max, y_max = coords + cx_min, cy_min, cx_max, cy_max = candidate_coords + nx_min, ny_min = min(x_min, cx_min), min(y_min, cy_min) + nx_max, ny_max = max(x_max, cx_max), max(y_max, cy_max) + if nx_min + width < nx_max + padding or ny_min + height < ny_max + padding: + return False, coords + return True, (nx_min, ny_min, nx_max, ny_max) + + separated = [] + final_masks = [] + for b in range(B): + m = masks[b] + rows, cols = m.any(dim=1), m.any(dim=0) + y_min, y_max = torch.where(rows)[0][[0, -1]] + x_min, x_max = torch.where(cols)[0][[0, -1]] + w = x_max - x_min + 1 + h = y_max - y_min + 1 + separated.append(((x_min.item(), y_min.item(), x_max.item(), y_max.item()), m)) + + separated.sort(key=lambda x: x[0]) + fits = [] + for i, masks in enumerate(separated): + coord = masks[0] + fits_in_box = [] + for j, cand_mask in enumerate(separated): + if i == j: + continue + r, coord = mask_fits(coord, cand_mask[0]) + if r: + fits_in_box.append(j) + fits.append((i, fits_in_box)) + fits.sort(key=lambda x: -len(x[1])) + seen = [] + unique_fits = [] + for idx, fs in fits: + uniq = [i for i in fs if i not in seen] + unique_fits.append((idx, fs, uniq)) + seen.extend(uniq) + unique_fits.sort(key=lambda x: (-len(x[1]), -len(x[2]))) + merged = [] + for mask_idx, fitting_masks, _ in unique_fits: + if mask_idx in merged: + continue + fitting_masks = [i for i in fitting_masks if i not in merged] + combined_mask = separated[mask_idx][1].clone() + for i in fitting_masks: + combined_mask += separated[i][1] + merged.append(i) + merged.append(mask_idx) + final_masks.append(combined_mask) + + logging.info(f"Consolidated {B} masks into {len(final_masks)}") + return (torch.stack(final_masks, dim=0),) + + +class DrawMaskOnImage: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE", ), + "mask": ("MASK", ), + "color": ("STRING", {"default": "0, 0, 0", "tooltip": "Color as RGB/RGBA values in range 0-255 or 0.0-1.0, separated by commas. Ex: 255, 0, 0, 128"}), + }, + "optional": { + "device": (["cpu", "gpu"], {"default": "cpu", "tooltip": "Device to use for processing"}), + } + } + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("images",) + FUNCTION = "apply" + CATEGORY = "KJNodes/masking" + DESCRIPTION = "Applies the provided masks to the input images with Alpha Blending support." + + def apply(self, image, mask, color, device="cpu"): + B, H, W, C = image.shape + BM, HM, WM = mask.shape + + processing_device = main_device if device == "gpu" else torch.device("cpu") + + in_masks = mask.clone().to(processing_device) + in_images = image.clone().to(processing_device) + + # Resize mask if dimensions don't match + if HM != H or WM != W: + in_masks = F.interpolate(mask.unsqueeze(1), size=(H, W), mode='nearest-exact').squeeze(1) + # Handle batch size mismatch + if B > BM: + in_masks = in_masks.repeat((B + BM - 1) // BM, 1, 1)[:B] + elif BM > B: + in_masks = in_masks[:B] + + output_images = [] + + # Parse Color String (Handle RGB, RGBA, and Hex formats) + color = color.strip() + color_values = [] + + if color.startswith('#'): + # Handle hex format (#RGB, #RGBA, #RRGGBB, #RRGGBBAA) + hex_color = color.lstrip('#') + if len(hex_color) == 3: # #RGB + color_values = [int(c*2, 16) / 255.0 for c in hex_color] + elif len(hex_color) == 4: # #RGBA + color_values = [int(c*2, 16) / 255.0 for c in hex_color] + elif len(hex_color) == 6: # #RRGGBB + color_values = [int(hex_color[i:i+2], 16) / 255.0 for i in (0, 2, 4)] + elif len(hex_color) == 8: # #RRGGBBAA + color_values = [int(hex_color[i:i+2], 16) / 255.0 for i in (0, 2, 4, 6)] + else: + raise ValueError(f"Invalid hex color format: {color}") + else: + # Handle comma-separated RGB/RGBA format + for x in color.split(","): + val = float(x.strip()) + color_values.append(val / 255.0 if val > 1.0 else val) + + rgb = color_values[:3] + alpha_val = color_values[3] if len(color_values) == 4 else 1.0 + + fill_color = torch.tensor(rgb, dtype=torch.float32, device=processing_device) + + for i in tqdm(range(B), desc="DrawMaskOnImage batch"): + curr_mask = in_masks[i] # [H, W] + img_idx = min(i, B - 1) + curr_image = in_images[img_idx] # [H, W, C] + + blend_factor = curr_mask.unsqueeze(-1) * alpha_val + img_channels = curr_image.shape[-1] + + if img_channels == 4: + img_rgb = curr_image[..., :3] + img_a = curr_image[..., 3:] + out_rgb = img_rgb * (1 - blend_factor) + fill_color * blend_factor + out_a = torch.maximum(img_a, blend_factor) + masked_image = torch.cat((out_rgb, out_a), dim=-1) + else: + masked_image = curr_image * (1 - blend_factor) + fill_color * blend_factor + output_images.append(masked_image) + + if not output_images: + return (torch.zeros((0, H, W, C), dtype=image.dtype),) + + out_tensor = torch.stack(output_images, dim=0).cpu() + + return (out_tensor, ) + +class BlockifyMask: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "masks": ("MASK",), + "block_size": ("INT", {"default": 32, "min": 8, "max": 512, "step": 1, "tooltip": "Size of blocks in pixels (smaller = smaller blocks)"}), + }, + "optional": { + "device": (["cpu", "gpu"], {"default": "cpu", "tooltip": "Device to use for processing"}), + } + } + + RETURN_TYPES = ("MASK", ) + RETURN_NAMES = ("mask",) + FUNCTION = "process" + CATEGORY = "KJNodes/masking" + DESCRIPTION = "Creates a block mask by dividing the bounding box of each mask into blocks of the specified size and filling in blocks that contain any part of the original mask." + + def process(self, masks, block_size, device="cpu"): + processing_device = main_device if device == "gpu" else torch.device("cpu") + + masks = masks.to(processing_device) + batch_size, height, width = masks.shape + + result_masks = torch.zeros_like(masks) + + for i in tqdm(range(batch_size), desc="BlockifyMask batch"): + mask = masks[i] + + # Find bounding box efficiently + mask_bool = mask > 0 + if not mask_bool.any(): + continue + + y_indices = torch.nonzero(mask_bool.any(dim=1), as_tuple=True)[0] + x_indices = torch.nonzero(mask_bool.any(dim=0), as_tuple=True)[0] + + if len(y_indices) == 0 or len(x_indices) == 0: + continue + + y_min, y_max = y_indices[0], y_indices[-1] + x_min, x_max = x_indices[0], x_indices[-1] + + bbox_width = x_max - x_min + 1 + bbox_height = y_max - y_min + 1 + + # Calculate block grid + w_divisions = max(1, bbox_width // block_size) + h_divisions = max(1, bbox_height // block_size) + + w_slice = bbox_width // w_divisions + h_slice = bbox_height // h_divisions + + # Create coordinate grids only for bbox region + y_coords = torch.arange(y_min, y_max + 1, device=processing_device).view(-1, 1) + x_coords = torch.arange(x_min, x_max + 1, device=processing_device).view(1, -1) + + # Calculate block indices for bbox region + w_block_indices = (x_coords - x_min) // w_slice + h_block_indices = (y_coords - y_min) // h_slice + + # Clamp to valid range + w_block_indices = w_block_indices.clamp(0, w_divisions - 1) + h_block_indices = h_block_indices.clamp(0, h_divisions - 1) + + # Create unique block IDs by combining h and w indices + block_ids = h_block_indices * w_divisions + w_block_indices + + # Get mask region within bbox + mask_region = mask[y_min:y_max+1, x_min:x_max+1] + + # Find which blocks have content using scatter_add + max_blocks = h_divisions * w_divisions + block_content = torch.zeros(max_blocks, device=processing_device) + block_content.scatter_add_(0, block_ids.flatten(), mask_region.flatten()) + + # Create result for blocks that have content + has_content = block_content > 0 + block_mask = has_content[block_ids] + + # Fill the result + result_masks[i, y_min:y_max+1, x_min:x_max+1] = block_mask.float() + + return (result_masks.clamp(0, 1),) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/model_optimization_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/model_optimization_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..59d2e3921b2260d168275e2cfc7da9b96ddc4134 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/model_optimization_nodes.py @@ -0,0 +1,2375 @@ +import os +import sys +import logging +import torch +import importlib +import math +import datetime +from tqdm import tqdm + +import folder_paths +import comfy.model_management as mm +from comfy.cli_args import args +from comfy.ldm.modules.attention import wrap_attn, optimized_attention, attention_pytorch +import comfy.utils +import comfy.sd +import comfy.ops + +try: + from comfy_api.latest import io + v3_available = True +except ImportError: + v3_available = False + logging.warning("ComfyUI v3 node API not available, please update ComfyUI to access latest v3 nodes.") + +sageattn_modes = ["disabled", "auto", "sageattn_qk_int8_pv_fp16_cuda", "sageattn_qk_int8_pv_fp16_triton", "sageattn_qk_int8_pv_fp8_cuda", "sageattn_qk_int8_pv_fp8_cuda++", "sageattn3", "sageattn3_per_block_mean"] + +def get_sage_func(sage_attention, allow_compile=False): + logging.info(f"Using sage attention mode: {sage_attention}") + if sage_attention == "auto": + from sageattention import sageattn + def sage_func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD"): + return sageattn(q, k, v, is_causal=is_causal, attn_mask=attn_mask, tensor_layout=tensor_layout) + elif sage_attention == "sageattn_qk_int8_pv_fp16_cuda": + from sageattention import sageattn_qk_int8_pv_fp16_cuda + def sage_func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD"): + return sageattn_qk_int8_pv_fp16_cuda(q, k, v, is_causal=is_causal, attn_mask=attn_mask, pv_accum_dtype="fp32", tensor_layout=tensor_layout) + elif sage_attention == "sageattn_qk_int8_pv_fp16_triton": + from sageattention import sageattn_qk_int8_pv_fp16_triton + def sage_func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD"): + return sageattn_qk_int8_pv_fp16_triton(q, k, v, is_causal=is_causal, attn_mask=attn_mask, tensor_layout=tensor_layout) + elif sage_attention == "sageattn_qk_int8_pv_fp8_cuda": + from sageattention import sageattn_qk_int8_pv_fp8_cuda + def sage_func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD"): + return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, attn_mask=attn_mask, pv_accum_dtype="fp32+fp32", tensor_layout=tensor_layout) + elif sage_attention == "sageattn_qk_int8_pv_fp8_cuda++": + from sageattention import sageattn_qk_int8_pv_fp8_cuda + def sage_func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD"): + return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, attn_mask=attn_mask, pv_accum_dtype="fp32+fp16", tensor_layout=tensor_layout) + elif "sageattn3" in sage_attention: + from sageattn3 import sageattn3_blackwell + def sage_func(q, k, v, is_causal=False, attn_mask=None, tensor_layout="NHD", **kwargs): + q, k, v = [x.transpose(1, 2) if tensor_layout == "NHD" else x for x in (q, k, v)] + out = sageattn3_blackwell(q, k, v, is_causal=is_causal, attn_mask=attn_mask, per_block_mean=(sage_attention == "sageattn3_per_block_mean")) + return out.transpose(1, 2) if tensor_layout == "NHD" else out + + if not allow_compile: + sage_func = torch.compiler.disable()(sage_func) + + @wrap_attn + def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + if kwargs.get("low_precision_attention", True) is False: + return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs) + in_dtype = v.dtype + if q.dtype == torch.float32 or k.dtype == torch.float32 or v.dtype == torch.float32: + q, k, v = q.to(torch.float16), k.to(torch.float16), v.to(torch.float16) + if skip_reshape: + b, _, _, dim_head = q.shape + tensor_layout="HND" + else: + b, _, dim_head = q.shape + dim_head //= heads + q, k, v = map( + lambda t: t.view(b, -1, heads, dim_head), + (q, k, v), + ) + tensor_layout="NHD" + if mask is not None: + # add a batch dimension if there isn't already one + if mask.ndim == 2: + mask = mask.unsqueeze(0) + # add a heads dimension if there isn't already one + if mask.ndim == 3: + mask = mask.unsqueeze(1) + out = sage_func(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout).to(in_dtype) + if tensor_layout == "HND": + if not skip_output_reshape: + out = ( + out.transpose(1, 2).reshape(b, -1, heads * dim_head) + ) + else: + if skip_output_reshape: + out = out.transpose(1, 2) + else: + out = out.reshape(b, -1, heads * dim_head) + return out + return attention_sage + + +from comfy.patcher_extension import CallbacksMP +class PathchSageAttentionKJ(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "sage_attention": (sageattn_modes, {"default": False, "tooltip": "Patch the attention of the model passing through this node to use sageattn. To revert, run this node again with the disabled option. Requires the sageattention library to be installed."}), + }, + "optional": { + "allow_compile": ("BOOLEAN", {"default": False, "tooltip": "Allow the use of torch.compile for the sage attention function, requires latest sageattn 2.2.0 or higher."}) + } + } + + RETURN_TYPES = ("MODEL", ) + FUNCTION = "patch" + DESCRIPTION = "Experimental node for patching attention mode. This doesn't use the model patching system and thus can't be disabled without running the node again with 'disabled' option." + EXPERIMENTAL = True + CATEGORY = "KJNodes/experimental" + + def patch(self, model, sage_attention, allow_compile=False): + if sage_attention == "disabled": + return model, + + model_clone = model.clone() + + new_attention = get_sage_func(sage_attention, allow_compile=allow_compile) + def attention_override_sage(func, *args, **kwargs): + return new_attention.__wrapped__(*args, **kwargs) + + # attention override + model_clone.model_options["transformer_options"]["optimized_attention_override"] = attention_override_sage + + return model_clone, + + +def get_flash_func(allow_compile=False, cast_dtype=torch.float16): + # Prefer FA2 (broad arch support, plain tensor return); fall back to FA3 + # (flash_attn_interface), which has no dropout arg and returns (out, lse). + is_fa3 = False + try: + from flash_attn import flash_attn_func + except ImportError: + try: + from flash_attn_interface import flash_attn_func + is_fa3 = True + except ImportError: + raise ImportError( + "Flash attention not found. Install either FA2 ('flash_attn') or " + "FA3 ('flash_attn_interface', pip package 'flash-attn-3')." + ) + logging.info(f"Using flash attention {'3' if is_fa3 else '2'}: cast_dtype={cast_dtype}") + + # q, k, v in NHD layout (b, seq, heads, dim_head) + def flash_func(q, k, v): + if is_fa3: + out = flash_attn_func(q, k, v, causal=False) + else: + out = flash_attn_func(q, k, v, dropout_p=0.0, causal=False) + # FA3 returns (out, softmax_lse); FA2 returns the tensor directly + return out[0] if isinstance(out, tuple) else out + + if not allow_compile: + flash_func = torch.compiler.disable()(flash_func) + + if torch.cuda.is_available(): + probe = torch.zeros(1, 8, 2, 64, dtype=cast_dtype, device="cuda") + flash_func(probe, probe, probe) + + @wrap_attn + def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + if mask is not None: + raise RuntimeError("Flash attention does not support attention masks") + in_dtype = v.dtype + # flash_attn only supports fp16/bf16 + if q.dtype == torch.float32 or k.dtype == torch.float32 or v.dtype == torch.float32: + q, k, v = q.to(cast_dtype), k.to(cast_dtype), v.to(cast_dtype) + # flash_attn wants NHD layout (b, seq, heads, dim_head) + if skip_reshape: + # input is HND (b, heads, seq, dim_head) + b, _, _, dim_head = q.shape + q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) + else: + b, _, dim_head = q.shape + dim_head //= heads + q, k, v = map(lambda t: t.view(b, -1, heads, dim_head), (q, k, v)) + out = flash_func(q, k, v).to(in_dtype) + if skip_output_reshape: + out = out.transpose(1, 2) # NHD -> HND + else: + out = out.reshape(b, -1, heads * dim_head) + return out + return attention_flash + + +class PatchFlashAttentionKJ(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + }, + "optional": { + "allow_compile": ("BOOLEAN", {"default": False, "tooltip": "Allow torch.compile to trace into the flash attention function. If disabled (default), the function is wrapped with torch.compiler.disable() for compatibility, matching the sage attention node."}), + }} + + RETURN_TYPES = ("MODEL", ) + FUNCTION = "patch" + DESCRIPTION = "Experimental node for patching attention to use flash attention, without the silent SDPA fallback the ComfyUI default does. Patches the attention of the model passing through this node; to disable, bypass or disconnect this node. Requires the flash_attn library to be installed." + EXPERIMENTAL = True + CATEGORY = "KJNodes/experimental" + + def patch(self, model, allow_compile=False): + # match the model's compute dtype for the fp32 downcast, fall back to fp16 + inference_dtype = model.model.get_dtype_inference() if hasattr(model.model, "get_dtype_inference") else torch.float16 + cast_dtype = inference_dtype if inference_dtype in (torch.float16, torch.bfloat16) else torch.float16 + + new_attention = get_flash_func( + allow_compile=allow_compile, + cast_dtype=cast_dtype, + ) + + model_clone = model.clone() + def attention_override_flash(func, *args, **kwargs): + return new_attention.__wrapped__(*args, **kwargs) + + # attention override + model_clone.model_options["transformer_options"]["optimized_attention_override"] = attention_override_flash + + return model_clone, + + +class CheckpointLoaderKJ(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}), + "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"],), + "compute_dtype": (["default", "fp16", "bf16", "fp32"], {"default": "default", "tooltip": "The compute dtype to use for the model."}), + "patch_cublaslinear": ("BOOLEAN", {"default": False, "tooltip": "Enable or disable the cublas_ops arg"}), + "sage_attention": (sageattn_modes, {"default": False, "tooltip": "Patch comfy attention to use sageattn."}), + "enable_fp16_accumulation": ("BOOLEAN", {"default": False, "tooltip": "Enable torch.backends.cuda.matmul.allow_fp16_accumulation, required minimum pytorch version 2.7.1"}), + }} + + RETURN_TYPES = ("MODEL", "CLIP", "VAE") + FUNCTION = "load" + DESCRIPTION = "Experimental node for patching torch.nn.Linear with CublasLinear." + EXPERIMENTAL = True + CATEGORY = "KJNodes/model_loaders" + + def load(self, ckpt_name, weight_dtype, compute_dtype, patch_cublaslinear, sage_attention, enable_fp16_accumulation): + DTYPE_MAP = { + "fp8_e4m3fn": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32 + } + model_options = {} + if dtype := DTYPE_MAP.get(weight_dtype): + model_options["dtype"] = dtype + logging.info(f"Setting {ckpt_name} weight dtype to {dtype}") + + if weight_dtype == "fp8_e4m3fn_fast": + model_options["dtype"] = torch.float8_e4m3fn + model_options["fp8_optimizations"] = True + + if patch_cublaslinear: + args.fast.add("cublas_ops") + else: + args.fast.discard("cublas_ops") + + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) + model, clip, vae, _ = comfy.sd.load_checkpoint_guess_config( + ckpt_path, + output_vae=True, + output_clip=True, + embedding_directory=folder_paths.get_folder_paths("embeddings"), + model_options=model_options) + + if dtype := DTYPE_MAP.get(compute_dtype): + model.set_model_compute_dtype(dtype) + model.force_cast_weights = False + logging.info(f"Setting {ckpt_name} compute dtype to {dtype}") + + if enable_fp16_accumulation: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + torch.backends.cuda.matmul.allow_fp16_accumulation = True + else: + raise RuntimeError("Failed to set fp16 accumulation, requires pytorch version 2.7.1 or higher") + else: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + torch.backends.cuda.matmul.allow_fp16_accumulation = False + + if sage_attention != "disabled": + new_attention = get_sage_func(sage_attention) + def attention_override_sage(func, *args, **kwargs): + return new_attention.__wrapped__(*args, **kwargs) + + # attention override + model.model_options["transformer_options"]["optimized_attention_override"] = attention_override_sage + + return model, clip, vae + + +class DiffusionModelSelector(): + @classmethod + def INPUT_TYPES(s): + ltx2_connector_models = folder_paths.get_filename_list("text_encoders") + ltx2_connector_models = [m for m in ltx2_connector_models if "connector" in m.lower()] + return {"required": { + "model_name": (folder_paths.get_filename_list("diffusion_models") + ltx2_connector_models, {"tooltip": "The name of the checkpoint (model) to load."}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("model_path",) + FUNCTION = "get_path" + DESCRIPTION = "Returns the path to the model as a string." + EXPERIMENTAL = True + CATEGORY = "KJNodes/model_loaders" + + def get_path(self, model_name): + if "connector" in model_name.lower(): + model_path = folder_paths.get_full_path_or_raise("text_encoders", model_name) + else: + model_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name) + return (model_path,) + +def _load_diffusion_model_kj(unet_path, model_options=None, extra_state_dict=None, disable_dynamic=False): + model_options = {} if model_options is None else dict(model_options) + + sd, metadata = comfy.utils.load_torch_file(unet_path, return_metadata=True) + if extra_state_dict is not None: + extra_sd = comfy.utils.load_torch_file(extra_state_dict) + sd.update(extra_sd) + del extra_sd + + diffusion_model_prefix = comfy.sd.model_detection.unet_prefix_from_state_dict(sd) + sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=False) + + model = comfy.sd.load_diffusion_model_state_dict( + sd, + model_options=model_options, + metadata=metadata, + disable_dynamic=disable_dynamic, + ) + + model.cached_patcher_init = (_load_diffusion_model_kj, (unet_path, model_options, extra_state_dict)) + return model + +class DiffusionModelLoaderKJ(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model_name": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "The name of the checkpoint (model) to load."}), + "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"],), + "compute_dtype": (["default", "fp16", "bf16", "fp32"], {"default": "default", "tooltip": "The compute dtype to use for the model."}), + "patch_cublaslinear": ("BOOLEAN", {"default": False, "tooltip": "Enable or disable the cublas_ops arg"}), + "sage_attention": (sageattn_modes, {"default": False, "tooltip": "Patch comfy attention to use sageattn."}), + "enable_fp16_accumulation": ("BOOLEAN", {"default": False, "tooltip": "Enable torch.backends.cuda.matmul.allow_fp16_accumulation, requires pytorch 2.7.0 nightly."}), + }, + "optional": { + "extra_state_dict": ("STRING", {"forceInput": True, "tooltip": "The full path to an additional state dict to load, this will be merged with the main state dict. Useful for example to add VACE module to a WanVideoModel. You can use DiffusionModelSelector to easily get the path."}), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch_and_load" + DESCRIPTION = "Node for patching torch.nn.Linear with CublasLinear." + EXPERIMENTAL = True + CATEGORY = "KJNodes/model_loaders" + + def patch_and_load(self, model_name, weight_dtype, compute_dtype, patch_cublaslinear, sage_attention, enable_fp16_accumulation, extra_state_dict=None): + DTYPE_MAP = { + "fp8_e4m3fn": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + "fp16": torch.float16, + "bf16": torch.bfloat16, + "fp32": torch.float32 + } + model_options = {} + if dtype := DTYPE_MAP.get(weight_dtype): + model_options["dtype"] = dtype + logging.info(f"Setting {model_name} weight dtype to {dtype}") + + if weight_dtype == "fp8_e4m3fn_fast": + model_options["dtype"] = torch.float8_e4m3fn + model_options["fp8_optimizations"] = True + + if enable_fp16_accumulation: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + torch.backends.cuda.matmul.allow_fp16_accumulation = True + else: + raise RuntimeError("Failed to set fp16 accumulation, this requires pytorch 2.7.1 or higher") + else: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + torch.backends.cuda.matmul.allow_fp16_accumulation = False + + if patch_cublaslinear: + args.fast.add("cublas_ops") + else: + args.fast.discard("cublas_ops") + + unet_path = folder_paths.get_full_path_or_raise("diffusion_models", model_name) + + model = _load_diffusion_model_kj(unet_path, model_options=model_options, extra_state_dict=extra_state_dict) + if dtype := DTYPE_MAP.get(compute_dtype): + model.set_model_compute_dtype(dtype) + model.force_cast_weights = False + logging.info(f"Setting {model_name} compute dtype to {dtype}") + + if sage_attention != "disabled": + new_attention = get_sage_func(sage_attention) + def attention_override_sage(func, *args, **kwargs): + return new_attention.__wrapped__(*args, **kwargs) + + # attention override + model.model_options["transformer_options"]["optimized_attention_override"] = attention_override_sage + + return (model,) + +class ModelPatchTorchSettings: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "enable_fp16_accumulation": ("BOOLEAN", {"default": False, "tooltip": "Enable torch.backends.cuda.matmul.allow_fp16_accumulation, requires pytorch 2.7.0 nightly."}), + }} + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + DESCRIPTION = "Adds callbacks to model to set torch settings before and after running the model." + EXPERIMENTAL = True + CATEGORY = "KJNodes/experimental" + + def patch(self, model, enable_fp16_accumulation): + model_clone = model.clone() + + def patch_enable_fp16_accum(model): + logging.info("Patching torch settings: torch.backends.cuda.matmul.allow_fp16_accumulation = True") + torch.backends.cuda.matmul.allow_fp16_accumulation = True + def patch_disable_fp16_accum(model): + logging.info("Patching torch settings: torch.backends.cuda.matmul.allow_fp16_accumulation = False") + torch.backends.cuda.matmul.allow_fp16_accumulation = False + + if enable_fp16_accumulation: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + model_clone.add_callback(CallbacksMP.ON_PRE_RUN, patch_enable_fp16_accum) + model_clone.add_callback(CallbacksMP.ON_CLEANUP, patch_disable_fp16_accum) + else: + raise RuntimeError("Failed to set fp16 accumulation, this requires pytorch 2.7.1 or higher") + else: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + model_clone.add_callback(CallbacksMP.ON_PRE_RUN, patch_disable_fp16_accum) + else: + raise RuntimeError("Failed to set fp16 accumulation, this requires pytorch 2.7.1 or higher") + + return (model_clone,) + + +class PatchModelPatcherOrder: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "patch_order": (["object_patch_first", "weight_patch_first"], {"default": "weight_patch_first", "tooltip": "Patch the comfy patch_model function to load weight patches (LoRAs) before compiling the model"}), + "full_load": (["enabled", "disabled", "auto"], {"default": "auto", "tooltip": "Disabling may help with memory issues when loading large models, when changing this you should probably force model reload to avoid issues!"}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "KJNodes/deprecated" + DESCRIPTION = "NO LONGER NECESSARY OR FUNCTIONAL, keeping node for backwards compatibility. Use the TorchCompileModelAdvanced to use LoRA with torch.compile." + DEPRECATED = True + + def patch(self, model, patch_order, full_load): + return model, + + +class TorchCompileModelFluxAdvancedV2: + def __init__(self): + self._compiled = False + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "backend": (["inductor", "cudagraphs"],), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + "double_blocks": ("BOOLEAN", {"default": True, "tooltip": "Compile double blocks"}), + "single_blocks": ("BOOLEAN", {"default": True, "tooltip": "Compile single blocks"}), + "dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}), + }, + "optional": { + "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), + "force_parameter_static_shapes": ("BOOLEAN", {"default": True, "tooltip": "torch._dynamo.config.force_parameter_static_shapes"}), + } + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "KJNodes/torchcompile" + EXPERIMENTAL = True + DEPRECATED = True + DESCRIPTION = "Deprecated, use TorchCompileModelAdvanced instead." + + def patch(self, model, backend, mode, fullgraph, single_blocks, double_blocks, dynamic, dynamo_cache_size_limit=64, force_parameter_static_shapes=True): + from comfy_api.torch_helpers import set_torch_compile_wrapper + m = model.clone() + diffusion_model = m.get_model_object("diffusion_model") + torch._dynamo.config.cache_size_limit = dynamo_cache_size_limit + torch._dynamo.config.force_parameter_static_shapes = force_parameter_static_shapes + + compile_key_list = [] + + try: + if double_blocks: + for i, block in enumerate(diffusion_model.double_blocks): + print("Adding double block to compile list", i) + compile_key_list.append(f"diffusion_model.double_blocks.{i}") + if single_blocks: + for i, block in enumerate(diffusion_model.single_blocks): + compile_key_list.append(f"diffusion_model.single_blocks.{i}") + + set_torch_compile_wrapper(model=m, keys=compile_key_list, backend=backend, mode=mode, dynamic=dynamic, fullgraph=fullgraph) + except Exception as e: + raise RuntimeError("Failed to compile model") from e + + return (m, ) + + +class TorchCompileModelWanVideoV2: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "backend": (["inductor","cudagraphs"], {"default": "inductor"}), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + "dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}), + "compile_transformer_blocks_only": ("BOOLEAN", {"default": True, "tooltip": "Compile only transformer blocks, faster compile and less error prone"}), + "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), + + }, + "optional": { + "force_parameter_static_shapes": ("BOOLEAN", {"default": True, "tooltip": "torch._dynamo.config.force_parameter_static_shapes"}), + }, + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "KJNodes/torchcompile" + EXPERIMENTAL = True + DEPRECATED = True + DESCRIPTION = "Deprecated, use TorchCompileModelAdvanced instead." + + def patch(self, model, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_transformer_blocks_only, force_parameter_static_shapes=True): + from comfy_api.torch_helpers import set_torch_compile_wrapper + m = model.clone() + diffusion_model = m.get_model_object("diffusion_model") + torch._dynamo.config.cache_size_limit = dynamo_cache_size_limit + torch._dynamo.config.force_parameter_static_shapes = force_parameter_static_shapes + try: + if compile_transformer_blocks_only: + compile_key_list = [] + for i, block in enumerate(diffusion_model.blocks): + compile_key_list.append(f"diffusion_model.blocks.{i}") + else: + compile_key_list =["diffusion_model"] + + set_torch_compile_wrapper(model=m, keys=compile_key_list, backend=backend, mode=mode, dynamic=dynamic, fullgraph=fullgraph) + except Exception as e: + raise RuntimeError("Failed to compile model") from e + + return (m, ) + + +_aimdo_patched = False + +def patch_aimdo_for_compile(): + # reduce recompiles with dynamic VRAM + global _aimdo_patched + if _aimdo_patched: + return + _aimdo_patched = True + names = ("cast_bias_weight", "uncast_bias_weight", "cast_modules_with_vbar", "resolve_cast_module_with_vbar") + for name in names: + fn = getattr(comfy.ops, name, None) + if fn is not None: + setattr(comfy.ops, name, torch._dynamo.disable(fn)) + try: + import comfy_aimdo.torch as _at + _at.get_tensor_from_raw_ptr = torch._dynamo.disable(_at.get_tensor_from_raw_ptr) + except Exception: + pass + logging.info("KJNodes dynamic-compile: comfy.ops weight cast marked as eager graph break (recompile fix active).") + + +def skip_torch_compile_dict(guard_entries): + # don't recompile when transformer_options change + return [("transformer_options" not in entry.name) for entry in guard_entries] + + +def build_compile_kwargs(backend, mode, fullgraph, dynamic, use_guard_filter=True): + # torch.compile forbids passing mode and options together; an explicit mode wins, + # otherwise attach the guard filter via options on the default mode. + kw = {"backend": backend, "fullgraph": fullgraph, "dynamic": dynamic} + if mode and mode != "default": + kw["mode"] = mode + elif use_guard_filter: + kw["options"] = {"guard_filter_fn": skip_torch_compile_dict} + return kw + +import weakref as _kj_weakref + +_KJ_COMPILE_KEY = "torch.compile" +_KJ_COMPILED_BY_MODEL = _kj_weakref.WeakKeyDictionary() # BaseModel instance -> {key: compiled_module} + +# resolve compiled modules by the BaseModel actually executing +def _kj_apply_torch_compile_wrapper(executor, *args, **kwargs): + compiled = _KJ_COMPILED_BY_MODEL.get(executor.class_obj) + if not compiled: + return executor(*args, **kwargs) # this BaseModel wasn't compiled -> run eager, no swap + orig = {} + try: + for key, value in compiled.items(): + orig[key] = comfy.utils.get_attr(executor.class_obj, key) + comfy.utils.set_attr(executor.class_obj, key, value) + return executor(*args, **kwargs) + finally: + for key, value in orig.items(): + comfy.utils.set_attr(executor.class_obj, key, value) + + +def kj_set_torch_compile_wrapper(model, backend, options=None, mode=None, fullgraph=False, dynamic=None, keys=("diffusion_model",)): + WrappersMP = comfy.patcher_extension.WrappersMP + model.remove_wrappers_with_key(WrappersMP.APPLY_MODEL, _KJ_COMPILE_KEY) + if not keys: + keys = ["diffusion_model"] + compile_kwargs = {"backend": backend, "options": options, "mode": mode, "fullgraph": fullgraph, "dynamic": dynamic} + compiled_modules = {key: torch.compile(model=model.get_model_object(key), **compile_kwargs) for key in keys} + _KJ_COMPILED_BY_MODEL[model.model] = compiled_modules # register by the BaseModel that will execute + model.add_wrapper_with_key(WrappersMP.APPLY_MODEL, _KJ_COMPILE_KEY, _kj_apply_torch_compile_wrapper) + model.model_options["torch_compile_kwargs"] = compile_kwargs + + +class TorchCompileModelAdvanced: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "backend": (["inductor","cudagraphs"], {"default": "inductor"}), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + "dynamic": ( + ["auto", "true", "false"], + {"default": "auto", "tooltip": "Use dynamic shape tracing."}, + ), + "compile_transformer_blocks_only": ("BOOLEAN", {"default": True, "tooltip": "Compile only transformer blocks, faster compile and less error prone"}), + "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), + "debug_compile_keys": ("BOOLEAN", {"default": False, "tooltip": "Print the compile keys used for torch.compile"}), + }, + "optional": { + "disable_dynamic_vram": ("BOOLEAN", {"default": False, "tooltip": "Disable dynamic VRAM feature as it can cause issues with compile"}), + } + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "KJNodes/torchcompile" + DESCRIPTION = "Advanced torch.compile patching for diffusion models." + EXPERIMENTAL = True + + def patch(self, model, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_transformer_blocks_only, debug_compile_keys, disable_dynamic_vram=False): + if disable_dynamic_vram: + try: + m = model.clone(disable_dynamic=True) + except TypeError: + logging.warning("This ComfyUI version do not support disabling dynamic VRAM through a node. This may cause issues with torch.compile.") + m = model.clone() + else: + m = model.clone() + + diffusion_model = m.get_model_object("diffusion_model") + torch._dynamo.config.cache_size_limit = dynamo_cache_size_limit + + try: + compile_key_list = [] + if compile_transformer_blocks_only: + layer_types = ["double_blocks", "single_blocks", "layers", "transformer_blocks", "blocks", "visual_transformer_blocks", "text_transformer_blocks", "patch_blocks", "pixel_blocks"] + for layer_name in layer_types: + if hasattr(diffusion_model, layer_name): + blocks = getattr(diffusion_model, layer_name) + for i in range(len(blocks)): + compile_key_list.append(f"diffusion_model.{layer_name}.{i}") + if not compile_key_list: + logging.warning("No known transformer blocks found to compile, compiling entire diffusion model instead") + elif debug_compile_keys: + logging.info("TorchCompileModelAdvanced: Compile key list:") + for key in compile_key_list: + logging.info(f" - {key}") + if not compile_key_list: + compile_key_list =["diffusion_model"] + + dynamic_kv = {"true": True, "false": False, "auto": None} + try: + dynamic = dynamic_kv[dynamic] + except KeyError: + raise ValueError(f"Invalid dynamic arg {dynamic}") + + if not disable_dynamic_vram and getattr(m, "is_dynamic", lambda: False)(): + patch_aimdo_for_compile() # reduce recompiles with dynamic VRAM, will break the graph still but better than nothing + kj_set_torch_compile_wrapper(model=m, keys=compile_key_list, **build_compile_kwargs(backend, mode, fullgraph, dynamic)) + except Exception as e: + raise RuntimeError("Failed to compile model") from e + + return (m, ) + + +class TorchCompileModelQwenImage: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "backend": (["inductor","cudagraphs"], {"default": "inductor"}), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + "dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}), + "compile_transformer_blocks_only": ("BOOLEAN", {"default": True, "tooltip": "Compile only transformer blocks, faster compile and less error prone"}), + "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), + }, + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "KJNodes/torchcompile" + EXPERIMENTAL = True + DEPRECATED = True + DESCRIPTION = "Deprecated, use TorchCompileModelAdvanced instead." + + def patch(self, model, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_transformer_blocks_only): + from comfy_api.torch_helpers import set_torch_compile_wrapper + m = model.clone() + diffusion_model = m.get_model_object("diffusion_model") + torch._dynamo.config.cache_size_limit = dynamo_cache_size_limit + try: + if compile_transformer_blocks_only: + compile_key_list = [] + for i, block in enumerate(diffusion_model.transformer_blocks): + compile_key_list.append(f"diffusion_model.transformer_blocks.{i}") + else: + compile_key_list =["diffusion_model"] + + set_torch_compile_wrapper(model=m, keys=compile_key_list, backend=backend, mode=mode, dynamic=dynamic, fullgraph=fullgraph) + except Exception as e: + raise RuntimeError("Failed to compile model") from e + + return (m, ) + +class TorchCompileVAE: + def __init__(self): + self._compiled_encoder = False + self._compiled_decoder = False + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "vae": ("VAE",), + "backend": (["inductor", "cudagraphs"],), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + "compile_encoder": ("BOOLEAN", {"default": True, "tooltip": "Compile encoder"}), + "compile_decoder": ("BOOLEAN", {"default": True, "tooltip": "Compile decoder"}), + }} + RETURN_TYPES = ("VAE",) + FUNCTION = "compile" + + CATEGORY = "KJNodes/torchcompile" + EXPERIMENTAL = True + + def compile(self, vae, backend, mode, fullgraph, compile_encoder, compile_decoder): + if compile_encoder: + if not self._compiled_encoder: + encoder_name = "encoder" + if hasattr(vae.first_stage_model, "taesd_encoder"): + encoder_name = "taesd_encoder" + + try: + setattr( + vae.first_stage_model, + encoder_name, + torch.compile( + getattr(vae.first_stage_model, encoder_name), + mode=mode, + fullgraph=fullgraph, + backend=backend, + ), + ) + self._compiled_encoder = True + except Exception as e: + raise RuntimeError("Failed to compile model") from e + if compile_decoder: + if not self._compiled_decoder: + decoder_name = "decoder" + if hasattr(vae.first_stage_model, "taesd_decoder"): + decoder_name = "taesd_decoder" + + try: + setattr( + vae.first_stage_model, + decoder_name, + torch.compile( + getattr(vae.first_stage_model, decoder_name), + mode=mode, + fullgraph=fullgraph, + backend=backend, + ), + ) + self._compiled_decoder = True + except Exception as e: + raise RuntimeError("Failed to compile model") from e + return (vae, ) + +class TorchCompileControlNet: + def __init__(self): + self._compiled= False + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "controlnet": ("CONTROL_NET",), + "backend": (["inductor", "cudagraphs"],), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + }} + RETURN_TYPES = ("CONTROL_NET",) + FUNCTION = "compile" + + CATEGORY = "KJNodes/torchcompile" + EXPERIMENTAL = True + + def compile(self, controlnet, backend, mode, fullgraph): + if not self._compiled: + try: + # for i, block in enumerate(controlnet.control_model.double_blocks): + # print("Compiling controlnet double_block", i) + # controlnet.control_model.double_blocks[i] = torch.compile(block, mode=mode, fullgraph=fullgraph, backend=backend) + controlnet.control_model = torch.compile(controlnet.control_model, mode=mode, fullgraph=fullgraph, backend=backend) + self._compiled = True + except Exception as e: + self._compiled = False + raise RuntimeError("Failed to compile model") from e + + return (controlnet, ) + +class WanVideoTeaCacheKJ: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "rel_l1_thresh": ("FLOAT", {"default": 0.275, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Threshold for to determine when to apply the cache, compromise between speed and accuracy. When using coefficients a good value range is something between 0.2-0.4 for all but 1.3B model, which should be about 10 times smaller, same as when not using coefficients."}), + "start_percent": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The start percentage of the steps to use with TeaCache."}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The end percentage of the steps to use with TeaCache."}), + "cache_device": (["main_device", "offload_device"], {"default": "offload_device", "tooltip": "Device to cache to"}), + "coefficients": (["disabled", "1.3B", "14B", "i2v_480", "i2v_720"], {"default": "i2v_480", "tooltip": "Coefficients for rescaling the relative l1 distance, if disabled the threshold value should be about 10 times smaller than the value used with coefficients."}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "patch_teacache" + CATEGORY = "KJNodes/deprecated" + DEPRECATED = True + DESCRIPTION = """DEPRECATED, use the native EasyCache or alternative custom node that's up to date instead of this.""" + EXPERIMENTAL = True + + def patch_teacache(self, model, rel_l1_thresh, start_percent, end_percent, cache_device, coefficients): + return model, + + +from comfy.ldm.flux.math import apply_rope + +def modified_wan_self_attention_forward(self, x, freqs, transformer_options={}): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n * d) + return q, k, v + + q, k, v = qkv_fn(x) + q, k = apply_rope(q, k, freqs) + feta_scores = get_feta_scores(q, k, self.num_frames, self.enhance_weight) + + x = comfy.ldm.modules.attention.optimized_attention( + q.view(b, s, n * d), + k.view(b, s, n * d), + v, + heads=self.num_heads, + transformer_options=transformer_options, + ) + + x = self.o(x) + + x *= feta_scores + + return x + +from einops import rearrange +def get_feta_scores(query, key, num_frames, enhance_weight, num_heads=12): + img_q, img_k = query, key #torch.Size([2, 9216, 12, 128]) + + if img_q.ndim == 4: + B, ST, num_heads, head_dim = img_q.shape + elif img_q.ndim == 3: + B, ST, hidden_dim = img_q.shape + head_dim = hidden_dim // num_heads + + # Reshape from [B, ST, hidden_dim] to [B, ST, num_heads, head_dim] + img_q = img_q.view(B, ST, num_heads, head_dim) + img_k = img_k.view(B, ST, num_heads, head_dim) + + spatial_dim = ST // num_frames + + query_image = rearrange( + img_q, "B (T S) N C -> (B S) N T C", T=num_frames, S=spatial_dim, N=num_heads, C=head_dim + ) + key_image = rearrange( + img_k, "B (T S) N C -> (B S) N T C", T=num_frames, S=spatial_dim, N=num_heads, C=head_dim + ) + + return feta_score(query_image, key_image, head_dim, num_frames, enhance_weight) + +def feta_score(query_image, key_image, head_dim, num_frames, enhance_weight): + scale = head_dim**-0.5 + query_image = query_image * scale + attn_temp = query_image @ key_image.transpose(-2, -1) # translate attn to float32 + attn_temp = attn_temp.to(torch.float32) + attn_temp = attn_temp.softmax(dim=-1) + + # Reshape to [batch_size * num_tokens, num_frames, num_frames] + attn_temp = attn_temp.reshape(-1, num_frames, num_frames) + + # Create a mask for diagonal elements + diag_mask = torch.eye(num_frames, device=attn_temp.device).bool() + diag_mask = diag_mask.unsqueeze(0).expand(attn_temp.shape[0], -1, -1) + + # Zero out diagonal elements + attn_wo_diag = attn_temp.masked_fill(diag_mask, 0) + + # Calculate mean for each token's attention matrix + # Number of off-diagonal elements per matrix is n*n - n + num_off_diag = num_frames * num_frames - num_frames + mean_scores = attn_wo_diag.sum(dim=(1, 2)) / num_off_diag + + enhance_scores = mean_scores.mean() * (num_frames + enhance_weight) + enhance_scores = enhance_scores.clamp(min=1) + return enhance_scores + +import types +class WanAttentionPatch: + def __init__(self, num_frames, weight): + self.num_frames = num_frames + self.enhance_weight = weight + + def __get__(self, obj, objtype=None): + # Create bound method with stored parameters + def wrapped_attention(self_module, *args, **kwargs): + self_module.num_frames = self.num_frames + self_module.enhance_weight = self.enhance_weight + return modified_wan_self_attention_forward(self_module, *args, **kwargs) + return types.MethodType(wrapped_attention, obj) + +class WanVideoEnhanceAVideoKJ: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT", {"tooltip": "Only used to get the latent count"}), + "weight": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Strength of the enhance effect"}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "enhance" + CATEGORY = "KJNodes/wan" + DESCRIPTION = "https://github.com/NUS-HPC-AI-Lab/Enhance-A-Video" + EXPERIMENTAL = True + + def enhance(self, model, weight, latent): + if weight == 0: + return (model,) + + num_frames = latent["samples"].shape[2] + + model_clone = model.clone() + if 'transformer_options' not in model_clone.model_options: + model_clone.model_options['transformer_options'] = {} + model_clone.model_options["transformer_options"]["enhance_weight"] = weight + diffusion_model = model_clone.get_model_object("diffusion_model") + + compile_settings = getattr(model.model, "compile_settings", None) + for idx, block in enumerate(diffusion_model.blocks): + patched_attn = WanAttentionPatch(num_frames, weight).__get__(block.self_attn, block.__class__) + if compile_settings is not None: + patched_attn = torch.compile(patched_attn, mode=compile_settings["mode"], dynamic=compile_settings["dynamic"], fullgraph=compile_settings["fullgraph"], backend=compile_settings["backend"]) + + model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.self_attn.forward", patched_attn) + + return (model_clone,) + +try: + from comfy.ldm.lightricks.model import apply_rotary_emb +except ImportError: + apply_rotary_emb = None + +try: + from comfy.ldm.lightricks.model import GuideAttentionMask as _GuideAttentionMask, _attention_with_guide_mask as _ltx_attn_with_guide_mask +except ImportError: + _GuideAttentionMask = None + _ltx_attn_with_guide_mask = None + + +def ltxv_feta_forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}): + q = self.to_q(x) + context = x if context is None else context + k = self.to_k(context) + v = self.to_v(context) + + q = self.q_norm(q) + k = self.k_norm(k) + + 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) + + feta_scores = get_feta_scores(q, k, self.num_frames, self.enhance_weight, self.heads) + + 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) + elif _GuideAttentionMask is not None and isinstance(mask, _GuideAttentionMask): + out = _ltx_attn_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options) + else: + out = comfy.ldm.modules.attention.optimized_attention_masked(q, k, v, self.heads, 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, H) + b, t, _ = out.shape + out = out.view(b, t, self.heads, self.dim_head) + gates = 2.0 * torch.sigmoid(gate_logits) # zero-init -> identity + out = out * gates.unsqueeze(-1) + out = out.view(b, t, self.heads * self.dim_head) + + return self.to_out(out) * feta_scores + + +class LTXCrossAttentionPatch: + def __init__(self, num_frames, weight): + self.num_frames = num_frames + self.enhance_weight = weight + + def __get__(self, obj, objtype=None): + # Create bound method with stored parameters + def wrapped_attention(self_module, *args, **kwargs): + self_module.num_frames = self.num_frames + self_module.enhance_weight = self.enhance_weight + return ltxv_feta_forward(self_module, *args, **kwargs) + return types.MethodType(wrapped_attention, obj) + +class LTXVEnhanceAVideoKJ: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT", {"tooltip": "Only used to get the latent count"}), + "weight": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 100.0, "step": 0.001, "tooltip": "Strength of the enhance effect"}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "enhance" + CATEGORY = "KJNodes/ltxv" + DESCRIPTION = "https://github.com/NUS-HPC-AI-Lab/Enhance-A-Video" + EXPERIMENTAL = True + + def enhance(self, model, weight, latent): + if weight == 0: + return (model,) + + num_frames = latent["samples"].shape[2] + + model_clone = model.clone() + if 'transformer_options' not in model_clone.model_options: + model_clone.model_options['transformer_options'] = {} + model_clone.model_options["transformer_options"]["enhance_weight"] = weight + diffusion_model = model_clone.get_model_object("diffusion_model") + + for idx, block in enumerate(diffusion_model.transformer_blocks): + patched_attn1 = LTXCrossAttentionPatch(num_frames, weight).__get__(block.attn1, block.__class__) + model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.attn1.forward", patched_attn1) + return (model_clone,) + +def _wan_compute_attention(self, query, context, transformer_options={}): + k = self.norm_k(self.k(context)) + v = self.v(context) + return comfy.ldm.modules.attention.optimized_attention(query, k, v, heads=self.num_heads, transformer_options=transformer_options).flatten(2) + +def wan_nag_attention(self, query, context_positive, nag_context, transformer_options={}): + x_positive = _wan_compute_attention(self, query, context_positive, transformer_options) + x_negative = _wan_compute_attention(self, query, nag_context, transformer_options) + return x_positive, x_negative + +def normalized_attention_guidance(self, x_positive, x_negative): + if self.inplace: + nag_guidance = x_negative.mul_(self.nag_scale - 1).neg_().add_(x_positive, alpha=self.nag_scale) + del x_negative + else: + nag_guidance = x_negative * (self.nag_scale - 1) + del x_negative + nag_guidance = (x_positive * self.nag_scale).sub_(nag_guidance) + + norm_positive = torch.norm(x_positive, p=1, dim=-1, keepdim=True) + norm_guidance = torch.norm(nag_guidance, p=1, dim=-1, keepdim=True) + + scale = norm_guidance / norm_positive + torch.nan_to_num_(scale, nan=10.0) + mask = scale > self.nag_tau + del scale + + adjustment = (norm_positive * self.nag_tau) / (norm_guidance + 1e-7) + del norm_positive, norm_guidance + + nag_guidance.mul_(torch.where(mask, adjustment, 1.0)) + del mask, adjustment + + if self.inplace: + return nag_guidance.sub_(x_positive).mul_(self.nag_alpha).add_(x_positive) + else: + nag_guidance.mul_(self.nag_alpha) + return nag_guidance.add_(x_positive * (1 - self.nag_alpha)) + +#region NAG +def wan_crossattn_forward_nag(self, x, context, transformer_options={}, **kwargs): + # Determine batch splitting and context handling + if self.input_type == "default": + # Single or [pos, neg] pair + if context.shape[0] == 1: + x_pos, context_pos = x, context + x_neg, context_neg = None, None + else: + x_pos, x_neg = torch.chunk(x, 2, dim=0) + context_pos, context_neg = torch.chunk(context, 2, dim=0) + elif self.input_type == "batch": + # Standard batch, no CFG + x_pos, context_pos = x, context + x_neg, context_neg = None, None + + # Positive branch + q_pos = self.norm_q(self.q(x_pos)) + nag_context = self.nag_context + if self.input_type == "batch": + nag_context = nag_context.repeat(x_pos.shape[0], 1, 1) + del x_pos + + x_positive, x_negative = wan_nag_attention(self, q_pos, context_pos, nag_context, transformer_options=transformer_options) + del context_pos, q_pos + + x_pos_out = normalized_attention_guidance(self, x_positive, x_negative) + del x_positive, x_negative + + # Negative branch + if x_neg is not None and context_neg is not None: + q_neg = self.norm_q(self.q(x_neg)) + k_neg = self.norm_k(self.k(context_neg)) + v_neg = self.v(context_neg) + x_neg_out = comfy.ldm.modules.attention.optimized_attention(q_neg, k_neg, v_neg, heads=self.num_heads, transformer_options=transformer_options) + x = torch.cat([x_pos_out, x_neg_out], dim=0) + else: + x = x_pos_out + + return self.o(x) + +def wan_i2v_crossattn_forward_nag(self, x, context, context_img_len, transformer_options={}, **kwargs): + context_img = context[:, :context_img_len] + context = context[:, context_img_len:] + + q_img = 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_img, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options) + del q_img, k_img, v_img, context_img + + if context.shape[0] == 2: + x, x_real_negative = torch.chunk(x, 2, dim=0) + context_positive, context_negative = torch.chunk(context, 2, dim=0) + else: + context_positive = context + context_negative = None + + q = self.norm_q(self.q(x)) + + x_positive, x_negative = wan_nag_attention(self, q, context_positive, self.nag_context, transformer_options=transformer_options) + del q, context_positive + x = normalized_attention_guidance(self, x_positive, x_negative) + del x_positive, x_negative + + if context_negative is not None: + q_real_negative = self.norm_q(self.q(x_real_negative)) + k_real_negative = self.norm_k(self.k(context_negative)) + v_real_negative = self.v(context_negative) + x_real_negative = comfy.ldm.modules.attention.optimized_attention(q_real_negative, k_real_negative, v_real_negative, heads=self.num_heads, transformer_options=transformer_options) + x = torch.cat([x, x_real_negative], dim=0) + + return self.o(x + img_x) + + +class WanCrossAttentionPatch: + def __init__(self, context, nag_scale, nag_alpha, nag_tau, i2v=False, input_type="default", inplace=True): + self.nag_context = context + self.nag_scale = nag_scale + self.nag_alpha = nag_alpha + self.nag_tau = nag_tau + self.i2v = i2v + self.input_type = input_type + self.inplace = inplace + def __get__(self, obj, objtype=None): + # Create bound method with stored parameters + def wrapped_attention(self_module, *args, **kwargs): + self_module.nag_context = self.nag_context + self_module.nag_scale = self.nag_scale + self_module.nag_alpha = self.nag_alpha + self_module.nag_tau = self.nag_tau + self_module.input_type = self.input_type + self_module.inplace = self.inplace + if self.i2v: + return wan_i2v_crossattn_forward_nag(self_module, *args, **kwargs) + else: + return wan_crossattn_forward_nag(self_module, *args, **kwargs) + return types.MethodType(wrapped_attention, obj) + + +class WanVideoNAG: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "conditioning": ("CONDITIONING",), + "nag_scale": ("FLOAT", {"default": 11.0, "min": 0.0, "max": 100.0, "step": 0.001, "tooltip": "Strength of negative guidance effect"}), + "nag_alpha": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "Mixing coefficient in that controls the balance between the normalized guided representation and the original positive representation."}), + "nag_tau": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Clipping threshold that controls how much the guided attention can deviate from the positive attention."}), + }, + "optional": { + "input_type": (["default", "batch"], {"tooltip": "Type of the model input"}), + "inplace": ("BOOLEAN", {"default": False, "tooltip": "If true, modifies tensors in place to save memory. Leads to different numerical results which may change the output slightly."}), + }, + + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "patch" + CATEGORY = "KJNodes/wan" + DESCRIPTION = "https://github.com/ChenDarYen/Normalized-Attention-Guidance" + EXPERIMENTAL = True + + def patch(self, model, conditioning, nag_scale, nag_alpha, nag_tau, input_type="default", inplace=False): + if nag_scale == 0: + return (model,) + + device = mm.get_torch_device() + dtype = mm.unet_dtype() + + model_clone = model.clone() + + diffusion_model = model_clone.get_model_object("diffusion_model") + + diffusion_model.text_embedding.to(device) + context = diffusion_model.text_embedding(conditioning[0][0].to(device, dtype)) + + type_str = str(type(model.model.model_config).__name__) + i2v = True if "WAN21_I2V" in type_str else False + + for idx, block in enumerate(diffusion_model.blocks): + patched_attn = WanCrossAttentionPatch(context, nag_scale, nag_alpha, nag_tau, i2v, input_type=input_type, inplace=inplace).__get__(block.cross_attn, block.__class__) + + model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.cross_attn.forward", patched_attn) + + return (model_clone,) + +class SkipLayerGuidanceWanVideo: + @classmethod + def INPUT_TYPES(s): + return {"required": {"model": ("MODEL", ), + "blocks": ("STRING", {"default": "10", "multiline": False}), + "start_percent": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "slg" + EXPERIMENTAL = True + DESCRIPTION = "Simplified skip layer guidance that only skips the uncond on selected blocks" + DEPRECATED = True + CATEGORY = "advanced/guidance" + + def slg(self, model, start_percent, end_percent, blocks): + def skip(args, extra_args): + transformer_options = extra_args.get("transformer_options", {}) + original_block = extra_args["original_block"] + + if not transformer_options: + raise ValueError("transformer_options not found in extra_args, currently SkipLayerGuidanceWanVideo only works with TeaCacheKJ") + if start_percent <= transformer_options["current_percent"] <= end_percent: + if args["img"].shape[0] == 2: + prev_img_uncond = args["img"][0].unsqueeze(0) + + new_args = { + "img": args["img"][1].unsqueeze(0), + "txt": args["txt"][1].unsqueeze(0), + "vec": args["vec"][1].unsqueeze(0), + "pe": args["pe"][1].unsqueeze(0) + } + + block_out = original_block(new_args) + + out = { + "img": torch.cat([prev_img_uncond, block_out["img"]], dim=0), + "txt": args["txt"], + "vec": args["vec"], + "pe": args["pe"] + } + else: + if transformer_options.get("cond_or_uncond") == [0]: + out = original_block(args) + else: + out = args + else: + out = original_block(args) + return out + + block_list = [int(x.strip()) for x in blocks.split(",")] + blocks = [int(i) for i in block_list] + logging.info(f"Selected blocks to skip uncond on: {blocks}") + + m = model.clone() + + for b in blocks: + #m.set_model_patch_replace(skip, "dit", "double_block", b) + model_options = m.model_options["transformer_options"].copy() + if "patches_replace" not in model_options: + model_options["patches_replace"] = {} + else: + model_options["patches_replace"] = model_options["patches_replace"].copy() + + if "dit" not in model_options["patches_replace"]: + model_options["patches_replace"]["dit"] = {} + else: + model_options["patches_replace"]["dit"] = model_options["patches_replace"]["dit"].copy() + + block = ("double_block", b) + + model_options["patches_replace"]["dit"][block] = skip + m.model_options["transformer_options"] = model_options + + + return (m, ) + +class CFGZeroStarAndInit: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "use_zero_init": ("BOOLEAN", {"default": True}), + "zero_init_steps": ("INT", {"default": 0, "min": 0, "tooltip": "for zero init, starts from 0 so first step is always zeroed out if use_zero_init enabled"}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + DESCRIPTION = "https://github.com/WeichenFan/CFG-Zero-star" + CATEGORY = "KJNodes/experimental" + EXPERIMENTAL = True + + def patch(self, model, use_zero_init, zero_init_steps): + def cfg_zerostar(args): + #zero init + cond = args["cond"] + timestep = args["timestep"] + sigmas = args["model_options"]["transformer_options"]["sample_sigmas"] + matched_step_index = (sigmas == timestep[0]).nonzero() + if len(matched_step_index) > 0: + current_step_index = matched_step_index.item() + else: + for i in range(len(sigmas) - 1): + if (sigmas[i] - timestep[0]) * (sigmas[i + 1] - timestep[0]) <= 0: + current_step_index = i + break + else: + current_step_index = 0 + + if (current_step_index <= zero_init_steps) and use_zero_init: + return cond * 0 + + uncond = args["uncond"] + cond_scale = args["cond_scale"] + + batch_size = cond.shape[0] + + positive_flat = cond.view(batch_size, -1) + negative_flat = uncond.view(batch_size, -1) + + dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) + squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 + alpha = dot_product / squared_norm + alpha = alpha.view(batch_size, *([1] * (len(cond.shape) - 1))) + + noise_pred = uncond * alpha + cond_scale * (cond - uncond * alpha) + return noise_pred + + m = model.clone() + m.set_model_sampler_cfg_function(cfg_zerostar) + return (m, ) + +class GGUFLoaderKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + # Get GGUF models safely, fallback to empty list if unet_gguf folder doesn't exist + try: + gguf_models = folder_paths.get_filename_list("unet_gguf") + ltx2_connector_models = folder_paths.get_filename_list("text_encoders") + ltx2_connector_models = [m for m in ltx2_connector_models if "connector" in m.lower()] + except KeyError: + gguf_models = [] + ltx2_connector_models = [] + + return io.Schema( + node_id="GGUFLoaderKJ", + category="KJNodes/model_loaders", + description="Loads a GGUF model with advanced options, requires [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) to be installed.", + is_experimental=True, + inputs=[ + io.Combo.Input("model_name", options=gguf_models), + io.Combo.Input("extra_model_name", options=gguf_models + ltx2_connector_models + ["none"], default="none", tooltip="An extra gguf model to load and merge into the main model, for example VACE module"), + io.Combo.Input("dequant_dtype", options=["default", "target", "float32", "float16", "bfloat16"], default="default"), + io.Combo.Input("patch_dtype", options=["default", "target", "float32", "float16", "bfloat16"], default="default"), + io.Boolean.Input("patch_on_device", default=False), + io.Boolean.Input("enable_fp16_accumulation", default=False, tooltip="Enable torch.backends.cuda.matmul.allow_fp16_accumulation, required minimum pytorch version 2.7.1"), + io.Combo.Input("attention_override", options=["none", "sdpa", "sageattn", "xformers", "flashattn"], default="none", tooltip="Overrides the used attention implementation, requires the respective library to be installed"), + + ], + outputs=[io.Model.Output(),], + ) + + def attention_override_pytorch(func, *args, **kwargs): + new_attention = comfy.ldm.modules.attention.attention_pytorch + return new_attention.__wrapped__(*args, **kwargs) + def attention_override_sage(func, *args, **kwargs): + new_attention = comfy.ldm.modules.attention.attention_sage + return new_attention.__wrapped__(*args, **kwargs) + def attention_override_xformers(func, *args, **kwargs): + new_attention = comfy.ldm.modules.attention.attention_xformers + return new_attention.__wrapped__(*args, **kwargs) + def attention_override_flash(func, *args, **kwargs): + new_attention = comfy.ldm.modules.attention.attention_flash + return new_attention.__wrapped__(*args, **kwargs) + + ATTENTION_OVERRIDES = { + "sdpa": attention_override_pytorch, + "sageattn": attention_override_sage, + "xformers": attention_override_xformers, + "flashattn": attention_override_flash, + } + + + @classmethod + def _get_gguf_module(cls): + """Import GGUF module with version validation""" + for key, mod in sys.modules.items(): + if key.endswith("ComfyUI-GGUF") or key.endswith("comfyui-gguf"): + if hasattr(mod, "ops") and hasattr(mod, "nodes"): + return mod + + gguf_path = os.path.join(folder_paths.folder_names_and_paths["custom_nodes"][0][0], "ComfyUI-GGUF") + for module_name in ["ComfyUI-GGUF", "custom_nodes.ComfyUI-GGUF", "comfyui-gguf", "custom_nodes.comfyui-gguf", gguf_path, gguf_path.lower()]: + try: + module = importlib.import_module(module_name) + return module + except ImportError: + continue + + raise ImportError( + "Compatible ComfyUI-GGUF not found. " + "Please install/update from: https://github.com/city96/ComfyUI-GGUF" + ) + + @classmethod + def execute(cls, model_name, extra_model_name, dequant_dtype, patch_dtype, patch_on_device, attention_override, enable_fp16_accumulation): + gguf_nodes = cls._get_gguf_module() + ops = gguf_nodes.ops.GGMLOps() + + def set_linear_dtype(attr, value): + if value == "default": + setattr(ops.Linear, attr, None) + elif value == "target": + setattr(ops.Linear, attr, value) + else: + setattr(ops.Linear, attr, getattr(torch, value)) + + set_linear_dtype("dequant_dtype", dequant_dtype) + set_linear_dtype("patch_dtype", patch_dtype) + + # init model + extra = {} + model_path = folder_paths.get_full_path("unet", model_name) + try: + sd, extra = gguf_nodes.loader.gguf_sd_loader(model_path) + except TypeError: + sd = gguf_nodes.loader.gguf_sd_loader(model_path) + + if extra_model_name is not None and extra_model_name != "none": + if extra_model_name.endswith(".gguf"): + extra_model_full_path = folder_paths.get_full_path("unet", extra_model_name) + try: + extra_model, _ = gguf_nodes.loader.gguf_sd_loader(extra_model_full_path) + except TypeError: + extra_model = gguf_nodes.loader.gguf_sd_loader(extra_model_full_path) + elif "connector" in extra_model_name.lower(): + extra_model_full_path = folder_paths.get_full_path("text_encoders", extra_model_name) + extra_model = comfy.utils.load_torch_file(extra_model_full_path) + diffusion_model_prefix = comfy.model_detection.unet_prefix_from_state_dict(extra_model) + if diffusion_model_prefix == "model.diffusion_model.": + temp_sd = comfy.utils.state_dict_prefix_replace(extra_model, {diffusion_model_prefix: ""}, filter_keys=True) + if len(temp_sd) > 0: + extra_model = temp_sd + else: + raise ValueError("Extra model must also be a .gguf file") + sd.update(extra_model) + + model = comfy.sd.load_diffusion_model_state_dict( + sd, model_options={"custom_operations": ops}, metadata=extra.get("metadata", {}) + ) + if model is None: + raise RuntimeError(f"ERROR: Could not detect model type of: {model_path}") + + model = gguf_nodes.nodes.GGUFModelPatcher.clone(model) + model.patch_on_device = patch_on_device + + # attention override + if attention_override in cls.ATTENTION_OVERRIDES: + model.model_options["transformer_options"]["optimized_attention_override"] = cls.ATTENTION_OVERRIDES[attention_override] + + if enable_fp16_accumulation: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + torch.backends.cuda.matmul.allow_fp16_accumulation = True + else: + raise RuntimeError("Failed to set fp16 accumulation, requires pytorch version 2.7.1 or higher") + else: + if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): + torch.backends.cuda.matmul.allow_fp16_accumulation = False + + return io.NodeOutput(model,) + +try: + from torch.nn.attention.flex_attention import flex_attention, BlockMask +except ImportError: + flex_attention = None + BlockMask = None + +class NABLA_AttentionKJ(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "latent": ("LATENT", {"tooltip": "Only used to get the latent shape"}), + "window_time": ("INT", {"default": 11, "min": 1, "tooltip": "Temporal attention window size"}), + "window_width": ("INT", {"default": 3, "min": 1, "tooltip": "Spatial attention window size"}), + "window_height": ("INT", {"default": 3, "min": 1, "tooltip": "Spatial attention window size"}), + "sparsity": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}), + "torch_compile": ("BOOLEAN", {"default": True, "tooltip": "Most likely required for reasonable memory usage"}) + }, + } + + RETURN_TYPES = ("MODEL", ) + FUNCTION = "patch" + DESCRIPTION = "Experimental node for patching attention mode to use NABLA sparse attention for video models, currently only works with Kadinsky5" + CATEGORY = "KJNodes/experimental" + + def patch(self, model, latent, window_time, window_width, window_height, sparsity, torch_compile): + if flex_attention is None or BlockMask is None: + raise RuntimeError("can't import flex_attention from torch.nn.attention, requires newer pytorch version") + + model_clone = model.clone() + samples = latent["samples"] + + sparse_params = get_sparse_params(samples, window_time, window_height, window_width, sparsity) + nabla_attention = NABLA_Attention(sparse_params) + + def attention_override_nabla(func, *args, **kwargs): + return nabla_attention(*args, **kwargs) + + if torch_compile: + attention_override_nabla = torch.compile(attention_override_nabla, mode="max-autotune-no-cudagraphs", dynamic=True) + + # attention override + model_clone.model_options["transformer_options"]["optimized_attention_override"] = attention_override_nabla + + return model_clone, + + +class NABLA_Attention(): + def __init__(self, sparse_params): + self.sparse_params = sparse_params + + def __call__(self, q, k, v, heads, **kwargs): + if q.shape[-2] < 3000 or k.shape[-2] < 3000: + return optimized_attention(q, k, v, heads, **kwargs) + block_mask = self.nablaT_v2(q, k, self.sparse_params["sta_mask"], thr=self.sparse_params["P"]) + out = flex_attention(q, k, v, block_mask=block_mask).transpose(1, 2).contiguous().flatten(-2, -1) + return out + + def nablaT_v2(self, q, k, sta, thr=0.9): + # Map estimation + BLOCK_SIZE = 64 + B, h, S, D = q.shape + s1 = S // BLOCK_SIZE + qa = q.reshape(B, h, s1, BLOCK_SIZE, D).mean(-2) + ka = k.reshape(B, h, s1, BLOCK_SIZE, D).mean(-2).transpose(-2, -1) + map = qa @ ka + + map = torch.softmax(map / math.sqrt(D), dim=-1) + # Map binarization + vals, inds = map.sort(-1) + cvals = vals.cumsum_(-1) + mask = (cvals >= 1 - thr).int() + mask = mask.gather(-1, inds.argsort(-1)) + + mask = torch.logical_or(mask, sta) + + # BlockMask creation + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=BLOCK_SIZE, mask_mod=None) + +def fast_sta_nabla(T, H, W, wT=3, wH=3, wW=3): + l = torch.Tensor([T, H, W]).amax() + r = torch.arange(0, l, 1, dtype=torch.int16, device=mm.get_torch_device()) + mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() + sta_t, sta_h, sta_w = ( + mat[:T, :T].flatten(), + mat[:H, :H].flatten(), + mat[:W, :W].flatten(), + ) + sta_t = sta_t <= wT // 2 + sta_h = sta_h <= wH // 2 + sta_w = sta_w <= wW // 2 + sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() + sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) + return sta.reshape(T * H * W, T * H * W) + + +def get_sparse_params(x, wT, wH, wW, sparsity=0.9): + B, C, T, H, W = x.shape + #print("x shape:", x.shape) + patch_size = (1, 2, 2) + T, H, W = ( + T // patch_size[0], + H // patch_size[1], + W // patch_size[2], + ) + sta_mask = fast_sta_nabla(T, H // 8, W // 8, wT, wH, wW) + sparse_params = { + "sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0), + "to_fractal": True, + "P": sparsity, + "wT": wT, + "wH": wH, + "wW": wW, + "add_sta": True, + "visual_shape": (T, H, W), + "method": "topcdf", + } + + return sparse_params + +from comfy.comfy_types.node_typing import IO +class StartRecordCUDAMemoryHistory(): + # @classmethod + # def IS_CHANGED(s): + # return True + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input": (IO.ANY,), + "enabled": (["all", "state", "None"], {"default": "all", "tooltip": "None: disable, 'state': keep info for allocated memory, 'all': keep history of all alloc/free calls"}), + "context": (["all", "state", "alloc", "None"], {"default": "all", "tooltip": "None: no tracebacks, 'state': tracebacks for allocated memory, 'alloc': for alloc calls, 'all': for free calls"}), + "stacks": (["python", "all"], {"default": "all", "tooltip": "'python': Python/TorchScript/inductor frames, 'all': also C++ frames"}), + "max_entries": ("INT", {"default": 100000, "min": 1000, "max": 10000000, "tooltip": "Maximum number of entries to record"}), + }, + } + + RETURN_TYPES = (IO.ANY, ) + RETURN_NAMES = ("input", "output_path",) + FUNCTION = "start" + CATEGORY = "KJNodes/memory" + DESCRIPTION = "THIS NODE ALWAYS RUNS. Starts recording CUDA memory allocation history, can be ended and saved with EndRecordCUDAMemoryHistory. " + + def start(self, input, enabled, context, stacks, max_entries): + mm.soft_empty_cache() + torch.cuda.reset_peak_memory_stats(mm.get_torch_device()) + torch.cuda.memory._record_memory_history( + max_entries=max_entries, + enabled=enabled if enabled != "None" else None, + context=context if context != "None" else None, + stacks=stacks + ) + return input, + +class EndRecordCUDAMemoryHistory(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "input": (IO.ANY,), + "output_path": ("STRING", {"default": "comfy_cuda_memory_history"}, "Base path for saving the CUDA memory history file, timestamp and .pt extension will be added"), + }, + } + + RETURN_TYPES = (IO.ANY, "STRING",) + RETURN_NAMES = ("input", "output_path",) + FUNCTION = "end" + CATEGORY = "KJNodes/memory" + DESCRIPTION = "Records CUDA memory allocation history between start and end, saves to a file that can be analyzed here: https://docs.pytorch.org/memory_viz or with VisualizeCUDAMemoryHistory node" + + def end(self, input, output_path): + mm.soft_empty_cache() + time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + output_path = f"{output_path}{time}.pt" + torch.cuda.memory._dump_snapshot(output_path) + torch.cuda.memory._record_memory_history(enabled=None) + return input, output_path + +try: + from server import PromptServer +except ImportError: + PromptServer = None + +class VisualizeCUDAMemoryHistory(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "snapshot_path": ("STRING", ), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("output_path",) + FUNCTION = "visualize" + CATEGORY = "KJNodes/memory" + DESCRIPTION = "Visualizes a CUDA memory allocation history file, opens in browser" + OUTPUT_NODE = True + + def visualize(self, snapshot_path, unique_id): + import pickle + from torch.cuda import _memory_viz + import uuid + + from folder_paths import get_output_directory + output_dir = get_output_directory() + + with open(snapshot_path, "rb") as f: + snapshot = pickle.load(f) + + html = _memory_viz.trace_plot(snapshot) + html_filename = f"cuda_memory_history_{uuid.uuid4().hex}.html" + output_path = os.path.join(output_dir, "memory_history", html_filename) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + + with open(output_path, "w", encoding="utf-8") as f: + f.write(html) + + api_url = f"http://localhost:8188/api/view?type=output&filename={html_filename}&subfolder=memory_history" + + # Progress UI + if unique_id and PromptServer is not None: + try: + PromptServer.instance.send_progress_text( + api_url, + unique_id + ) + except: + pass + + return api_url, + + +class ModelMemoryUseReportPatch: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + }} + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + DESCRIPTION = "Adds callbacks to model to report memory usage during after sampling" + EXPERIMENTAL = True + CATEGORY = "KJNodes/memory" + + def patch(self, model): + model_clone = model.clone() + device = mm.get_torch_device() + + def reset_mem_usage(model): + torch.cuda.reset_peak_memory_stats(device) + def report_mem_usage(model): + max_memory = torch.cuda.max_memory_allocated(device) / 1024**3 + max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 + logging.info(f"Sampling max allocated memory: {max_memory=:.3f} GB") + logging.info(f"Sampling max reserved memory: {max_reserved=:.3f} GB") + + model_clone.add_callback(CallbacksMP.ON_PRE_RUN, reset_mem_usage) + model_clone.add_callback(CallbacksMP.ON_CLEANUP, report_mem_usage) + + return (model_clone,) + + +class MemoryUsageFactorAdjustWrapper: + def __init__(self, memory_usage_factor, original_factor): + self.memory_usage_factor = memory_usage_factor + self.original_factor = original_factor + + def __call__(self, executor, model, noise_shape: torch.Tensor, *args, **kwargs): + m = model.clone() + m.model.memory_usage_factor = self.memory_usage_factor + logging.info(f"Temporarily set memory usage factor to {self.memory_usage_factor}") + try: + result = executor(m, noise_shape, *args, **kwargs) + finally: + logging.info(f"Model memory usage calculated, restoring original memory usage factor: {self.original_factor}") + m.model.memory_usage_factor = self.original_factor + return result + +class ModelMemoryUsageFactorOverride: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "memory_usage_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.001}), + }} + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + DESCRIPTION = "Overrides the memory usage factor of the model during sampling." + EXPERIMENTAL = True + CATEGORY = "KJNodes/memory" + + def patch(self, model, memory_usage_factor): + model_clone = model.clone() + original_memory_usage_factor = model_clone.model.memory_usage_factor + logging.info(f"Original memory usage factor: {original_memory_usage_factor}") + + wrapper = MemoryUsageFactorAdjustWrapper(memory_usage_factor, original_memory_usage_factor) + model_clone.add_wrapper_with_key( + comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING, + "memory_usage_factor_adjust_prepare_sampling", + wrapper + ) + return (model_clone,) + +def wan_ffn_chunked_forward(self, x): + if x.shape[1] > self.dim_threshold: + chunks = torch.chunk(x, self.num_chunks, dim=1) + output_chunks = [] + for chunk in chunks: + output_chunks.append(torch.nn.Sequential.forward(self, chunk)) + chunked = torch.cat(output_chunks, dim=1) + return chunked + else: + return torch.nn.Sequential.forward(self, x) + +class WanffnChunkPatch: + def __init__(self, num_chunks, dim_threshold=4096): + self.num_chunks = num_chunks + self.dim_threshold = dim_threshold + + def __get__(self, obj, objtype=None): + def wrapped_forward(self_module, *args, **kwargs): + self_module.num_chunks = self.num_chunks + self_module.dim_threshold = self.dim_threshold + return wan_ffn_chunked_forward(self_module, *args, **kwargs) + return types.MethodType(wrapped_forward, obj) + +class WanChunkFeedForward(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WanChunkFeedForward", + display_name="Wan Chunk FeedForward", + category="KJNodes/wan", + description="EXPERIMENTAL AND MAY CHANGE THE MODEL OUTPUT!! Chunks feedforward activations to reduce peak VRAM usage.", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + io.Int.Input("chunks", default=2, min=1, max=100, step=1, tooltip="Number of chunks to split the feedforward activations into to reduce peak VRAM usage."), + io.Int.Input("dim_threshold", default=4096, min=1024, max=16384, step=256, tooltip="Dimension threshold above which to apply chunking."), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, chunks, dim_threshold) -> io.NodeOutput: + if chunks == 1: + return io.NodeOutput(model) + + model_clone = model.clone() + diffusion_model = model_clone.get_model_object("diffusion_model") + + for idx, block in enumerate(diffusion_model.blocks): + patched_ffn = WanffnChunkPatch(chunks, dim_threshold).__get__(block.ffn, block.__class__) + model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.ffn.forward", patched_ffn) + + return io.NodeOutput(model_clone) + + +# Ideogram4 peak-VRAM patches (FFN sequence chunking + bf16 RoPE) +from comfy.ldm.lumina.model import FeedForward as _Ideogram4FeedForward +from comfy.ldm.modules.attention import optimized_attention_masked as _ideogram4_attn + + +def ideogram4_ffn_chunked_forward(self, x): + # x: (B, L, dim). Chunk over the token dim so the (B, L, hidden) SwiGLU + if x.shape[1] > self.kj_dim_threshold and self.kj_num_chunks > 1: + out = [_Ideogram4FeedForward.forward(self, c) for c in torch.chunk(x, self.kj_num_chunks, dim=1)] + return torch.cat(out, dim=1) + return _Ideogram4FeedForward.forward(self, x) + + +class Ideogram4FFNChunkPatch: + def __init__(self, num_chunks, dim_threshold): + self.num_chunks = num_chunks + self.dim_threshold = dim_threshold + + def __get__(self, obj, objtype=None): + def wrapped_forward(self_module, *args, **kwargs): + self_module.kj_num_chunks = self.num_chunks + self_module.kj_dim_threshold = self.dim_threshold + return ideogram4_ffn_chunked_forward(self_module, *args, **kwargs) + return types.MethodType(wrapped_forward, obj) + + +def _ideogram4_apply_rope_lowp(xq, xk, freqs_cis): + # (bf16/fp16) instead of being upcast to fp32 -> ~halves RoPE activation memory. + cos = freqs_cis[0].to(xq.dtype) + sin = freqs_cis[1].to(xq.dtype) + nsin = freqs_cis[2].to(xq.dtype) + + q_embed = xq * cos + qs = q_embed.shape[-1] // 2 + q_embed[..., :qs].addcmul_(xq[..., qs:], nsin) + q_embed[..., qs:].addcmul_(xq[..., :qs], sin) + + k_embed = xk * cos + ks = k_embed.shape[-1] // 2 + k_embed[..., :ks].addcmul_(xk[..., ks:], nsin) + k_embed[..., ks:].addcmul_(xk[..., :ks], sin) + return q_embed, k_embed + + +def ideogram4_attention_lowp_rope_forward(self, x, attn_mask, freqs_cis, transformer_options={}): + batch_size, seq_len, _ = x.shape + q, k, v = self.qkv(x).view(batch_size, seq_len, 3, self.num_heads, self.head_dim).unbind(dim=2) + q = self.norm_q(q).transpose(1, 2) + k = self.norm_k(k).transpose(1, 2) + v = v.transpose(1, 2) + q, k = _ideogram4_apply_rope_lowp(q, k, freqs_cis) + out = _ideogram4_attn(q, k, v, self.num_heads, attn_mask, skip_reshape=True, transformer_options=transformer_options) + return self.o(out) + + +class Ideogram4RopePatch: + def __get__(self, obj, objtype=None): + return types.MethodType(ideogram4_attention_lowp_rope_forward, obj) + + +class Ideogram4OptimizationsKJ(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="Ideogram4OptimizationsKJ", + display_name="Ideogram4 Optimizations KJ", + category="KJNodes/experimental", + description="EXPERIMENTAL AND MAY CHANGE THE MODEL OUTPUT!! Reduces peak VRAM of the Ideogram4 forward. " + "chunk_ffn splits the SwiGLU activations over the token dim; bf16_rope applies RoPE in the model " + "dtype instead of upcasting to fp32. Both target the two largest transient tensors in the block.", + is_experimental=True, + inputs=[ + io.Model.Input("model"), + io.Boolean.Input("chunk_ffn", default=True, + tooltip="Chunk the feedforward activations over the sequence dim to cap the (B, L, hidden) intermediate."), + io.Int.Input("ffn_chunks", default=2, min=1, max=64, step=1, + tooltip="Number of chunks to split the feedforward sequence into. More chunks = lower peak, slightly more overhead."), + io.Int.Input("ffn_seq_threshold", default=1024, min=256, max=65536, step=256, + tooltip="Only chunk when the token sequence length exceeds this (skips chunking for tiny sequences)."), + io.Boolean.Input("bf16_rope", default=True, + tooltip="Apply RoPE in the input dtype instead of fp32. ~Halves RoPE activation memory; matches the HF reference dtype."), + ], + outputs=[ + io.Model.Output(display_name="model"), + ], + ) + + @classmethod + def execute(cls, model, chunk_ffn, ffn_chunks, ffn_seq_threshold, bf16_rope) -> io.NodeOutput: + if not chunk_ffn and not bf16_rope: + return io.NodeOutput(model) + + m = model.clone() + diffusion_model = m.get_model_object("diffusion_model") + + layers = getattr(diffusion_model, "layers", None) + if not layers or not hasattr(layers[0], "feed_forward") or not hasattr(layers[0], "attention"): + logging.warning("Ideogram4OptimizationsKJ: model does not look like Ideogram4 " + "(expected diffusion_model.layers[*].feed_forward/.attention); returning model unchanged.") + return io.NodeOutput(model) + + for idx, block in enumerate(layers): + if chunk_ffn and ffn_chunks > 1: + patched_ffn = Ideogram4FFNChunkPatch(ffn_chunks, ffn_seq_threshold).__get__(block.feed_forward, block.feed_forward.__class__) + m.add_object_patch(f"diffusion_model.layers.{idx}.feed_forward.forward", patched_ffn) + if bf16_rope: + patched_attn = Ideogram4RopePatch().__get__(block.attention, block.attention.__class__) + m.add_object_patch(f"diffusion_model.layers.{idx}.attention.forward", patched_attn) + + return io.NodeOutput(m) + + +from comfy.samplers import KSAMPLER +from comfy.k_diffusion.sampling import to_d + +def sample_selfrefinevideo(model, x, sigmas, stochastic_step_map, certain_percentage=0.999, uncertainty_threshold=0.25, extra_args=None, callback=None, disable=None, verbose=False, video_shape=None, seed=None): + extra_args = {} if extra_args is None else extra_args + sigma_in = x.new_ones([x.shape[0]]) + + if seed is not None: + generator = torch.Generator(torch.device("cpu")).manual_seed(seed) + + pbar = tqdm(total=len(sigmas) - 1, disable=disable, desc="Sampling") + + for i in range(len(sigmas) - 1): + + # Get stochastic steps for this noise level + current_num_anneal_steps = stochastic_step_map.get(i, 0) + use_stochastic = current_num_anneal_steps > 0 + m = current_num_anneal_steps + 1 if use_stochastic else 1 + + sigma, sigma_next = sigmas[i], sigmas[i + 1] + + prev_certain_mask = None + prev_denoised = None + prev_denoised_full = None + prev_x_next = None + prev_x_next_video = None + is_certain = False + + for ii in range(m): + if m > 1: + pbar.set_description(f"Step {i}/{len(sigmas)-1} (substep {ii+1}/{m})") + # Early exit if certain threshold reached + if is_certain: + x = prev_x_next + break + + # Determine input + noise = torch.randn(x.shape, device=torch.device("cpu"), generator=generator).to(x) + x_in = x if ii == 0 else (1.0 - sigma) * prev_denoised_full + sigma * noise + if ii > 0: + x = x_in + + denoised = model(x_in, sigmas[i] * sigma_in, **extra_args) + + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + + # Compute next latents + d = to_d(x, sigma, denoised) + x_next = x + (sigma_next - sigma) * d + + # Separate video and audio if joint model + if d.ndim == 3 and video_shape is not None: + cut = math.prod(video_shape[1:]) + denoised_video = denoised[:, :, :cut].reshape([denoised.shape[0]] + list(video_shape)[1:]) + x_next_video = x_next[:, :, :cut].reshape([denoised.shape[0]] + list(video_shape)[1:]) + denoised_audio = denoised[:, :, cut:] + x_next_audio = x_next[:, :, cut:] + if verbose: + tqdm.write(f"Video shape: {denoised_video.shape}, Audio shape: {denoised_audio.shape}") + else: + denoised_video = denoised + x_next_video = x_next + denoised_audio = None + x_next_audio = None + + # Stochastic sampling with uncertainty masking + if use_stochastic and prev_denoised is not None: + # Compute uncertainty and masking on video part + diff = denoised_video - prev_denoised + uncertainty = torch.sqrt(torch.sum(diff ** 2, dim=1)) / denoised_video.shape[1] + certain_mask = uncertainty < uncertainty_threshold + + if verbose: + tqdm.write(f"Step {i}/{len(sigmas)-1} substep {ii+1}/{m}:") + tqdm.write(f"Uncertainty: min {uncertainty.min():.4f}, max {uncertainty.max():.4f}, threshold {uncertainty_threshold}") + tqdm.write(f"Certain pixels: {certain_mask.sum()}/{certain_mask.numel()} = {certain_mask.sum()/certain_mask.numel():.4f}") + + # Update certain mask (union with previous) + if prev_certain_mask is not None: + certain_mask = certain_mask | prev_certain_mask + + # Check certainty threshold + if certain_mask.sum() / certain_mask.numel() > certain_percentage: + is_certain = True + if verbose: + tqdm.write(f"{ii}/{current_num_anneal_steps}: Certain region is more than {certain_percentage}, we are certain") + + # Apply masking to video + certain_mask_float = certain_mask.float().unsqueeze(1) + x_next_video = certain_mask_float * prev_x_next_video + (1.0 - certain_mask_float) * x_next_video + denoised_video = certain_mask_float * prev_denoised + (1.0 - certain_mask_float) * denoised_video + + # Reconstruct full latents by replacing the video portion + if x_next_audio is not None: + # Flatten masked video back to match original format and replace video portion + x_next = x_next.clone() + x_next[:, :, :cut] = x_next_video.reshape([x_next_video.shape[0], x_next.shape[1], -1]) + # Also reconstruct full denoised for next iteration input + denoised_full = denoised.clone() + denoised_full[:, :, :cut] = denoised_video.reshape([denoised_video.shape[0], denoised.shape[1], -1]) + else: + # No audio separation + x_next = x_next_video + denoised_full = denoised_video + + prev_certain_mask = certain_mask + prev_denoised = denoised_video + prev_denoised_full = denoised_full + prev_x_next_video = x_next_video + prev_x_next = x_next + elif use_stochastic: + # For first stochastic step, create denoised_full if we have audio + if x_next_audio is not None: + denoised_full = denoised.clone() + denoised_full[:, :, :cut] = denoised_video.reshape([denoised_video.shape[0], denoised.shape[1], -1]) + else: + denoised_full = denoised_video + + prev_certain_mask = None + prev_denoised = denoised_video + prev_denoised_full = denoised_full + prev_x_next_video = x_next_video + prev_x_next = x_next + + # Update x for final step + if use_stochastic and ii == m - 1: + x = prev_x_next + elif not use_stochastic: + x = x_next + + pbar.update(1) + if m == 1: + pbar.set_description("Sampling") + pbar.close() + return x + +class SamplerSelfRefineVideo(io.ComfyNode): + @classmethod + def define_schema(cls): + default_ranges = [ + (2, 5, 3), # Range 1 + (6, 14, 1), # Range 2 + ] + + options = [] + + # Option 1: 2 ranges + range_inputs_2 = [] + for i in range(1, 3): + start_default, end_default, steps_default = default_ranges[i - 1] + range_inputs_2.extend([ + io.Int.Input(f"start_step{i}", default=start_default, min=0, max=999, step=1, tooltip=f"Start step for range {i}"), + io.Int.Input(f"end_step{i}", default=end_default, min=0, max=999, step=1, tooltip=f"End step for range {i}"), + io.Int.Input(f"steps_{i}", default=steps_default, min=1, max=100, step=1, tooltip=f"Number of P&P steps for range {i}"), + ]) + options.append(io.DynamicCombo.Option(key="2 ranges", inputs=range_inputs_2)) + + # Option 2: 1 range + range_inputs_1 = [] + for i in range(1, 2): + start_default, end_default, steps_default = default_ranges[i - 1] + range_inputs_1.extend([ + io.Int.Input(f"start_step{i}", default=start_default, min=0, max=999, step=1, tooltip=f"Start step for range {i}"), + io.Int.Input(f"end_step{i}", default=end_default, min=0, max=999, step=1, tooltip=f"End step for range {i}"), + io.Int.Input(f"steps_{i}", default=steps_default, min=1, max=100, step=1, tooltip=f"Number of P&P steps for range {i}"), + ]) + options.append(io.DynamicCombo.Option(key="1 range", inputs=range_inputs_1)) + + # Option 3: Manual string input + options.append(io.DynamicCombo.Option( + key="from_string", + inputs=[ + io.String.Input( + "stochastic_plan", + default="2-5:3,6-14:1", + multiline=True, + tooltip="Format: 'start-end:steps,start-end:steps' e.g. '2-5:3,6-14:1'" + ) + ] + )) + return io.Schema( + node_id="SamplerSelfRefineVideo", + category="KJNodes/samplers", + description="Attempt to implement https://github.com/agwmon/self-refine-video, for testing only, MAY NOT WORK AS INTENDED.", + is_experimental=True, + inputs=[ + io.DynamicCombo.Input("input_mode", options=options, tooltip="How to configure the step plan"), + io.Float.Input("certain_percentage", default=0.999, min=0.0, max=1.0, step=0.001, round=False, tooltip="Percentage of certain pixels to consider the frame as certain and skip further refinement"), + io.Float.Input("uncertainty_threshold", default=0.2, min=0.0, max=1.0, step=0.01, round=False, tooltip="Threshold of uncertainty to consider a pixel uncertain"), + io.Boolean.Input("verbose", default=False, tooltip="Enable verbose logging during sampling"), + io.Latent.Input("latent", optional=True, tooltip="Optional latent input to get input shape for LTX2 audio/video separation"), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, step=1, tooltip="Seed for stochastic sampling"), + ], + outputs=[io.Sampler.Output()] + ) + + @classmethod + def execute(cls, input_mode, certain_percentage, uncertainty_threshold, seed, verbose, latent=None) -> io.NodeOutput: + video_shape = None + if latent is not None: + video_shape = latent["samples"].shape + + range_keys = sorted([k for k in input_mode.keys() if k.startswith('start_step')]) + stochastic_step_map = {} + if "stochastic_plan" in input_mode: + # Parse manual string format: "2-5:3,6-14:1" + plan_str = input_mode["stochastic_plan"] + ranges = plan_str.split(",") + for range_spec in ranges: + range_spec = range_spec.strip() + if not range_spec: + continue + try: + range_part, steps_part = range_spec.split(":") + start, end = range_part.split("-") + start, end, steps = int(start), int(end), int(steps_part) + for idx in range(start, end + 1): + stochastic_step_map[idx] = steps + except ValueError: + raise ValueError(f"Invalid format in stochastic_plan: '{range_spec}'. Expected format: 'start-end:steps'") + else: + range_keys = [k for k in input_mode.keys() if k.startswith('start_step')] + for start_key in range_keys: + i = start_key.replace('start_step', '') + start = input_mode.get(f"start_step{i}") + end = input_mode.get(f"end_step{i}") + steps = input_mode.get(f"steps_{i}") + + if start is not None and end is not None and steps is not None: + for idx in range(start, end + 1): + stochastic_step_map[idx] = steps + + sampler = KSAMPLER(sample_selfrefinevideo, { + "stochastic_step_map": stochastic_step_map, + "certain_percentage": certain_percentage, + "uncertainty_threshold": uncertainty_threshold, + "verbose": verbose, + "video_shape": video_shape, + "seed": seed, + }) + return io.NodeOutput(sampler) + + +# Multi-feature linear bias corrector for PiD (Flux2 backbone, 4-step). +# Calibrated on 124 natural-image samples (LOO-CV per-channel RMSE 0.027/0.026/0.024). +# Held-out validation on 20 unseen images: 60% reduction in total drift vs uncorrected. +# Features per row: [R_mean, G_mean, B_mean, R_std, G_std, B_std, +# R_mean*G_mean, R_mean*B_mean, G_mean*B_mean, intercept(1.0)] +# Columns: predicted bias for R, G, B (subtract from x0_pred at step 0). +PID_BIAS_COEF_FLUX2 = torch.tensor([ + [-0.130306, +0.127184, +0.014058], # R_mean + [-0.053279, -0.408929, +0.004243], # G_mean + [-0.009386, +0.109546, -0.134091], # B_mean + [-0.033373, -0.011615, -0.026129], # R_std + [+0.180052, +0.062021, +0.071317], # G_std + [-0.067958, -0.058595, -0.098645], # B_std + [-0.248116, -0.240633, -0.105600], # R_mean*G_mean + [+0.304035, +0.322566, +0.093224], # R_mean*B_mean + [-0.157648, -0.227127, -0.112368], # G_mean*B_mean + [-0.062814, +0.030765, +0.062735], # intercept +], dtype=torch.float32) + + +class PiDColorBiasCorrection: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "strength": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01, + "tooltip": "Correction strength. 1.0 = full predicted bias subtracted. <1 = milder, >1 = stronger, 0 = disabled."}), + "backbone": (["flux2"], {"default": "flux2", + "tooltip": "Calibrated PiD backbone (currently only flux2 — others use the same model but coefficients differ)."}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "KJNodes/experimental" + EXPERIMENTAL = True + DESCRIPTION = ( + "PiD 4-step decoder color/brightness drift corrector. " + "Subtracts a per-channel bias from x0_pred at the first sampling step, " + "using a small linear model calibrated against the model's systematic drift " + "(model tends to brighten dark scenes and add a blue cast)." + ) + + def patch(self, model, strength, backbone): + if strength == 0.0 or backbone != "flux2": + return (model,) + coef_cpu = PID_BIAS_COEF_FLUX2 # (10, 3) + + def pid_bias_post_cfg(args): + denoised = args["denoised"] + # Step detection: only apply at the first sampling step. + # Use sample_sigmas like CFGZeroStarAndInit for robustness across schedules. + try: + sigmas = args["model_options"]["transformer_options"]["sample_sigmas"] + sigma = args.get("sigma", args.get("timestep")) + # First step matches sigmas[0] + if sigma is None or not torch.isclose(sigma.max(), sigmas[0]).item(): + return denoised + except (KeyError, AttributeError): + # Fallback heuristic: PiD's first step has sigma=0.999 + sigma = args.get("sigma") + if sigma is None or sigma.max().item() < 0.95: + return denoised + + coef = coef_cpu.to(denoised.device, dtype=denoised.dtype) + rgb_m = denoised.mean(dim=(0, 2, 3)) + rgb_s = denoised.std(dim=(0, 2, 3)) + one = torch.tensor(1.0, device=denoised.device, dtype=denoised.dtype) + feats = torch.stack([ + rgb_m[0], rgb_m[1], rgb_m[2], + rgb_s[0], rgb_s[1], rgb_s[2], + rgb_m[0] * rgb_m[1], rgb_m[0] * rgb_m[2], rgb_m[1] * rgb_m[2], + one, + ]) + bias = feats @ coef # (3,) + return denoised - strength * bias.view(1, 3, 1, 1) + + m = model.clone() + m.set_model_sampler_post_cfg_function(pid_bias_post_cfg) + return (m,) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..5eebc21da4069b8e2cf0cc2713e8fcde5d8b2b70 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/nodes.py @@ -0,0 +1,3357 @@ +import torch +import torch.nn as nn +import numpy as np +from PIL import Image +import json +import re +import os +import time +import math +import importlib +import logging + +from comfy import model_management +import folder_paths +from nodes import MAX_RESOLUTION +from comfy.utils import common_upscale, ProgressBar, load_torch_file, save_torch_file, state_dict_prefix_replace +from comfy.comfy_types.node_typing import IO +from comfy_api.latest import io, ui +import comfy.latent_formats +import node_helpers +from io import BytesIO + +script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +folder_paths.add_model_folder_path("kjnodes_fonts", os.path.join(script_directory, "fonts")) + +class BOOLConstant: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "value": ("BOOLEAN", {"default": True}), + }, + } + RETURN_TYPES = ("BOOLEAN",) + RETURN_NAMES = ("value",) + FUNCTION = "get_value" + CATEGORY = "KJNodes/constants" + SEARCH_ALIASES = ["boolean", "value"] + + def get_value(self, value): + return (value,) + +class INTConstant: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "value": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff}), + }, + } + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("value",) + FUNCTION = "get_value" + CATEGORY = "KJNodes/constants" + SEARCH_ALIASES = ["integer", "value"] + + def get_value(self, value): + return (value,) + +class FloatConstant: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "value": ("FLOAT", {"default": 0.0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 0.00001}), + }, + } + + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("value",) + FUNCTION = "get_value" + CATEGORY = "KJNodes/constants" + SEARCH_ALIASES = ["float", "value"] + + def get_value(self, value): + return (round(value, 6),) + +class StringConstant: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "string": ("STRING", {"default": '', "multiline": False}), + } + } + RETURN_TYPES = ("STRING",) + FUNCTION = "passtring" + CATEGORY = "KJNodes/constants" + SEARCH_ALIASES = ["text", "value"] + + def passtring(self, string): + return (string, ) + +class StringConstantMultiline: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "string": ("STRING", {"default": "", "multiline": True}), + "strip_newlines": ("BOOLEAN", {"default": True}), + } + } + RETURN_TYPES = ("STRING",) + FUNCTION = "stringify" + CATEGORY = "KJNodes/constants" + SEARCH_ALIASES = ["text", "value"] + + def stringify(self, string, strip_newlines): + new_string = string + if strip_newlines: + new_string = new_string.replace('\n', '').strip() + return (new_string,) + + + +class ScaleBatchPromptSchedule: + + RETURN_TYPES = ("STRING",) + FUNCTION = "scaleschedule" + CATEGORY = "KJNodes/misc" + DESCRIPTION = """ +Scales a batch schedule from Fizz' nodes BatchPromptSchedule +to a different frame count. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_str": ("STRING", {"forceInput": True,"default": "0:(0.0),\n7:(1.0),\n15:(0.0)\n"}), + "old_frame_count": ("INT", {"forceInput": True,"default": 1,"min": 1, "max": 4096, "step": 1}), + "new_frame_count": ("INT", {"forceInput": True,"default": 1,"min": 1, "max": 4096, "step": 1}), + + }, + } + + def scaleschedule(self, old_frame_count, input_str, new_frame_count): + pattern = r'"(\d+)"\s*:\s*"(.*?)"(?:,|\Z)' + frame_strings = dict(re.findall(pattern, input_str)) + + # Calculate the scaling factor + scaling_factor = (new_frame_count - 1) / (old_frame_count - 1) + + # Initialize a dictionary to store the new frame numbers and strings + new_frame_strings = {} + + # Iterate over the frame numbers and strings + for old_frame, string in frame_strings.items(): + # Calculate the new frame number + new_frame = int(round(int(old_frame) * scaling_factor)) + + # Store the new frame number and corresponding string + new_frame_strings[new_frame] = string + + # Format the output string + output_str = ', '.join([f'"{k}":"{v}"' for k, v in sorted(new_frame_strings.items())]) + return (output_str,) + + +class GetLatentsFromBatchIndexed: + + RETURN_TYPES = ("LATENT",) + FUNCTION = "indexedlatentsfrombatch" + CATEGORY = "KJNodes/latents" + DESCRIPTION = """ +Selects and returns the latents at the specified indices as an latent batch. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "latents": ("LATENT",), + "indexes": ("STRING", {"default": "0, 1, 2", "multiline": True}), + "latent_format": (["BCHW", "BTCHW", "BCTHW"], {"default": "BCHW"}), + }, + } + + def indexedlatentsfrombatch(self, latents, indexes, latent_format): + + samples = latents.copy() + latent_samples = samples["samples"] + + # Parse the indexes string into a list of integers + index_list = [int(index.strip()) for index in indexes.split(',')] + + # Convert list of indices to a PyTorch tensor + indices_tensor = torch.tensor(index_list, dtype=torch.long) + + # Select the latents at the specified indices + if latent_format == "BCHW": + chosen_latents = latent_samples[indices_tensor] + elif latent_format == "BTCHW": + chosen_latents = latent_samples[:, indices_tensor] + elif latent_format == "BCTHW": + chosen_latents = latent_samples[:, :, indices_tensor] + + samples["samples"] = chosen_latents + return (samples,) + + +class ConditioningMultiCombine: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 20, "step": 1}), + "operation": (["combine", "concat"], {"default": "combine"}), + "conditioning_1": ("CONDITIONING", ), + "conditioning_2": ("CONDITIONING", ), + }, + } + + RETURN_TYPES = ("CONDITIONING", "INT") + RETURN_NAMES = ("combined", "inputcount") + FUNCTION = "combine" + CATEGORY = "KJNodes/masking/conditioning" + DESCRIPTION = """ +Combines multiple conditioning nodes into one +""" + + def combine(self, inputcount, operation, **kwargs): + from nodes import ConditioningCombine + from nodes import ConditioningConcat + cond_combine_node = ConditioningCombine() + cond_concat_node = ConditioningConcat() + cond = kwargs["conditioning_1"] + for c in range(1, inputcount): + new_cond = kwargs[f"conditioning_{c + 1}"] + if operation == "combine": + cond = cond_combine_node.combine(cond, new_cond)[0] + elif operation == "concat": + cond = cond_concat_node.concat(cond, new_cond)[0] + return (cond, inputcount,) + +class AppendStringsToList: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "string1": ("STRING", {"default": '', "forceInput": True}), + "string2": ("STRING", {"default": '', "forceInput": True}), + } + } + RETURN_TYPES = ("STRING",) + FUNCTION = "joinstring" + CATEGORY = "KJNodes/text" + + def joinstring(self, string1, string2): + if not isinstance(string1, list): + string1 = [string1] + if not isinstance(string2, list): + string2 = [string2] + + joined_string = string1 + string2 + return (joined_string, ) + +class JoinStrings: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "delimiter": ("STRING", {"default": ' ', "multiline": False}), + }, + "optional": { + "string1": ("STRING", {"default": '', "forceInput": True}), + "string2": ("STRING", {"default": '', "forceInput": True}), + } + } + RETURN_TYPES = ("STRING",) + FUNCTION = "joinstring" + CATEGORY = "KJNodes/text" + + def joinstring(self, delimiter, string1="", string2=""): + joined_string = string1 + delimiter + string2 + return (joined_string, ) + +class JoinStringMulti: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "string_1": ("STRING", {"default": '', "forceInput": True}), + "delimiter": ("STRING", {"default": ' ', "multiline": False}), + "return_list": ("BOOLEAN", {"default": False}), + }, + "optional": { + "string_2": ("STRING", {"default": '', "forceInput": True}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("string",) + FUNCTION = "combine" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +Creates single string, or a list of strings, from +multiple input strings. +You can set how many inputs the node has, +with the **inputcount** and clicking update. +""" + + def combine(self, inputcount, delimiter, **kwargs): + string = kwargs["string_1"] + return_list = kwargs["return_list"] + strings = [string] # Initialize a list with the first string + for c in range(1, inputcount): + new_string = kwargs.get(f"string_{c + 1}", "") + if not new_string: + continue + if return_list: + strings.append(new_string) # Add new string to the list + else: + string = string + delimiter + new_string + if return_list: + return (strings,) # Return the list of strings + else: + return (string,) # Return the combined string + +class CondPassThrough: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + }, + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING",) + RETURN_NAMES = ("positive", "negative") + FUNCTION = "passthrough" + CATEGORY = "KJNodes/misc" + DESCRIPTION = """ + Simply passes through the positive and negative conditioning, + workaround for Set node not allowing bypassed inputs. +""" + + def passthrough(self, positive=None, negative=None): + return (positive, negative,) + +class ModelPassThrough: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "model": ("MODEL", ), + }, + } + + RETURN_TYPES = ("MODEL", ) + RETURN_NAMES = ("model",) + FUNCTION = "passthrough" + CATEGORY = "KJNodes/misc" + DESCRIPTION = """ + Simply passes through the model, + workaround for Set node not allowing bypassed inputs. +""" + + def passthrough(self, model=None): + return (model,) + +def append_helper(t, mask, c, set_area_to_bounds, strength): + n = [t[0], t[1].copy()] + _, h, w = mask.shape + n[1]['mask'] = mask + n[1]['set_area_to_bounds'] = set_area_to_bounds + n[1]['mask_strength'] = strength + c.append(n) + +class ConditioningSetMaskAndCombine: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive_1": ("CONDITIONING", ), + "negative_1": ("CONDITIONING", ), + "positive_2": ("CONDITIONING", ), + "negative_2": ("CONDITIONING", ), + "mask_1": ("MASK", ), + "mask_2": ("MASK", ), + "mask_1_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "set_cond_area": (["default", "mask bounds"],), + } + } + + RETURN_TYPES = ("CONDITIONING","CONDITIONING",) + RETURN_NAMES = ("combined_positive", "combined_negative",) + FUNCTION = "append" + CATEGORY = "KJNodes/masking/conditioning" + DESCRIPTION = """ +Bundles multiple conditioning mask and combine nodes into one,functionality is identical to ComfyUI native nodes +""" + + def append(self, positive_1, negative_1, positive_2, negative_2, mask_1, mask_2, set_cond_area, mask_1_strength, mask_2_strength): + c = [] + c2 = [] + set_area_to_bounds = False + if set_cond_area != "default": + set_area_to_bounds = True + if len(mask_1.shape) < 3: + mask_1 = mask_1.unsqueeze(0) + if len(mask_2.shape) < 3: + mask_2 = mask_2.unsqueeze(0) + for t in positive_1: + append_helper(t, mask_1, c, set_area_to_bounds, mask_1_strength) + for t in positive_2: + append_helper(t, mask_2, c, set_area_to_bounds, mask_2_strength) + for t in negative_1: + append_helper(t, mask_1, c2, set_area_to_bounds, mask_1_strength) + for t in negative_2: + append_helper(t, mask_2, c2, set_area_to_bounds, mask_2_strength) + return (c, c2) + +class ConditioningSetMaskAndCombine3: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive_1": ("CONDITIONING", ), + "negative_1": ("CONDITIONING", ), + "positive_2": ("CONDITIONING", ), + "negative_2": ("CONDITIONING", ), + "positive_3": ("CONDITIONING", ), + "negative_3": ("CONDITIONING", ), + "mask_1": ("MASK", ), + "mask_2": ("MASK", ), + "mask_3": ("MASK", ), + "mask_1_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_3_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "set_cond_area": (["default", "mask bounds"],), + } + } + + RETURN_TYPES = ("CONDITIONING","CONDITIONING",) + RETURN_NAMES = ("combined_positive", "combined_negative",) + FUNCTION = "append" + CATEGORY = "KJNodes/masking/conditioning" + DESCRIPTION = """ +Bundles multiple conditioning mask and combine nodes into one,functionality is identical to ComfyUI native nodes +""" + + def append(self, positive_1, negative_1, positive_2, positive_3, negative_2, negative_3, mask_1, mask_2, mask_3, set_cond_area, mask_1_strength, mask_2_strength, mask_3_strength): + c = [] + c2 = [] + set_area_to_bounds = False + if set_cond_area != "default": + set_area_to_bounds = True + if len(mask_1.shape) < 3: + mask_1 = mask_1.unsqueeze(0) + if len(mask_2.shape) < 3: + mask_2 = mask_2.unsqueeze(0) + if len(mask_3.shape) < 3: + mask_3 = mask_3.unsqueeze(0) + for t in positive_1: + append_helper(t, mask_1, c, set_area_to_bounds, mask_1_strength) + for t in positive_2: + append_helper(t, mask_2, c, set_area_to_bounds, mask_2_strength) + for t in positive_3: + append_helper(t, mask_3, c, set_area_to_bounds, mask_3_strength) + for t in negative_1: + append_helper(t, mask_1, c2, set_area_to_bounds, mask_1_strength) + for t in negative_2: + append_helper(t, mask_2, c2, set_area_to_bounds, mask_2_strength) + for t in negative_3: + append_helper(t, mask_3, c2, set_area_to_bounds, mask_3_strength) + return (c, c2) + +class ConditioningSetMaskAndCombine4: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive_1": ("CONDITIONING", ), + "negative_1": ("CONDITIONING", ), + "positive_2": ("CONDITIONING", ), + "negative_2": ("CONDITIONING", ), + "positive_3": ("CONDITIONING", ), + "negative_3": ("CONDITIONING", ), + "positive_4": ("CONDITIONING", ), + "negative_4": ("CONDITIONING", ), + "mask_1": ("MASK", ), + "mask_2": ("MASK", ), + "mask_3": ("MASK", ), + "mask_4": ("MASK", ), + "mask_1_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_3_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_4_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "set_cond_area": (["default", "mask bounds"],), + } + } + + RETURN_TYPES = ("CONDITIONING","CONDITIONING",) + RETURN_NAMES = ("combined_positive", "combined_negative",) + FUNCTION = "append" + CATEGORY = "KJNodes/masking/conditioning" + DESCRIPTION = """ +Bundles multiple conditioning mask and combine nodes into one,functionality is identical to ComfyUI native nodes +""" + + def append(self, positive_1, negative_1, positive_2, positive_3, positive_4, negative_2, negative_3, negative_4, mask_1, mask_2, mask_3, mask_4, set_cond_area, mask_1_strength, mask_2_strength, mask_3_strength, mask_4_strength): + c = [] + c2 = [] + set_area_to_bounds = False + if set_cond_area != "default": + set_area_to_bounds = True + if len(mask_1.shape) < 3: + mask_1 = mask_1.unsqueeze(0) + if len(mask_2.shape) < 3: + mask_2 = mask_2.unsqueeze(0) + if len(mask_3.shape) < 3: + mask_3 = mask_3.unsqueeze(0) + if len(mask_4.shape) < 3: + mask_4 = mask_4.unsqueeze(0) + for t in positive_1: + append_helper(t, mask_1, c, set_area_to_bounds, mask_1_strength) + for t in positive_2: + append_helper(t, mask_2, c, set_area_to_bounds, mask_2_strength) + for t in positive_3: + append_helper(t, mask_3, c, set_area_to_bounds, mask_3_strength) + for t in positive_4: + append_helper(t, mask_4, c, set_area_to_bounds, mask_4_strength) + for t in negative_1: + append_helper(t, mask_1, c2, set_area_to_bounds, mask_1_strength) + for t in negative_2: + append_helper(t, mask_2, c2, set_area_to_bounds, mask_2_strength) + for t in negative_3: + append_helper(t, mask_3, c2, set_area_to_bounds, mask_3_strength) + for t in negative_4: + append_helper(t, mask_4, c2, set_area_to_bounds, mask_4_strength) + return (c, c2) + +class ConditioningSetMaskAndCombine5: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive_1": ("CONDITIONING", ), + "negative_1": ("CONDITIONING", ), + "positive_2": ("CONDITIONING", ), + "negative_2": ("CONDITIONING", ), + "positive_3": ("CONDITIONING", ), + "negative_3": ("CONDITIONING", ), + "positive_4": ("CONDITIONING", ), + "negative_4": ("CONDITIONING", ), + "positive_5": ("CONDITIONING", ), + "negative_5": ("CONDITIONING", ), + "mask_1": ("MASK", ), + "mask_2": ("MASK", ), + "mask_3": ("MASK", ), + "mask_4": ("MASK", ), + "mask_5": ("MASK", ), + "mask_1_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_3_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_4_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "mask_5_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "set_cond_area": (["default", "mask bounds"],), + } + } + + RETURN_TYPES = ("CONDITIONING","CONDITIONING",) + RETURN_NAMES = ("combined_positive", "combined_negative",) + FUNCTION = "append" + CATEGORY = "KJNodes/masking/conditioning" + DESCRIPTION = """ +Bundles multiple conditioning mask and combine nodes into one,functionality is identical to ComfyUI native nodes +""" + + def append(self, positive_1, negative_1, positive_2, positive_3, positive_4, positive_5, negative_2, negative_3, negative_4, negative_5, mask_1, mask_2, mask_3, mask_4, mask_5, set_cond_area, mask_1_strength, mask_2_strength, mask_3_strength, mask_4_strength, mask_5_strength): + c = [] + c2 = [] + set_area_to_bounds = False + if set_cond_area != "default": + set_area_to_bounds = True + if len(mask_1.shape) < 3: + mask_1 = mask_1.unsqueeze(0) + if len(mask_2.shape) < 3: + mask_2 = mask_2.unsqueeze(0) + if len(mask_3.shape) < 3: + mask_3 = mask_3.unsqueeze(0) + if len(mask_4.shape) < 3: + mask_4 = mask_4.unsqueeze(0) + if len(mask_5.shape) < 3: + mask_5 = mask_5.unsqueeze(0) + for t in positive_1: + append_helper(t, mask_1, c, set_area_to_bounds, mask_1_strength) + for t in positive_2: + append_helper(t, mask_2, c, set_area_to_bounds, mask_2_strength) + for t in positive_3: + append_helper(t, mask_3, c, set_area_to_bounds, mask_3_strength) + for t in positive_4: + append_helper(t, mask_4, c, set_area_to_bounds, mask_4_strength) + for t in positive_5: + append_helper(t, mask_5, c, set_area_to_bounds, mask_5_strength) + for t in negative_1: + append_helper(t, mask_1, c2, set_area_to_bounds, mask_1_strength) + for t in negative_2: + append_helper(t, mask_2, c2, set_area_to_bounds, mask_2_strength) + for t in negative_3: + append_helper(t, mask_3, c2, set_area_to_bounds, mask_3_strength) + for t in negative_4: + append_helper(t, mask_4, c2, set_area_to_bounds, mask_4_strength) + for t in negative_5: + append_helper(t, mask_5, c2, set_area_to_bounds, mask_5_strength) + return (c, c2) + +class VRAM_Debug: + + @classmethod + + def INPUT_TYPES(s): + return { + "required": { + + "empty_cache": ("BOOLEAN", {"default": True}), + "gc_collect": ("BOOLEAN", {"default": True}), + "unload_all_models": ("BOOLEAN", {"default": False}), + }, + "optional": { + "any_input": (IO.ANY,), + "image_pass": ("IMAGE",), + "model_pass": ("MODEL",), + } + } + + RETURN_TYPES = (IO.ANY, "IMAGE","MODEL","INT", "INT",) + RETURN_NAMES = ("any_output", "image_pass", "model_pass", "freemem_before", "freemem_after") + FUNCTION = "VRAMdebug" + CATEGORY = "KJNodes/memory" + DESCRIPTION = """ +Returns the inputs unchanged, they are only used as triggers, +and performs comfy model management functions and garbage collection, +reports free VRAM before and after the operations. +""" + + def VRAMdebug(self, gc_collect, empty_cache, unload_all_models, image_pass=None, model_pass=None, any_input=None): + freemem_before = model_management.get_free_memory() + logging.info(f"VRAMdebug: free memory before: {freemem_before:,.0f}") + if empty_cache: + model_management.soft_empty_cache() + if unload_all_models: + model_management.unload_all_models() + if gc_collect: + import gc + gc.collect() + freemem_after = model_management.get_free_memory() + logging.info(f"VRAMdebug: free memory after: {freemem_after:,.0f}") + logging.info(f"VRAMdebug: freed memory: {freemem_after - freemem_before:,.0f}") + return {"ui": { + "text": [f"{freemem_before:,.0f}x{freemem_after:,.0f}"]}, + "result": (any_input, image_pass, model_pass, freemem_before, freemem_after) + } + +class SomethingToString: + @classmethod + + def INPUT_TYPES(s): + return { + "required": { + "input": (IO.ANY, ), + }, + "optional": { + "prefix": ("STRING", {"default": ""}), + "suffix": ("STRING", {"default": ""}), + } + } + RETURN_TYPES = ("STRING",) + FUNCTION = "stringify" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +Converts any type to a string. +""" + + def stringify(self, input, prefix="", suffix=""): + if isinstance(input, (int, float, bool, str)): + stringified = str(input) + elif isinstance(input, list): + stringified = ', '.join(str(item) for item in input) + else: + return input, + if prefix: # Check if prefix is not empty + stringified = prefix + stringified # Add the prefix + if suffix: # Check if suffix is not empty + stringified = stringified + suffix # Add the suffix + + return (stringified,) + +class Sleep: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input": (IO.ANY, ), + "minutes": ("INT", {"default": 0, "min": 0, "max": 1439}), + "seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 59.99, "step": 0.01}), + }, + } + RETURN_TYPES = (IO.ANY,) + FUNCTION = "sleepdelay" + CATEGORY = "KJNodes/misc" + DESCRIPTION = """ +Delays the execution for the input amount of time. +""" + + def sleepdelay(self, input, minutes, seconds): + total_seconds = minutes * 60 + seconds + time.sleep(total_seconds) + return input, + +class EmptyLatentImagePresets: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "dimensions": ( + [ + '512 x 512 (1:1)', + '768 x 512 (1.5:1)', + '960 x 512 (1.875:1)', + '1024 x 512 (2:1)', + '1024 x 576 (1.778:1)', + '1536 x 640 (2.4:1)', + '1344 x 768 (1.75:1)', + '1216 x 832 (1.46:1)', + '1152 x 896 (1.286:1)', + '1024 x 1024 (1:1)', + ], + { + "default": '512 x 512 (1:1)' + }), + + "invert": ("BOOLEAN", {"default": False}), + "batch_size": ("INT", { + "default": 1, + "min": 1, + "max": 4096 + }), + }, + } + + RETURN_TYPES = ("LATENT", "INT", "INT") + RETURN_NAMES = ("Latent", "Width", "Height") + FUNCTION = "generate" + CATEGORY = "KJNodes/latents" + + def generate(self, dimensions, invert, batch_size): + from nodes import EmptyLatentImage + result = [x.strip() for x in dimensions.split('x')] + + # Remove the aspect ratio part + result[0] = result[0].split('(')[0].strip() + result[1] = result[1].split('(')[0].strip() + + if invert: + width = int(result[1].split(' ')[0]) + height = int(result[0]) + else: + width = int(result[0]) + height = int(result[1].split(' ')[0]) + latent = EmptyLatentImage().generate(width, height, batch_size)[0] + + return (latent, int(width), int(height),) + +class EmptyLatentImageCustomPresets: + @classmethod + def INPUT_TYPES(cls): + try: + with open(os.path.join(script_directory, 'custom_dimensions.json')) as f: + dimensions_dict = json.load(f) + except FileNotFoundError: + dimensions_dict = [] + return { + "required": { + "dimensions": ( + [f"{d['label']} - {d['value']}" for d in dimensions_dict], + ), + + "invert": ("BOOLEAN", {"default": False}), + "batch_size": ("INT", { + "default": 1, + "min": 1, + "max": 4096 + }), + }, + } + + RETURN_TYPES = ("LATENT", "INT", "INT") + RETURN_NAMES = ("Latent", "Width", "Height") + FUNCTION = "generate" + CATEGORY = "KJNodes/latents" + DESCRIPTION = """ +Generates an empty latent image with the specified dimensions. +The choices are loaded from 'custom_dimensions.json' in the nodes folder. +""" + + def generate(self, dimensions, invert, batch_size): + from nodes import EmptyLatentImage + # Split the string into label and value + label, value = dimensions.split(' - ') + # Split the value into width and height + width, height = [x.strip() for x in value.split('x')] + + if invert: + width, height = height, width + + latent = EmptyLatentImage().generate(int(width), int(height), batch_size)[0] + + return (latent, int(width), int(height),) + +class WidgetToString: + @classmethod + def IS_CHANGED(cls,*,id,node_title,any_input,**kwargs): + if any_input is not None and (id != 0 or node_title != ""): + return float("NaN") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "id": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}), + "widget_name": ("STRING", {"multiline": False}), + "return_all": ("BOOLEAN", {"default": False}), + }, + "optional": { + "any_input": (IO.ANY, ), + "node_title": ("STRING", {"multiline": False}), + "allowed_float_decimals": ("INT", {"default": 2, "min": 0, "max": 10, "tooltip": "Number of decimal places to display for float values"}), + }, + "hidden": {"extra_pnginfo": "EXTRA_PNGINFO", + "prompt": "PROMPT", + "unique_id": "UNIQUE_ID",}, + } + + RETURN_TYPES = ("STRING", ) + FUNCTION = "get_widget_value" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +Selects a node and it's specified widget and outputs the value as a string. +If no node id or title is provided it will use the 'any_input' link and use that node. +To see node id's, enable "Node ID Badge Mode" in main settings. +Alternatively you can search with the node title. Node titles ONLY exist if they +are manually edited! +'widget_name' can be a comma separated list. +The 'any_input' is required for making sure the node you want the value from exists in the workflow. +""" + + def get_widget_value(self, id, widget_name, extra_pnginfo, prompt, unique_id, return_all=False, any_input=None, node_title="", allowed_float_decimals=2): + workflow = extra_pnginfo["workflow"] + #print(json.dumps(workflow, indent=4)) + results = [] + node_id = link_id = subgraph_prefix = None + link_to_node_map = {} + node_to_subgraph_map = {} # Track which subgraph each node belongs to + + # Parse unique_id - handle both "parent:id" format and simple int format + if isinstance(unique_id, str) and ":" in unique_id: + unique_id_parts = unique_id.split(":") + unique_id_int = int(unique_id_parts[-1]) # Use the last part as the node id + subgraph_prefix = ":".join(unique_id_parts[:-1]) # Store the parent prefix (e.g., "14") + else: + unique_id_int = int(unique_id) + + # Collect all nodes from main workflow and subgraphs + all_nodes = list(workflow.get("nodes", [])) + definitions = workflow.get("definitions", {}) + subgraphs = definitions.get("subgraphs", []) + + # Find which main workflow node references each subgraph + subgraph_id_to_parent = {} + for node in workflow.get("nodes", []): + node_type = node.get("type", "") + # Subgraph nodes have a UUID as their type + if "-" in node_type and len(node_type) == 36: # UUID format check + subgraph_id_to_parent[node_type] = node["id"] + + for subgraph in subgraphs: + subgraph_id = subgraph.get("id", "") + parent_node_id = subgraph_id_to_parent.get(subgraph_id) + + subgraph_nodes = subgraph.get("nodes", []) + for node in subgraph_nodes: + # Track which subgraph (parent node) this node belongs to + if parent_node_id is not None: + node_to_subgraph_map[node["id"]] = parent_node_id + all_nodes.extend(subgraph_nodes) + + # Also build link_to_node_map from subgraph links + subgraph_links = subgraph.get("links", []) + for link in subgraph_links: + # link format: [link_id, origin_id, origin_slot, target_id, target_slot, type] + if isinstance(link, dict): + link_to_node_map[link["id"]] = link["origin_id"] + elif isinstance(link, list) and len(link) >= 2: + link_to_node_map[link[0]] = link[1] + + for node in all_nodes: + if node_title: + if "title" in node: + if node["title"] == node_title: + node_id = node["id"] + break + else: + logging.warning("Node title not found.") + elif id != 0: + if node["id"] == id: + node_id = id + break + elif any_input is not None: + if node["type"] == "WidgetToString" and node["id"] == unique_id_int and not link_id: + for node_input in node["inputs"]: + if node_input["name"] == "any_input": + link_id = node_input["link"] + + # Construct a map of links to node IDs for future reference + node_outputs = node.get("outputs", None) + if not node_outputs: + continue + for output in node_outputs: + node_links = output.get("links", None) + if not node_links: + continue + for link in node_links: + link_to_node_map[link] = node["id"] + if link_id and link == link_id: + break + + if link_id: + node_id = link_to_node_map.get(link_id, None) + + if node_id is None: + raise ValueError("No matching node found for the given title or id") + + # Determine the correct prompt key + # First check if the target node is in a subgraph + target_subgraph_parent = node_to_subgraph_map.get(node_id) + + if target_subgraph_parent is not None: + # Target node is in a subgraph, use the parent node id as prefix + prompt_key = f"{target_subgraph_parent}:{node_id}" + elif subgraph_prefix is not None: + # We're in a subgraph, use our prefix + prompt_key = f"{subgraph_prefix}:{node_id}" + else: + prompt_key = str(node_id) + + # Try the prefixed key first, then fall back to just the node_id + if prompt_key not in prompt: + prompt_key = str(node_id) + + if prompt_key not in prompt: + raise KeyError(f"Node not found in prompt. Tried keys: '{target_subgraph_parent}:{node_id}' and '{node_id}'") + + values = prompt[prompt_key] + if "inputs" in values: + inputs = values["inputs"] + + # support comma-separated list and trim whitespace + widget_names = [] + if widget_name: + widget_names = [w.strip() for w in widget_name.split(",") if w.strip()] + + if return_all: + # Format items based on type + formatted_items = [] + for k, v in inputs.items(): + if isinstance(v, float): + item = f"{k}: {v:.{allowed_float_decimals}f}" + else: + item = f"{k}: {str(v)}" + formatted_items.append(item) + results.append(", ".join(formatted_items)) + + # Single widget name (trimmed) + elif len(widget_names) == 1: + name = widget_names[0] + if name in inputs: + v = inputs[name] + if isinstance(v, float): + v = f"{v:.{allowed_float_decimals}f}" + else: + v = str(v) + return (v, ) + else: + raise NameError(f"Widget not found: {node_id}.{name}") + + # Multiple widget names: return "name: value" pairs + elif len(widget_names) > 1: + formatted_items = [] + for name in widget_names: + if name not in inputs: + raise NameError(f"Widget not found: {node_id}.{name}") + v = inputs[name] + if isinstance(v, float): + v = f"{v:.{allowed_float_decimals}f}" + else: + v = str(v) + formatted_items.append(f"{name}: {v}") + return (", ".join(formatted_items), ) + + else: + # No valid widget name provided + raise NameError(f"Widget not found: {node_id}.{widget_name}") + + return (", ".join(results).strip(", "), ) + + +class DummyOut: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "any_input": (IO.ANY, ), + } + } + + RETURN_TYPES = (IO.ANY,) + FUNCTION = "dummy" + CATEGORY = "KJNodes/misc" + OUTPUT_NODE = True + DESCRIPTION = """ +Does nothing, used to trigger generic workflow output. +A way to get previews in the UI without saving anything to disk. +""" + + def dummy(self, any_input): + return (any_input,) + +class FlipSigmasAdjusted: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"sigmas": ("SIGMAS", ), + "divide_by_last_sigma": ("BOOLEAN", {"default": False}), + "divide_by": ("FLOAT", {"default": 1,"min": 1, "max": 255, "step": 0.01}), + "offset_by": ("INT", {"default": 1,"min": -100, "max": 100, "step": 1}), + } + } + RETURN_TYPES = ("SIGMAS", "STRING",) + RETURN_NAMES = ("SIGMAS", "sigmas_string",) + CATEGORY = "KJNodes/noise" + FUNCTION = "get_sigmas_adjusted" + + def get_sigmas_adjusted(self, sigmas, divide_by_last_sigma, divide_by, offset_by): + + sigmas = sigmas.flip(0) + if sigmas[0] == 0: + sigmas[0] = 0.0001 + adjusted_sigmas = sigmas.clone() + #offset sigma + for i in range(1, len(sigmas)): + offset_index = i - offset_by + if 0 <= offset_index < len(sigmas): + adjusted_sigmas[i] = sigmas[offset_index] + else: + adjusted_sigmas[i] = 0.0001 + if adjusted_sigmas[0] == 0: + adjusted_sigmas[0] = 0.0001 + if divide_by_last_sigma: + adjusted_sigmas = adjusted_sigmas / adjusted_sigmas[-1] + + sigma_np_array = adjusted_sigmas.numpy() + array_string = np.array2string(sigma_np_array, precision=2, separator=', ', threshold=np.inf) + adjusted_sigmas = adjusted_sigmas / divide_by + return (adjusted_sigmas, array_string,) + +class CustomSigmas: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "sigmas_string" :("STRING", {"default": "14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029","multiline": True}), + "interpolate_to_steps": ("INT", {"default": 10,"min": 0, "max": 255, "step": 1}), + } + } + RETURN_TYPES = ("SIGMAS",) + RETURN_NAMES = ("SIGMAS",) + CATEGORY = "KJNodes/noise" + FUNCTION = "customsigmas" + DESCRIPTION = """ +Creates a sigmas tensor from a string of comma separated values. +Examples: + +Nvidia's optimized AYS 10 step schedule for SD 1.5: +14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029 +SDXL: +14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.029 +SVD: +700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002 +""" + def customsigmas(self, sigmas_string, interpolate_to_steps): + sigmas_list = sigmas_string.split(', ') + sigmas_float_list = [float(sigma) for sigma in sigmas_list] + sigmas_tensor = torch.FloatTensor(sigmas_float_list) + if len(sigmas_tensor) != interpolate_to_steps + 1: + sigmas_tensor = self.loglinear_interp(sigmas_tensor, interpolate_to_steps + 1) + sigmas_tensor[-1] = 0 + return (sigmas_tensor.float(),) + + def loglinear_interp(self, t_steps, num_steps): + """ + Performs log-linear interpolation of a given array of decreasing numbers. + """ + t_steps_np = t_steps.numpy() + + xs = np.linspace(0, 1, len(t_steps_np)) + ys = np.log(t_steps_np[::-1]) + + new_xs = np.linspace(0, 1, num_steps) + new_ys = np.interp(new_xs, xs, ys) + + interped_ys = np.exp(new_ys)[::-1].copy() + interped_ys_tensor = torch.tensor(interped_ys) + return interped_ys_tensor + +class StringToFloatList: + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "string" :("STRING", {"default": "1, 2, 3", "multiline": True}), + } + } + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("FLOAT",) + CATEGORY = "KJNodes/misc" + FUNCTION = "createlist" + + def createlist(self, string): + float_list = [float(x.strip()) for x in string.split(',')] + return (float_list,) + + +class InjectNoiseToLatent: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "latents":("LATENT",), + "strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 200.0, "step": 0.0001}), + "noise": ("LATENT",), + "normalize": ("BOOLEAN", {"default": False}), + "average": ("BOOLEAN", {"default": False}), + }, + "optional":{ + "mask": ("MASK", ), + "mix_randn_amount": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.001}), + "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), + } + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "injectnoise" + CATEGORY = "KJNodes/noise" + + def injectnoise(self, latents, strength, noise, normalize, average, mix_randn_amount=0, seed=None, mask=None): + samples = latents["samples"].clone().cpu() + noise = noise["samples"].clone().cpu() + if samples.shape != samples.shape: + raise ValueError("InjectNoiseToLatent: Latent and noise must have the same shape") + if average: + noised = (samples + noise) / 2 + else: + noised = samples + noise * strength + if normalize: + noised = noised / noised.std() + if mask is not None: + mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(noised.shape[2], noised.shape[3]), mode="bilinear") + mask = mask.expand((-1,noised.shape[1],-1,-1)) + if mask.shape[0] < noised.shape[0]: + mask = mask.repeat((noised.shape[0] -1) // mask.shape[0] + 1, 1, 1, 1)[:noised.shape[0]] + noised = mask * noised + (1-mask) * samples + if mix_randn_amount > 0: + if seed is not None: + generator = torch.manual_seed(seed) + rand_noise = torch.randn(noised.size(), dtype=noised.dtype, layout=noised.layout, generator=generator, device="cpu") + noised = noised + (mix_randn_amount * rand_noise) + + return ({"samples":noised},) + +class SoundReactive: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "sound_level": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 99999, "step": 0.01}), + "start_range_hz": ("INT", {"default": 150, "min": 0, "max": 9999, "step": 1}), + "end_range_hz": ("INT", {"default": 2000, "min": 0, "max": 9999, "step": 1}), + "multiplier": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 99999, "step": 0.01}), + "smoothing_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "normalize": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("FLOAT","INT",) + RETURN_NAMES =("sound_level", "sound_level_int",) + FUNCTION = "react" + CATEGORY = "KJNodes/audio" + DESCRIPTION = """ +Reacts to the sound level of the input. +Uses your browsers sound input options and requires. +Meant to be used with realtime diffusion with autoqueue. +""" + + def react(self, sound_level, start_range_hz, end_range_hz, smoothing_factor, multiplier, normalize): + + sound_level *= multiplier + + if normalize: + sound_level /= 255 + + sound_level_int = int(sound_level) + return (sound_level, sound_level_int, ) + +class GenerateNoise: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), + "multiplier": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 4096, "step": 0.01}), + "constant_batch_noise": ("BOOLEAN", {"default": False}), + "normalize": ("BOOLEAN", {"default": False}), + }, + "optional": { + "model": ("MODEL", ), + "sigmas": ("SIGMAS", ), + "latent_channels": (['4', '16', ],), + "shape": (["BCHW", "BCTHW","BTCHW",],), + } + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "generatenoise" + CATEGORY = "KJNodes/noise" + DESCRIPTION = """ +Generates noise for injection or to be used as empty latents on samplers with add_noise off. +""" + + def generatenoise(self, batch_size, width, height, seed, multiplier, constant_batch_noise, normalize, sigmas=None, model=None, latent_channels=4, shape="BCHW"): + + generator = torch.manual_seed(seed) + if shape == "BCHW": + noise = torch.randn([batch_size, int(latent_channels), height // 8, width // 8], dtype=torch.float32, layout=torch.strided, generator=generator, device="cpu") + elif shape == "BCTHW": + noise = torch.randn([1, int(latent_channels), batch_size,height // 8, width // 8], dtype=torch.float32, layout=torch.strided, generator=generator, device="cpu") + elif shape == "BTCHW": + noise = torch.randn([1, batch_size, int(latent_channels), height // 8, width // 8], dtype=torch.float32, layout=torch.strided, generator=generator, device="cpu") + if sigmas is not None: + sigma = sigmas[0] - sigmas[-1] + sigma /= model.model.latent_format.scale_factor + noise *= sigma + + noise *=multiplier + + if normalize: + noise = noise / noise.std() + if constant_batch_noise: + noise = noise[0].repeat(batch_size, 1, 1, 1) + + + return ({"samples":noise}, ) + +def camera_embeddings(elevation, azimuth): + elevation = torch.as_tensor([elevation]) + azimuth = torch.as_tensor([azimuth]) + embeddings = torch.stack( + [ + torch.deg2rad( + (90 - elevation) - (90) + ), # Zero123 polar is 90-elevation + torch.sin(torch.deg2rad(azimuth)), + torch.cos(torch.deg2rad(azimuth)), + torch.deg2rad( + 90 - torch.full_like(elevation, 0) + ), + ], dim=-1).unsqueeze(1) + + return embeddings + +def interpolate_angle(start, end, fraction): + # Calculate the difference in angles and adjust for wraparound if necessary + diff = (end - start + 540) % 360 - 180 + # Apply fraction to the difference + interpolated = start + fraction * diff + # Normalize the result to be within the range of -180 to 180 + return (interpolated + 180) % 360 - 180 + + +class StableZero123_BatchSchedule: + @classmethod + def INPUT_TYPES(s): + return {"required": { "clip_vision": ("CLIP_VISION",), + "init_image": ("IMAGE",), + "vae": ("VAE",), + "width": ("INT", {"default": 256, "min": 16, "max": MAX_RESOLUTION, "step": 8}), + "height": ("INT", {"default": 256, "min": 16, "max": MAX_RESOLUTION, "step": 8}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), + "azimuth_points_string": ("STRING", {"default": "0:(0.0),\n7:(1.0),\n15:(0.0)\n", "multiline": True}), + "elevation_points_string": ("STRING", {"default": "0:(0.0),\n7:(0.0),\n15:(0.0)\n", "multiline": True}), + }} + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive", "negative", "latent") + FUNCTION = "encode" + CATEGORY = "KJNodes/experimental" + + def encode(self, clip_vision, init_image, vae, width, height, batch_size, azimuth_points_string, elevation_points_string, interpolation): + output = clip_vision.encode_image(init_image) + pooled = output.image_embeds.unsqueeze(0) + pixels = common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1) + encode_pixels = pixels[:,:,:,:3] + t = vae.encode(encode_pixels) + + def ease_in(t): + return t * t + def ease_out(t): + return 1 - (1 - t) * (1 - t) + def ease_in_out(t): + return 3 * t * t - 2 * t * t * t + + # Parse the azimuth input string into a list of tuples + azimuth_points = [] + azimuth_points_string = azimuth_points_string.rstrip(',\n') + for point_str in azimuth_points_string.split(','): + frame_str, azimuth_str = point_str.split(':') + frame = int(frame_str.strip()) + azimuth = float(azimuth_str.strip()[1:-1]) + azimuth_points.append((frame, azimuth)) + # Sort the points by frame number + azimuth_points.sort(key=lambda x: x[0]) + + # Parse the elevation input string into a list of tuples + elevation_points = [] + elevation_points_string = elevation_points_string.rstrip(',\n') + for point_str in elevation_points_string.split(','): + frame_str, elevation_str = point_str.split(':') + frame = int(frame_str.strip()) + elevation_val = float(elevation_str.strip()[1:-1]) + elevation_points.append((frame, elevation_val)) + # Sort the points by frame number + elevation_points.sort(key=lambda x: x[0]) + + # Index of the next point to interpolate towards + next_point = 1 + next_elevation_point = 1 + + positive_cond_out = [] + positive_pooled_out = [] + negative_cond_out = [] + negative_pooled_out = [] + + #azimuth interpolation + for i in range(batch_size): + # Find the interpolated azimuth for the current frame + while next_point < len(azimuth_points) and i >= azimuth_points[next_point][0]: + next_point += 1 + # If next_point is equal to the length of points, we've gone past the last point + if next_point == len(azimuth_points): + next_point -= 1 # Set next_point to the last index of points + prev_point = max(next_point - 1, 0) # Ensure prev_point is not less than 0 + + # Calculate fraction + if azimuth_points[next_point][0] != azimuth_points[prev_point][0]: # Prevent division by zero + fraction = (i - azimuth_points[prev_point][0]) / (azimuth_points[next_point][0] - azimuth_points[prev_point][0]) + if interpolation == "ease_in": + fraction = ease_in(fraction) + elif interpolation == "ease_out": + fraction = ease_out(fraction) + elif interpolation == "ease_in_out": + fraction = ease_in_out(fraction) + + # Use the new interpolate_angle function + interpolated_azimuth = interpolate_angle(azimuth_points[prev_point][1], azimuth_points[next_point][1], fraction) + else: + interpolated_azimuth = azimuth_points[prev_point][1] + # Interpolate the elevation + next_elevation_point = 1 + while next_elevation_point < len(elevation_points) and i >= elevation_points[next_elevation_point][0]: + next_elevation_point += 1 + if next_elevation_point == len(elevation_points): + next_elevation_point -= 1 + prev_elevation_point = max(next_elevation_point - 1, 0) + + if elevation_points[next_elevation_point][0] != elevation_points[prev_elevation_point][0]: + fraction = (i - elevation_points[prev_elevation_point][0]) / (elevation_points[next_elevation_point][0] - elevation_points[prev_elevation_point][0]) + if interpolation == "ease_in": + fraction = ease_in(fraction) + elif interpolation == "ease_out": + fraction = ease_out(fraction) + elif interpolation == "ease_in_out": + fraction = ease_in_out(fraction) + + interpolated_elevation = interpolate_angle(elevation_points[prev_elevation_point][1], elevation_points[next_elevation_point][1], fraction) + else: + interpolated_elevation = elevation_points[prev_elevation_point][1] + + cam_embeds = camera_embeddings(interpolated_elevation, interpolated_azimuth) + cond = torch.cat([pooled, cam_embeds.repeat((pooled.shape[0], 1, 1))], dim=-1) + + positive_pooled_out.append(t) + positive_cond_out.append(cond) + negative_pooled_out.append(torch.zeros_like(t)) + negative_cond_out.append(torch.zeros_like(pooled)) + + # Concatenate the conditions and pooled outputs + final_positive_cond = torch.cat(positive_cond_out, dim=0) + final_positive_pooled = torch.cat(positive_pooled_out, dim=0) + final_negative_cond = torch.cat(negative_cond_out, dim=0) + final_negative_pooled = torch.cat(negative_pooled_out, dim=0) + + # Structure the final output + final_positive = [[final_positive_cond, {"concat_latent_image": final_positive_pooled}]] + final_negative = [[final_negative_cond, {"concat_latent_image": final_negative_pooled}]] + + latent = torch.zeros([batch_size, 4, height // 8, width // 8]) + return (final_positive, final_negative, {"samples": latent}) + +def linear_interpolate(start, end, fraction): + return start + (end - start) * fraction + +class SV3D_BatchSchedule: + @classmethod + def INPUT_TYPES(s): + return {"required": { "clip_vision": ("CLIP_VISION",), + "init_image": ("IMAGE",), + "vae": ("VAE",), + "width": ("INT", {"default": 576, "min": 16, "max": MAX_RESOLUTION, "step": 8}), + "height": ("INT", {"default": 576, "min": 16, "max": MAX_RESOLUTION, "step": 8}), + "batch_size": ("INT", {"default": 21, "min": 1, "max": 4096}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), + "azimuth_points_string": ("STRING", {"default": "0:(0.0),\n9:(180.0),\n20:(360.0)\n", "multiline": True}), + "elevation_points_string": ("STRING", {"default": "0:(0.0),\n9:(0.0),\n20:(0.0)\n", "multiline": True}), + }} + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive", "negative", "latent") + FUNCTION = "encode" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = """ +Allow scheduling of the azimuth and elevation conditions for SV3D. +Note that SV3D is still a video model and the schedule needs to always go forward +https://huggingface.co/stabilityai/sv3d +""" + + def encode(self, clip_vision, init_image, vae, width, height, batch_size, azimuth_points_string, elevation_points_string, interpolation): + output = clip_vision.encode_image(init_image) + pooled = output.image_embeds.unsqueeze(0) + pixels = common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1) + encode_pixels = pixels[:,:,:,:3] + t = vae.encode(encode_pixels) + + def ease_in(t): + return t * t + def ease_out(t): + return 1 - (1 - t) * (1 - t) + def ease_in_out(t): + return 3 * t * t - 2 * t * t * t + + # Parse the azimuth input string into a list of tuples + azimuth_points = [] + azimuth_points_string = azimuth_points_string.rstrip(',\n') + for point_str in azimuth_points_string.split(','): + frame_str, azimuth_str = point_str.split(':') + frame = int(frame_str.strip()) + azimuth = float(azimuth_str.strip()[1:-1]) + azimuth_points.append((frame, azimuth)) + # Sort the points by frame number + azimuth_points.sort(key=lambda x: x[0]) + + # Parse the elevation input string into a list of tuples + elevation_points = [] + elevation_points_string = elevation_points_string.rstrip(',\n') + for point_str in elevation_points_string.split(','): + frame_str, elevation_str = point_str.split(':') + frame = int(frame_str.strip()) + elevation_val = float(elevation_str.strip()[1:-1]) + elevation_points.append((frame, elevation_val)) + # Sort the points by frame number + elevation_points.sort(key=lambda x: x[0]) + + # Index of the next point to interpolate towards + next_point = 1 + next_elevation_point = 1 + elevations = [] + azimuths = [] + # For azimuth interpolation + for i in range(batch_size): + # Find the interpolated azimuth for the current frame + while next_point < len(azimuth_points) and i >= azimuth_points[next_point][0]: + next_point += 1 + if next_point == len(azimuth_points): + next_point -= 1 + prev_point = max(next_point - 1, 0) + + if azimuth_points[next_point][0] != azimuth_points[prev_point][0]: + fraction = (i - azimuth_points[prev_point][0]) / (azimuth_points[next_point][0] - azimuth_points[prev_point][0]) + # Apply the ease function to the fraction + if interpolation == "ease_in": + fraction = ease_in(fraction) + elif interpolation == "ease_out": + fraction = ease_out(fraction) + elif interpolation == "ease_in_out": + fraction = ease_in_out(fraction) + + interpolated_azimuth = linear_interpolate(azimuth_points[prev_point][1], azimuth_points[next_point][1], fraction) + else: + interpolated_azimuth = azimuth_points[prev_point][1] + + # Interpolate the elevation + next_elevation_point = 1 + while next_elevation_point < len(elevation_points) and i >= elevation_points[next_elevation_point][0]: + next_elevation_point += 1 + if next_elevation_point == len(elevation_points): + next_elevation_point -= 1 + prev_elevation_point = max(next_elevation_point - 1, 0) + + if elevation_points[next_elevation_point][0] != elevation_points[prev_elevation_point][0]: + fraction = (i - elevation_points[prev_elevation_point][0]) / (elevation_points[next_elevation_point][0] - elevation_points[prev_elevation_point][0]) + # Apply the ease function to the fraction + if interpolation == "ease_in": + fraction = ease_in(fraction) + elif interpolation == "ease_out": + fraction = ease_out(fraction) + elif interpolation == "ease_in_out": + fraction = ease_in_out(fraction) + + interpolated_elevation = linear_interpolate(elevation_points[prev_elevation_point][1], elevation_points[next_elevation_point][1], fraction) + else: + interpolated_elevation = elevation_points[prev_elevation_point][1] + + azimuths.append(interpolated_azimuth) + elevations.append(interpolated_elevation) + + #print("azimuths", azimuths) + #print("elevations", elevations) + + # Structure the final output + final_positive = [[pooled, {"concat_latent_image": t, "elevation": elevations, "azimuth": azimuths}]] + final_negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t),"elevation": elevations, "azimuth": azimuths}]] + + latent = torch.zeros([batch_size, 4, height // 8, width // 8]) + return (final_positive, final_negative, {"samples": latent}) + + +class Superprompt: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "instruction_prompt": ("STRING", {"default": 'Expand the following prompt to add more detail', "multiline": True}), + "prompt": ("STRING", {"default": '', "multiline": True, "forceInput": True}), + "max_new_tokens": ("INT", {"default": 128, "min": 1, "max": 4096, "step": 1}), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "process" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +# SuperPrompt +A T5 model fine-tuned on the SuperPrompt dataset for +upsampling text prompts to more detailed descriptions. +Meant to be used as a pre-generation step for text-to-image +models that benefit from more detailed prompts. +https://huggingface.co/roborovski/superprompt-v1 +""" + + def process(self, instruction_prompt, prompt, max_new_tokens): + device = model_management.get_torch_device() + from transformers import T5Tokenizer, T5ForConditionalGeneration + + checkpoint_path = os.path.join(script_directory, "models","superprompt-v1") + if not os.path.exists(checkpoint_path): + logging.info(f"Downloading model to: {checkpoint_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="roborovski/superprompt-v1", + local_dir=checkpoint_path, + local_dir_use_symlinks=False) + tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small", legacy=False) + + model = T5ForConditionalGeneration.from_pretrained(checkpoint_path, device_map=device) + model.to(device) + input_text = instruction_prompt + ": " + prompt + + input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device) + outputs = model.generate(input_ids, max_new_tokens=max_new_tokens) + out = (tokenizer.decode(outputs[0])) + out = out.replace('', '') + out = out.replace('', '') + + return (out, ) + + +class CameraPoseVisualizer: + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "pose_file_path": ("STRING", {"default": '', "multiline": False}), + "base_xval": ("FLOAT", {"default": 0.2,"min": 0, "max": 100, "step": 0.01}), + "zval": ("FLOAT", {"default": 0.3,"min": 0, "max": 100, "step": 0.01}), + "scale": ("FLOAT", {"default": 1.0,"min": 0.01, "max": 10.0, "step": 0.01}), + "use_exact_fx": ("BOOLEAN", {"default": False}), + "relative_c2w": ("BOOLEAN", {"default": True}), + "use_viewer": ("BOOLEAN", {"default": False}), + }, + "optional": { + "cameractrl_poses": ("CAMERACTRL_POSES", {"default": None}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "plot" + CATEGORY = "KJNodes/misc" + DESCRIPTION = """ +Visualizes the camera poses, from Animatediff-Evolved CameraCtrl Pose +or a .txt file with RealEstate camera intrinsics and coordinates, in a 3D plot. +""" + + def plot(self, pose_file_path, scale, base_xval, zval, use_exact_fx, relative_c2w, use_viewer, cameractrl_poses=None): + import matplotlib as mpl + import matplotlib.pyplot as plt + from torchvision.transforms import ToTensor + + x_min = -2.0 * scale + x_max = 2.0 * scale + y_min = -2.0 * scale + y_max = 2.0 * scale + z_min = -2.0 * scale + z_max = 2.0 * scale + plt.rcParams['text.color'] = '#999999' + self.fig = plt.figure(figsize=(18, 7)) + self.fig.patch.set_facecolor('#353535') + self.ax = self.fig.add_subplot(projection='3d') + self.ax.set_facecolor('#353535') # Set the background color here + self.ax.grid(color='#999999', linestyle='-', linewidth=0.5) + self.plotly_data = None # plotly data traces + self.ax.set_aspect("auto") + self.ax.set_xlim(x_min, x_max) + self.ax.set_ylim(y_min, y_max) + self.ax.set_zlim(z_min, z_max) + self.ax.set_xlabel('x', color='#999999') + self.ax.set_ylabel('y', color='#999999') + self.ax.set_zlabel('z', color='#999999') + for text in self.ax.get_xticklabels() + self.ax.get_yticklabels() + self.ax.get_zticklabels(): + text.set_color('#999999') + logging.info('initialize camera pose visualizer') + + if pose_file_path != "": + with open(pose_file_path, 'r') as f: + poses = f.readlines() + w2cs = [np.asarray([float(p) for p in pose.strip().split(' ')[7:]]).reshape(3, 4) for pose in poses[1:]] + fxs = [float(pose.strip().split(' ')[1]) for pose in poses[1:]] + #print(poses) + elif cameractrl_poses is not None: + poses = cameractrl_poses + w2cs = [np.array(pose[7:]).reshape(3, 4) for pose in cameractrl_poses] + fxs = [pose[1] for pose in cameractrl_poses] + else: + raise ValueError("Please provide either pose_file_path or cameractrl_poses") + + total_frames = len(w2cs) + transform_matrix = np.asarray([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]]).reshape(4, 4) + last_row = np.zeros((1, 4)) + last_row[0, -1] = 1.0 + + w2cs = [np.concatenate((w2c, last_row), axis=0) for w2c in w2cs] + c2ws = self.get_c2w(w2cs, transform_matrix, relative_c2w) + + for frame_idx, c2w in enumerate(c2ws): + self.extrinsic2pyramid(c2w, frame_idx / total_frames, hw_ratio=1/1, base_xval=base_xval, + zval=(fxs[frame_idx] if use_exact_fx else zval)) + + # Create the colorbar + cmap = mpl.cm.rainbow + norm = mpl.colors.Normalize(vmin=0, vmax=total_frames) + colorbar = self.fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=self.ax, orientation='vertical') + + # Change the colorbar label + colorbar.set_label('Frame', color='#999999') # Change the label and its color + + # Change the tick colors + colorbar.ax.yaxis.set_tick_params(colors='#999999') # Change the tick color + + # Change the tick frequency + # Assuming you want to set the ticks at every 10th frame + ticks = np.arange(0, total_frames, 10) + colorbar.ax.yaxis.set_ticks(ticks) + + plt.title('') + plt.draw() + buf = BytesIO() + plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0) + buf.seek(0) + img = Image.open(buf) + tensor_img = ToTensor()(img) + buf.close() + tensor_img = tensor_img.permute(1, 2, 0).unsqueeze(0) + if use_viewer: + time.sleep(1) + plt.show() + return (tensor_img,) + + def extrinsic2pyramid(self, extrinsic, color_map='red', hw_ratio=1/1, base_xval=1, zval=3): + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d.art3d import Poly3DCollection + vertex_std = np.array([[0, 0, 0, 1], + [base_xval, -base_xval * hw_ratio, zval, 1], + [base_xval, base_xval * hw_ratio, zval, 1], + [-base_xval, base_xval * hw_ratio, zval, 1], + [-base_xval, -base_xval * hw_ratio, zval, 1]]) + vertex_transformed = vertex_std @ extrinsic.T + meshes = [[vertex_transformed[0, :-1], vertex_transformed[1][:-1], vertex_transformed[2, :-1]], + [vertex_transformed[0, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1]], + [vertex_transformed[0, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]], + [vertex_transformed[0, :-1], vertex_transformed[4, :-1], vertex_transformed[1, :-1]], + [vertex_transformed[1, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]]] + + color = color_map if isinstance(color_map, str) else plt.cm.rainbow(color_map) + + self.ax.add_collection3d( + Poly3DCollection(meshes, facecolors=color, linewidths=0.3, edgecolors=color, alpha=0.25)) + + def customize_legend(self, list_label): + from matplotlib.patches import Patch + import matplotlib.pyplot as plt + list_handle = [] + for idx, label in enumerate(list_label): + color = plt.cm.rainbow(idx / len(list_label)) + patch = Patch(color=color, label=label) + list_handle.append(patch) + plt.legend(loc='right', bbox_to_anchor=(1.8, 0.5), handles=list_handle) + + def get_c2w(self, w2cs, transform_matrix, relative_c2w): + if relative_c2w: + target_cam_c2w = np.array([ + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1] + ]) + abs2rel = target_cam_c2w @ w2cs[0] + ret_poses = [target_cam_c2w, ] + [abs2rel @ np.linalg.inv(w2c) for w2c in w2cs[1:]] + else: + ret_poses = [np.linalg.inv(w2c) for w2c in w2cs] + ret_poses = [transform_matrix @ x for x in ret_poses] + return np.array(ret_poses, dtype=np.float32) + + + +class CheckpointPerturbWeights: + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "joint_blocks": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 10.0, "step": 0.001}), + "final_layer": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 10.0, "step": 0.001}), + "rest_of_the_blocks": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 10.0, "step": 0.001}), + "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), + } + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "mod" + OUTPUT_NODE = True + + CATEGORY = "KJNodes/experimental" + + def mod(self, seed, model, joint_blocks, final_layer, rest_of_the_blocks): + import copy + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + device = model_management.get_torch_device() + model_copy = copy.deepcopy(model) + model_copy.model.to(device) + keys = model_copy.model.diffusion_model.state_dict().keys() + + dict = {} + for key in keys: + dict[key] = model_copy.model.diffusion_model.state_dict()[key] + + pbar = ProgressBar(len(keys)) + for k in keys: + v = dict[k] + logging.info(f'{k}: {v.std()}') + if k.startswith('joint_blocks'): + multiplier = joint_blocks + elif k.startswith('final_layer'): + multiplier = final_layer + else: + multiplier = rest_of_the_blocks + dict[k] += torch.normal(torch.zeros_like(v) * v.mean(), torch.ones_like(v) * v.std() * multiplier).to(device) + pbar.update(1) + model_copy.model.diffusion_model.load_state_dict(dict) + return model_copy, + +class DifferentialDiffusionAdvanced(): + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL", ), + "samples": ("LATENT",), + "mask": ("MASK",), + "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}), + }} + RETURN_TYPES = ("MODEL", "LATENT") + FUNCTION = "apply" + CATEGORY = "_for_testing" + INIT = False + + def apply(self, model, samples, mask, multiplier): + self.multiplier = multiplier + model = model.clone() + model.set_model_denoise_mask_function(self.forward) + s = samples.copy() + s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + return (model, s) + + def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict): + model = extra_options["model"] + step_sigmas = extra_options["sigmas"] + sigma_to = model.inner_model.model_sampling.sigma_min + if step_sigmas[-1] > sigma_to: + sigma_to = step_sigmas[-1] + sigma_from = step_sigmas[0] + + ts_from = model.inner_model.model_sampling.timestep(sigma_from) + ts_to = model.inner_model.model_sampling.timestep(sigma_to) + current_ts = model.inner_model.model_sampling.timestep(sigma[0]) + + threshold = (current_ts - ts_to) / (ts_from - ts_to) / self.multiplier + + return (denoise_mask >= threshold).to(denoise_mask.dtype) + +class FluxBlockLoraSelect: + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(s): + arg_dict = {} + argument = ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01}) + + for i in range(19): + arg_dict["double_blocks.{}.".format(i)] = argument + + for i in range(38): + arg_dict["single_blocks.{}.".format(i)] = argument + + return {"required": arg_dict} + + RETURN_TYPES = ("SELECTEDDITBLOCKS", ) + RETURN_NAMES = ("blocks", ) + OUTPUT_TOOLTIPS = ("The modified diffusion model.",) + FUNCTION = "load_lora" + + CATEGORY = "KJNodes/experimental" + DESCRIPTION = "Select individual block alpha values, value of 0 removes the block altogether" + + def load_lora(self, **kwargs): + return (kwargs,) + +class HunyuanVideoBlockLoraSelect: + @classmethod + def INPUT_TYPES(s): + arg_dict = {} + argument = ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01}) + + for i in range(20): + arg_dict["double_blocks.{}.".format(i)] = argument + + for i in range(40): + arg_dict["single_blocks.{}.".format(i)] = argument + + return {"required": arg_dict} + + RETURN_TYPES = ("SELECTEDDITBLOCKS", ) + RETURN_NAMES = ("blocks", ) + OUTPUT_TOOLTIPS = ("The modified diffusion model.",) + FUNCTION = "load_lora" + + CATEGORY = "KJNodes/hunyuanvideo" + DESCRIPTION = "Select individual block alpha values, value of 0 removes the block altogether" + + def load_lora(self, **kwargs): + return (kwargs,) + +class Wan21BlockLoraSelect: + @classmethod + def INPUT_TYPES(s): + arg_dict = {} + argument = ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01}) + + for i in range(40): + arg_dict["blocks.{}.".format(i)] = argument + + return {"required": arg_dict} + + RETURN_TYPES = ("SELECTEDDITBLOCKS", ) + RETURN_NAMES = ("blocks", ) + OUTPUT_TOOLTIPS = ("The modified diffusion model.",) + FUNCTION = "load_lora" + + CATEGORY = "KJNodes/wan" + DESCRIPTION = "Select individual block alpha values, value of 0 removes the block altogether" + + def load_lora(self, **kwargs): + return (kwargs,) + +class LTX2BlockLoraSelect: + @classmethod + def INPUT_TYPES(s): + arg_dict = {} + argument = ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10000.0, "step": 0.01}) + + for i in range(48): + arg_dict["blocks.{}.".format(i)] = argument + + return {"required": arg_dict} + + RETURN_TYPES = ("SELECTEDDITBLOCKS", ) + RETURN_NAMES = ("blocks", ) + OUTPUT_TOOLTIPS = ("The modified diffusion model.",) + FUNCTION = "load_lora" + + CATEGORY = "KJNodes/ltxv" + DESCRIPTION = "Select individual block alpha values, value of 0 removes the block altogether" + + def load_lora(self, **kwargs): + return (kwargs,) + + +class DiTBlockLoraLoader: + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}), + "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}), + + }, + "optional": { + "lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}), + "opt_lora_path": ("STRING", {"forceInput": True, "tooltip": "Absolute path of the LoRA."}), + "blocks": ("SELECTEDDITBLOCKS",), + } + } + + RETURN_TYPES = ("MODEL", "STRING", ) + RETURN_NAMES = ("model", "rank", ) + OUTPUT_TOOLTIPS = ("The modified diffusion model.", "possible rank of the LoRA.") + FUNCTION = "load_lora" + CATEGORY = "KJNodes/lora" + + def load_lora(self, model, strength_model, lora_name=None, opt_lora_path=None, blocks=None): + + import comfy.lora + + if opt_lora_path: + lora_path = opt_lora_path + else: + lora_path = folder_paths.get_full_path("loras", lora_name) + + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + self.loaded_lora = None + + if lora is None: + lora = load_torch_file(lora_path, safe_load=True) + self.loaded_lora = (lora_path, lora) + + # Find the first key that ends with "weight" + rank = "unknown" + weight_key = next((key for key in lora.keys() if key.endswith('weight')), None) + # Print the shape of the value corresponding to the key + if weight_key: + logging.info(f"Shape of the first 'weight' key ({weight_key}): {lora[weight_key].shape}") + rank = str(lora[weight_key].shape[0]) + else: + logging.warning("No key ending with 'weight' found.") + rank = "Couldn't find rank" + self.loaded_lora = (lora_path, lora) + + key_map = {} + if model is not None: + key_map = comfy.lora.model_lora_keys_unet(model.model, key_map) + + loaded = comfy.lora.load_lora(lora, key_map) + + if blocks is not None: + keys_to_delete = [] + + for block in blocks: + for key in list(loaded.keys()): + match = False + if isinstance(key, str) and block in key: + match = True + elif isinstance(key, tuple): + for k in key: + if block in k: + match = True + break + + if match: + ratio = blocks[block] + if ratio == 0: + keys_to_delete.append(key) + else: + # Only modify LoRA adapters, skip diff tuples + value = loaded[key] + if hasattr(value, 'weights'): + logging.info(f"Modifying LoRA adapter for key: {key}") + weights_list = list(value.weights) + weights_list[2] = ratio + loaded[key].weights = tuple(weights_list) + else: + logging.info(f"Skipping non-LoRA entry for key: {key}") + + for key in keys_to_delete: + del loaded[key] + + logging.info("loading lora keys:") + for key, value in loaded.items(): + if hasattr(value, 'weights'): + logging.info(f"Key: {key}, Alpha: {value.weights[2]}") + else: + logging.info(f"Key: {key}, Type: {type(value)}") + + if model is not None: + new_modelpatcher = model.clone() + k = new_modelpatcher.add_patches(loaded, strength_model) + + k = set(k) + for x in loaded: + if (x not in k): + logging.warning(f"NOT LOADED {x}") + + return (new_modelpatcher, rank) + +class CustomControlNetWeightsFluxFromList: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "list_of_floats": ("FLOAT", {"forceInput": True}, ), + }, + "optional": { + "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), + "autosize": ("ACNAUTOSIZE", {"padding": 0}), + } + } + + RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",) + RETURN_NAMES = ("CN_WEIGHTS", "TK_SHORTCUT") + FUNCTION = "load_weights" + DESCRIPTION = "Creates controlnet weights from a list of floats for Advanced-ControlNet" + + CATEGORY = "KJNodes/controlnet" + + def load_weights(self, list_of_floats: list[float], + uncond_multiplier: float=1.0, cn_extras: dict[str]={}): + + adv_control = importlib.import_module("ComfyUI-Advanced-ControlNet.adv_control") + ControlWeights = adv_control.utils.ControlWeights + TimestepKeyframeGroup = adv_control.utils.TimestepKeyframeGroup + TimestepKeyframe = adv_control.utils.TimestepKeyframe + + weights = ControlWeights.controlnet(weights_input=list_of_floats, uncond_multiplier=uncond_multiplier, extras=cn_extras) + logging.info(weights.weights_input) + return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) + +SHAKKERLABS_UNION_CONTROLNET_TYPES = { + "canny": 0, + "tile": 1, + "depth": 2, + "blur": 3, + "pose": 4, + "gray": 5, + "low quality": 6, +} + +class SetShakkerLabsUnionControlNetType: + @classmethod + def INPUT_TYPES(s): + return {"required": {"control_net": ("CONTROL_NET", ), + "type": (["auto"] + list(SHAKKERLABS_UNION_CONTROLNET_TYPES.keys()),) + }} + + CATEGORY = "conditioning/controlnet" + RETURN_TYPES = ("CONTROL_NET",) + + FUNCTION = "set_controlnet_type" + + def set_controlnet_type(self, control_net, type): + control_net = control_net.copy() + type_number = SHAKKERLABS_UNION_CONTROLNET_TYPES.get(type, -1) + if type_number >= 0: + control_net.set_extra_arg("control_type", [type_number]) + else: + control_net.set_extra_arg("control_type", []) + + return (control_net,) + +class ModelSaveKJ: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + "filename_prefix": ("STRING", {"default": "diffusion_models/ComfyUI"}), + "model_key_prefix": ("STRING", {"default": "model.diffusion_model."}), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} + RETURN_TYPES = () + FUNCTION = "save" + OUTPUT_NODE = True + + CATEGORY = "advanced/model_merging" + + def save(self, model, filename_prefix, model_key_prefix, prompt=None, extra_pnginfo=None): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + + output_checkpoint = f"{filename}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(full_output_folder, output_checkpoint) + + load_models = [model] + + model_management.load_models_gpu(load_models) + default_prefix = "model.diffusion_model." + + sd = model.state_dict_for_saving(None, None, None) + + new_sd = {} + for k in sd: + if k.startswith(default_prefix): + new_key = model_key_prefix + k[len(default_prefix):] + else: + new_key = k # In case the key doesn't start with the default prefix, keep it unchanged + t = sd[k] + if not t.is_contiguous(): + t = t.contiguous() + new_sd[new_key] = t + logging.info(f"full_output_folder: {full_output_folder}") + if not os.path.exists(full_output_folder): + os.makedirs(full_output_folder) + save_torch_file(new_sd, os.path.join(full_output_folder, output_checkpoint)) + return {} + +class StyleModelApplyAdvanced: + @classmethod + def INPUT_TYPES(s): + return {"required": {"conditioning": ("CONDITIONING", ), + "style_model": ("STYLE_MODEL", ), + "clip_vision_output": ("CLIP_VISION_OUTPUT", ), + "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}), + }} + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "apply_stylemodel" + CATEGORY = "KJNodes/experimental" + DESCRIPTION = "StyleModelApply but with strength parameter" + + def apply_stylemodel(self, clip_vision_output, style_model, conditioning, strength=1.0): + cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0) + cond = strength * cond + c = [] + for t in conditioning: + n = [torch.cat((t[0], cond), dim=1), t[1].copy()] + c.append(n) + return (c, ) + +class AudioConcatenate: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "audio1": ("AUDIO",), + "audio2": ("AUDIO",), + "direction": ( + [ 'right', + 'left', + ], + { + "default": 'right' + }), + }} + + RETURN_TYPES = ("AUDIO",) + FUNCTION = "concanate" + CATEGORY = "KJNodes/audio" + DESCRIPTION = """ +Concatenates the audio1 to audio2 in the specified direction. +""" + + def concanate(self, audio1, audio2, direction): + sample_rate_1 = audio1["sample_rate"] + sample_rate_2 = audio2["sample_rate"] + if sample_rate_1 != sample_rate_2: + raise ValueError("Sample rates of the two audios do not match") + + waveform_1 = audio1["waveform"] + waveform_2 = audio2["waveform"] + + # Concatenate based on the specified direction + if direction == 'right': + concatenated_audio = torch.cat((waveform_1, waveform_2), dim=2) # Concatenate along width + elif direction == 'left': + concatenated_audio= torch.cat((waveform_2, waveform_1), dim=2) # Concatenate along width + return ({"waveform": concatenated_audio, "sample_rate": sample_rate_1},) + +class LeapfusionHunyuanI2V: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT",), + "index": ("INT", {"default": 0, "min": -1, "max": 1000, "step": 1,"tooltip": "The index of the latent to be replaced. 0 for first frame and -1 for last"}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The start percentage of steps to apply"}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The end percentage of steps to apply"}), + "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "KJNodes/hunyuanvideo" + + def patch(self, model, latent, index, strength, start_percent, end_percent): + + def outer_wrapper(samples, index, start_percent, end_percent): + def unet_wrapper(apply_model, args): + steps = args["c"]["transformer_options"]["sample_sigmas"] + inp, timestep, c = args["input"], args["timestep"], args["c"] + matched_step_index = (steps == timestep).nonzero() + if len(matched_step_index) > 0: + current_step_index = matched_step_index.item() + else: + for i in range(len(steps) - 1): + # walk from beginning of steps until crossing the timestep + if (steps[i] - timestep[0]) * (steps[i + 1] - timestep[0]) <= 0: + current_step_index = i + break + else: + current_step_index = 0 + current_percent = current_step_index / (len(steps) - 1) + if samples is not None: + if start_percent <= current_percent <= end_percent: + inp[:, :, [index], :, :] = samples[:, :, [0], :, :].to(inp) + else: + inp[:, :, [index], :, :] = torch.zeros(1) + return apply_model(inp, timestep, **c) + return unet_wrapper + + samples = latent["samples"] * 0.476986 * strength + m = model.clone() + m.set_model_unet_function_wrapper(outer_wrapper(samples, index, start_percent, end_percent)) + + return (m,) + +class ImageNoiseAugmentation: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "noise_aug_strength": ("FLOAT", {"default": None, "min": 0.0, "max": 100.0, "step": 0.001}), + "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "add_noise" + CATEGORY = "KJNodes/image" + DESCRIPTION = """ + Add noise to an image. + """ + + def add_noise(self, image, noise_aug_strength, seed): + torch.manual_seed(seed) + sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * noise_aug_strength + image_noise = torch.randn_like(image) * sigma[:, None, None, None] + image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise) + image_out = image + image_noise + return image_out, + +class VAELoaderKJ: + video_taes = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5"] + image_taes = ["taesd", "taesdxl", "taesd3", "taef1"] + @staticmethod + def vae_list(s): + vaes = folder_paths.get_filename_list("vae") + approx_vaes = folder_paths.get_filename_list("vae_approx") + sdxl_taesd_enc = False + sdxl_taesd_dec = False + sd1_taesd_enc = False + sd1_taesd_dec = False + sd3_taesd_enc = False + sd3_taesd_dec = False + f1_taesd_enc = False + f1_taesd_dec = False + + for v in approx_vaes: + if v.startswith("taesd_decoder."): + sd1_taesd_dec = True + elif v.startswith("taesd_encoder."): + sd1_taesd_enc = True + elif v.startswith("taesdxl_decoder."): + sdxl_taesd_dec = True + elif v.startswith("taesdxl_encoder."): + sdxl_taesd_enc = True + elif v.startswith("taesd3_decoder."): + sd3_taesd_dec = True + elif v.startswith("taesd3_encoder."): + sd3_taesd_enc = True + elif v.startswith("taef1_encoder."): + f1_taesd_dec = True + elif v.startswith("taef1_decoder."): + f1_taesd_enc = True + else: + for tae in s.video_taes: + if v.startswith(tae): + vaes.append(v) + + if sd1_taesd_dec and sd1_taesd_enc: + vaes.append("taesd") + if sdxl_taesd_dec and sdxl_taesd_enc: + vaes.append("taesdxl") + if sd3_taesd_dec and sd3_taesd_enc: + vaes.append("taesd3") + if f1_taesd_dec and f1_taesd_enc: + vaes.append("taef1") + vaes.append("pixel_space") + return vaes + + @staticmethod + def load_taesd(name): + sd = {} + approx_vaes = folder_paths.get_filename_list("vae_approx") + + encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes)) + decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes)) + + enc = load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder)) + for k in enc: + sd["taesd_encoder.{}".format(k)] = enc[k] + + dec = load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder)) + for k in dec: + sd["taesd_decoder.{}".format(k)] = dec[k] + + if name == "taesd": + sd["vae_scale"] = torch.tensor(0.18215) + sd["vae_shift"] = torch.tensor(0.0) + elif name == "taesdxl": + sd["vae_scale"] = torch.tensor(0.13025) + sd["vae_shift"] = torch.tensor(0.0) + elif name == "taesd3": + sd["vae_scale"] = torch.tensor(1.5305) + sd["vae_shift"] = torch.tensor(0.0609) + elif name == "taef1": + sd["vae_scale"] = torch.tensor(0.3611) + sd["vae_shift"] = torch.tensor(0.1159) + return sd + + @classmethod + def INPUT_TYPES(s): + return { + "required": { "vae_name": (s.vae_list(s), ), + "device": (["main_device", "cpu"],), + "weight_dtype": (["bf16", "fp16", "fp32" ],), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "load_vae" + CATEGORY = "KJNodes/vae" + + def load_vae(self, vae_name, device, weight_dtype): + from comfy.sd import VAE + metadata = None + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[weight_dtype] + if device == "main_device": + device = model_management.get_torch_device() + elif device == "cpu": + device = torch.device("cpu") + + if vae_name == "pixel_space": + sd = {} + sd["pixel_space_vae"] = torch.tensor(1.0) + elif vae_name in self.image_taes: + sd = self.load_taesd(vae_name) + else: + if os.path.splitext(vae_name)[0] in self.video_taes: + vae_path = folder_paths.get_full_path_or_raise("vae_approx", vae_name) + else: + vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) + sd, metadata = load_torch_file(vae_path, return_metadata=True) + + + is_audio_vae = ( + "vocoder.conv_post.weight" in sd + or "vocoder.vocoder.conv_post.weight" in sd + or "vocoder.resblocks.0.convs1.0.weight" in sd + or "vocoder.vocoder.resblocks.0.convs1.0.weight" in sd + ) + if is_audio_vae: + sd_audio = state_dict_prefix_replace(dict(sd), {"audio_vae.": "autoencoder.", "vocoder.": "vocoder."}, filter_keys=True) + vae = VAE(sd=sd_audio, metadata=metadata) + else: + vae = VAE(sd=sd, device=device, dtype=dtype, metadata=metadata) + vae.throw_exception_if_invalid() + return (vae,) + +from comfy.samplers import sampling_function, CFGGuider +class Guider_ScheduledCFG(CFGGuider): + + def set_cfg(self, cfg, start_percent, end_percent): + self.cfg = cfg + self.start_percent = start_percent + self.end_percent = end_percent + + def predict_noise(self, x, timestep, model_options={}, seed=None): + steps = model_options["transformer_options"]["sample_sigmas"] + if isinstance(timestep, torch.Tensor): + timestep_value = timestep.reshape(-1)[0].to(steps) + else: + timestep_value = torch.tensor(timestep, device=steps.device, dtype=steps.dtype) + matched_step_index = torch.isclose(steps, timestep_value).nonzero() + assert not (isinstance(self.cfg, list) and len(self.cfg) != (len(steps) - 1)), "cfg list length must match step count" + if len(matched_step_index) > 0: + current_step_index = matched_step_index.item() + else: + for i in range(len(steps) - 1): + # walk from beginning of steps until crossing the timestep + if (steps[i] - timestep_value) * (steps[i + 1] - timestep_value) <= 0: + current_step_index = i + break + else: + current_step_index = 0 + current_percent = current_step_index / (len(steps) - 1) + + if self.start_percent <= current_percent <= self.end_percent: + if isinstance(self.cfg, list): + cfg = self.cfg[current_step_index] + else: + cfg = self.cfg + uncond = self.conds.get("negative", None) + else: + uncond = None + cfg = 1.0 + + return sampling_function(self.inner_model, x, timestep, uncond, self.conds.get("positive", None), cfg, model_options=model_options, seed=seed) + +class ScheduledCFGGuidance: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 100.0, "step": 0.01}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step":0.01}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.01}), + }, + } + RETURN_TYPES = ("GUIDER",) + FUNCTION = "get_guider" + CATEGORY = "KJNodes/experimental" + DESCRiPTION = """ +CFG Guider that allows for scheduled CFG changes over steps, the steps outside the range will use CFG 1.0 thus being processed faster. +cfg input can be a list of floats matching step count, or a single float for all steps. +""" + + def get_guider(self, model, cfg, positive, negative, start_percent, end_percent): + guider = Guider_ScheduledCFG(model) + guider.set_conds(positive, negative) + guider.set_cfg(cfg, start_percent, end_percent) + return (guider, ) + + +class ApplyRifleXRoPE_WanVideo: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT", {"tooltip": "Only used to get the latent count"}), + "k": ("INT", {"default": 6, "min": 1, "max": 100, "step": 1, "tooltip": "Index of intrinsic frequency"}), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "KJNodes/wan" + EXPERIMENTAL = True + DESCRIPTION = "Extends the potential frame count of HunyuanVideo using this method: https://github.com/thu-ml/RIFLEx" + + def patch(self, model, latent, k): + model_class = model.model.diffusion_model + + model_clone = model.clone() + num_frames = latent["samples"].shape[2] + d = model_class.dim // model_class.num_heads + + rope_embedder = EmbedND_RifleX( + d, + 10000.0, + [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)], + num_frames, + k + ) + + model_clone.add_object_patch(f"diffusion_model.rope_embedder", rope_embedder) + + return (model_clone, ) + +class ApplyRifleXRoPE_HunuyanVideo: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "latent": ("LATENT", {"tooltip": "Only used to get the latent count"}), + "k": ("INT", {"default": 4, "min": 1, "max": 100, "step": 1, "tooltip": "Index of intrinsic frequency"}), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "KJNodes/hunyuanvideo" + EXPERIMENTAL = True + DESCRIPTION = "Extends the potential frame count of HunyuanVideo using this method: https://github.com/thu-ml/RIFLEx" + + def patch(self, model, latent, k): + model_class = model.model.diffusion_model + + model_clone = model.clone() + num_frames = latent["samples"].shape[2] + + pe_embedder = EmbedND_RifleX( + model_class.params.hidden_size // model_class.params.num_heads, + model_class.params.theta, + model_class.params.axes_dim, + num_frames, + k + ) + + model_clone.add_object_patch(f"diffusion_model.pe_embedder", pe_embedder) + + return (model_clone, ) + +def rope_riflex(pos, dim, theta, L_test, k): + from einops import rearrange + assert dim % 2 == 0 + if model_management.is_device_mps(pos.device) or model_management.is_intel_xpu() or model_management.is_directml_enabled(): + device = torch.device("cpu") + else: + device = pos.device + + scale = torch.linspace(0, (dim - 2) / dim, steps=dim//2, dtype=torch.float64, device=device) + omega = 1.0 / (theta**scale) + + # RIFLEX modification - adjust last frequency component if L_test and k are provided + if k and L_test: + omega[k-1] = 0.9 * 2 * torch.pi / L_test + + out = torch.einsum("...n,d->...nd", pos.to(dtype=torch.float32, device=device), omega) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2) + return out.to(dtype=torch.float32, device=pos.device) + +class EmbedND_RifleX(nn.Module): + def __init__(self, dim, theta, axes_dim, num_frames, k): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + self.num_frames = num_frames + self.k = k + + def forward(self, ids): + n_axes = ids.shape[-1] + emb = torch.cat( + [rope_riflex(ids[..., i], self.axes_dim[i], self.theta, self.num_frames, self.k if i == 0 else 0) for i in range(n_axes)], + dim=-3, + ) + return emb.unsqueeze(1) + + +class Timer: + def __init__(self, name): + self.name = name + self.start_time = None + self.elapsed = 0 + +class TimerNodeKJ: + @classmethod + + def INPUT_TYPES(s): + return { + "required": { + "any_input": (IO.ANY, ), + "mode": (["start", "stop"],), + "name": ("STRING", {"default": "Timer"}), + }, + "optional": { + "timer": ("TIMER",), + }, + } + + RETURN_TYPES = (IO.ANY, "TIMER", "INT", ) + RETURN_NAMES = ("any_output", "timer", "time") + FUNCTION = "timer" + CATEGORY = "KJNodes/misc" + + def timer(self, mode, name, any_input=None, timer=None): + if timer is None: + if mode == "start": + timer = Timer(name=name) + timer.start_time = time.time() + return {"ui": { + "text": [f"{timer.start_time}"]}, + "result": (any_input, timer, 0) + } + elif mode == "stop" and timer is not None: + end_time = time.time() + timer.elapsed = int((end_time - timer.start_time) * 1000) + timer.start_time = None + return (any_input, timer, timer.elapsed) + +class HunyuanVideoEncodeKeyframesToCond: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "positive": ("CONDITIONING", ), + "vae": ("VAE", ), + "start_frame": ("IMAGE", ), + "end_frame": ("IMAGE", ), + "num_frames": ("INT", {"default": 33, "min": 2, "max": 4096, "step": 1}), + "tile_size": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}), + "temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to encode at a time."}), + "temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Only used for video VAEs: Amount of frames to overlap."}), + }, + "optional": { + "negative": ("CONDITIONING", ), + } + } + + RETURN_TYPES = ("MODEL", "CONDITIONING","CONDITIONING","LATENT") + RETURN_NAMES = ("model", "positive", "negative", "latent") + FUNCTION = "encode" + + CATEGORY = "KJNodes/hunyuanvideo" + + def encode(self, model, positive, start_frame, end_frame, num_frames, vae, tile_size, overlap, temporal_size, temporal_overlap, negative=None): + + model_clone = model.clone() + + model_clone.add_object_patch("concat_keys", ("concat_image",)) + + + x = (start_frame.shape[1] // 8) * 8 + y = (start_frame.shape[2] // 8) * 8 + + if start_frame.shape[1] != x or start_frame.shape[2] != y: + x_offset = (start_frame.shape[1] % 8) // 2 + y_offset = (start_frame.shape[2] % 8) // 2 + start_frame = start_frame[:,x_offset:x + x_offset, y_offset:y + y_offset,:] + if end_frame.shape[1] != x or end_frame.shape[2] != y: + x_offset = (start_frame.shape[1] % 8) // 2 + y_offset = (start_frame.shape[2] % 8) // 2 + end_frame = end_frame[:,x_offset:x + x_offset, y_offset:y + y_offset,:] + + video_frames = torch.zeros(num_frames-2, start_frame.shape[1], start_frame.shape[2], start_frame.shape[3], device=start_frame.device, dtype=start_frame.dtype) + video_frames = torch.cat([start_frame, video_frames, end_frame], dim=0) + + concat_latent = vae.encode_tiled(video_frames[:,:,:,:3], tile_x=tile_size, tile_y=tile_size, overlap=overlap, tile_t=temporal_size, overlap_t=temporal_overlap) + + out_latent = {} + out_latent["samples"] = torch.zeros_like(concat_latent) + + out = [] + for conditioning in [positive, negative if negative is not None else []]: + c = [] + for t in conditioning: + d = t[1].copy() + d["concat_latent_image"] = concat_latent + n = [t[0], d] + c.append(n) + out.append(c) + if len(out) == 1: + out.append(out[0]) + return (model_clone, out[0], out[1], out_latent) + + +class LazySwitchKJ: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "switch": ("BOOLEAN",), + "on_false": (IO.ANY, {"lazy": True}), + "on_true": (IO.ANY, {"lazy": True}), + }, + } + + RETURN_TYPES = (IO.ANY,) + FUNCTION = "switch" + CATEGORY = "KJNodes/misc" + DESCRIPTION = "Controls flow of execution based on a boolean switch." + + def check_lazy_status(self, switch, on_false=None, on_true=None): + if switch and on_true is None: + return ["on_true"] + if not switch and on_false is None: + return ["on_false"] + + def switch(self, switch, on_false = None, on_true=None): + value = on_true if switch else on_false + return (value,) + + +from comfy.patcher_extension import WrappersMP +from comfy.sampler_helpers import prepare_mask +class TTM_OuterSampleWrapper: + def __init__(self, mask, steps): + self.mask = mask + self.steps = steps + + def __call__(self, executor, noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes): + guider = executor.class_obj + guider.model_options + wrappers = guider.model_options["transformer_options"]["wrappers"] + w = wrappers.setdefault(WrappersMP.APPLY_MODEL, {}) + + if self.mask is not None: + motion_mask = self.mask.reshape((-1, 1, self.mask.shape[-2], self.mask.shape[-1])) + shape = latent_shapes[0] + motion_mask = prepare_mask(motion_mask, shape, noise.device) + + scale_latent_inpaint = guider.model_patcher.model.scale_latent_inpaint + w["TTM_ApplyModel_Wrapper"] = [TTM_ApplyModel_Wrapper(latent_image, noise, motion_mask, self.steps, scale_latent_inpaint)] + + out = executor(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes) + + return out + + +class TTM_ApplyModel_Wrapper: + def __init__(self, reference_samples, noise, motion_mask, steps, scale_latent_inpaint): + self.reference_samples = reference_samples + self.noise = noise + self.motion_mask = motion_mask + self.steps = steps + self.scale_latent_inpaint = scale_latent_inpaint + + def __call__(self, executor, x, t, c_concat, c_crossattn, control, transformer_options, **kwargs): + sigmas = transformer_options["sample_sigmas"] + + matched = (sigmas == t).nonzero(as_tuple=True)[0] + if matched.numel() > 0: + current_step_index = matched.item() + else: + crossing = ((sigmas[:-1] - t) * (sigmas[1:] - t) <= 0).nonzero(as_tuple=True)[0] + current_step_index = crossing.item() if crossing.numel() > 0 else 0 + + next_sigma = sigmas[current_step_index + 1] if current_step_index < len(sigmas) - 1 else sigmas[current_step_index] + + if current_step_index != 0 and current_step_index < self.steps: + noisy_latent = self.scale_latent_inpaint(x=x, sigma=torch.tensor([next_sigma]), noise=self.noise.to(x), latent_image=self.reference_samples.to(x)) + if self.motion_mask is not None: + x = x * (1-self.motion_mask).to(x) + noisy_latent * self.motion_mask.to(x) + else: + x = noisy_latent + + return executor(x, t, c_concat, c_crossattn, control, transformer_options, **kwargs) + + +class LatentInpaintTTM: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL", ), + "steps": ("INT", {"default": 7, "min": 0, "max": 888, "step": 1, "tooltip": "Number of steps to apply TTM inpainting for."}), + }, + "optional": { + "mask": ("MASK", {"tooltip": "Latent mask where white (1.0) is the area to inpaint and black (0.0) is the area to keep unchanged."}), + } + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + EXPERIMENTAL = True + DESCRIPTION = "https://github.com/time-to-move/TTM" + SEARCH_ALIASES = ["time to move"] + CATEGORY = "KJNodes/experimental" + + def patch(self, model, steps, mask=None): + m = model.clone() + m.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, "TTM_OuterSampleWrapper", TTM_OuterSampleWrapper(mask, steps)) + return (m, ) + + +class SimpleCalculatorKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + template = io.Autogrow.TemplateNames(input=io.MultiType.Input("var", [io.Int, io.Float, io.Boolean], optional=True), names=["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"], min=2) + return io.Schema( + node_id="SimpleCalculatorKJ", + category="KJNodes/misc", + description=""" +Calculator node that evaluates a mathematical expression using inputs a and b. + Supported operations: +, -, *, /, //, %, **, <<, >>, unary +/- + Supported comparisons: ==, !=, <, <=, >, >= + Supported logic: and, or, not + Supported functions: abs(), round(), min(), max(), pow(), sqrt(), sin(), cos(), tan(), log(), log10(), exp(), floor(), ceil() + Supported constants: pi, euler, True, False +""", + search_aliases=["math", "arithmetic", "expression", "logic"], + inputs=[ + io.String.Input("expression", default="a + b", multiline=True), + io.Autogrow.Input("variables", template=template), + ], + outputs=[ + io.Float.Output(), + io.Int.Output(), + io.Boolean.Output(), + ], + ) + + @classmethod + def execute(cls, variables, expression, a=None, b=None) -> io.NodeOutput: + import ast + import operator + + # Allowed operations + allowed_operators = { + ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: operator.pow, + ast.USub: operator.neg, ast.UAdd: operator.pos, ast.LShift: operator.lshift, + ast.RShift: operator.rshift, ast.Eq: operator.eq, ast.NotEq: operator.ne, ast.Lt: operator.lt, + ast.LtE: operator.le, ast.Gt: operator.gt, ast.GtE: operator.ge, ast.And: operator.and_, + ast.Or: operator.or_, ast.Not: operator.not_, + } + + # Allowed functions + allowed_functions = { + 'abs': abs, 'round': round, 'min': min, 'max': max, + 'pow': pow, 'sqrt': math.sqrt, 'sin': math.sin, + 'cos': math.cos, 'tan': math.tan, 'log': math.log, + 'log10': math.log10, 'exp': math.exp, 'floor': math.floor, + 'ceil': math.ceil + } + + # Allowed constants - start with pi, e, True, False + allowed_names = {'pi': math.pi, 'euler': math.e, 'True': True, 'False': False} + + # Add all variables from autogrow to allowed_names + for var_name, var_value in variables.items(): + allowed_names[var_name] = var_value + + # Backwards compatibility: add a and b if they're provided (for old workflows) + if a is not None: + allowed_names['a'] = a + if b is not None: + allowed_names['b'] = b + + def eval_node(node): + if isinstance(node, ast.Constant): # Numbers and booleans + return node.value + elif isinstance(node, ast.Name): # Variables + if node.id in allowed_names: + return allowed_names[node.id] + raise ValueError(f"Name '{node.id}' is not allowed") + elif isinstance(node, ast.BinOp): # Binary operations + if type(node.op) not in allowed_operators: + raise ValueError(f"Operator {type(node.op).__name__} is not allowed") + left = eval_node(node.left) + right = eval_node(node.right) + return allowed_operators[type(node.op)](left, right) + elif isinstance(node, ast.UnaryOp): # Unary operations + if type(node.op) not in allowed_operators: + raise ValueError(f"Operator {type(node.op).__name__} is not allowed") + operand = eval_node(node.operand) + return allowed_operators[type(node.op)](operand) + elif isinstance(node, ast.Compare): # Comparison operations + left = eval_node(node.left) + for op, comparator in zip(node.ops, node.comparators): + if type(op) not in allowed_operators: + raise ValueError(f"Operator {type(op).__name__} is not allowed") + right = eval_node(comparator) + result = allowed_operators[type(op)](left, right) + if not result: + return False + left = right + return True + elif isinstance(node, ast.BoolOp): # Boolean operations (and, or) + if type(node.op) not in allowed_operators: + raise ValueError(f"Operator {type(node.op).__name__} is not allowed") + values = [eval_node(value) for value in node.values] + if isinstance(node.op, ast.And): + return all(values) + elif isinstance(node.op, ast.Or): + return any(values) + elif isinstance(node, ast.Call): # Function calls + if not isinstance(node.func, ast.Name): + raise ValueError("Only simple function calls are allowed") + if node.func.id not in allowed_functions: + raise ValueError(f"Function '{node.func.id}' is not allowed") + args = [eval_node(arg) for arg in node.args] + return allowed_functions[node.func.id](*args) + else: + raise ValueError(f"Node type {type(node).__name__} is not allowed") + + try: + tree = ast.parse(expression, mode='eval') + result = eval_node(tree.body) + return io.NodeOutput(float(result), int(result), bool(result)) + except Exception as e: + logging.error(f"CalculatorKJ Error: {str(e)}") + return io.NodeOutput(0.0, 0, False) + + +class GetTrackRange(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="GetTrackRange", + category="conditioning/video_models", + inputs=[ + io.Tracks.Input("tracks"), + io.Int.Input("start_index", default=24, min=-10000, max=10000, step=1), + io.Int.Input("num_frames", default=10, min=1, max=10000, step=1), + ], + outputs=[ + io.Tracks.Output(), + ], + ) + + @classmethod + def execute(cls, tracks, start_index, num_frames) -> io.NodeOutput: + track_path = tracks["track_path"] + mask = tracks["track_visibility"] + total_frames = track_path.shape[0] + + if start_index < 0: + start_index = total_frames + start_index + start_index = max(0, min(start_index, total_frames)) + + # Clamp end_index + end_index = max(0, min(start_index + num_frames, total_frames)) + + tracks_out = track_path[start_index:end_index, ...] + mask_out = mask[start_index:end_index, ...] + + out_track = { + "track_path": tracks_out, + "track_visibility": mask_out, + } + return io.NodeOutput(out_track) + +class AddNoiseToTrackPath(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="AddNoiseToTrackPath", + category="conditioning/video_models", + inputs=[ + io.Tracks.Input("tracks"), + io.Float.Input("strength", default=1.0, min=0.0, max=100.0, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, step=1), + io.Float.Input("noise_x_ratio", default=1.0, min=0.0, max=100.0, step=0.01, + tooltip="Multiplier for horizontal noise component"), + io.Float.Input("noise_y_ratio", default=1.0, min=0.0, max=100.0, step=0.01, + tooltip="Multiplier for vertical noise component"), + io.Float.Input("noise_temporal_ratio", default=1.0, min=0.0, max=100.0, step=0.01, + tooltip="Multiplier for temporal (frame-to-frame) noise"), + ], + outputs=[ + io.Tracks.Output(), + ], + ) + + @classmethod + def execute(cls, tracks, strength, seed, noise_x_ratio, noise_y_ratio, noise_temporal_ratio) -> io.NodeOutput: + track_path = tracks["track_path"].clone() + mask = tracks["track_visibility"] + + torch.manual_seed(seed) + noise = torch.randn_like(track_path) * strength + + # Apply directional scaling to noise + noise[..., 0] *= noise_x_ratio # X coordinate noise + noise[..., 1] *= noise_y_ratio # Y coordinate noise + + # Apply temporal smoothing if temporal ratio is less than 1 + if noise_temporal_ratio < 1.0: + num_frames = track_path.shape[0] + smoothed_noise = noise.clone() + kernel_size = max(1, int((1.0 - noise_temporal_ratio) * 10)) + + for i in range(num_frames): + start_idx = max(0, i - kernel_size // 2) + end_idx = min(num_frames, i + kernel_size // 2 + 1) + smoothed_noise[i] = noise[start_idx:end_idx].mean(dim=0) + + noise = smoothed_noise * noise_temporal_ratio + noise * (1 - noise_temporal_ratio) + + track_path = track_path + noise + + out_track = { + "track_path": track_path, + "track_visibility": mask, + } + return io.NodeOutput(out_track) + + +class VAEDecodeLoopKJ: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT", {"tooltip": "The latent to be decoded."}), + "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."}), + "overlap_latent_frames": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1, "tooltip": "Number of frames to blend for seamless loop, for Wan 2 works and HunyuanVideo 1.5 should use 4"}), + } + } + RETURN_TYPES = ("IMAGE",) + OUTPUT_TOOLTIPS = ("The decoded images.",) + FUNCTION = "decode" + CATEGORY = "KJNodes/vae" + DESCRIPTION = "Video latent VAE decoding to fix artifacts on loop seams." + + def decode(self, vae, samples, overlap_latent_frames): + latents = samples["samples"] + + images = vae.decode(latents) + if overlap_latent_frames <= 0: + if len(images.shape) == 5: + images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + return (images, ) + + end_frames = overlap_latent_frames + 1 + start_frames = overlap_latent_frames + + temp_images = vae.decode(torch.cat([latents[:, :, -end_frames:]] + [latents[:, :, :start_frames]], dim=2)).cpu().float() + + total_concat = end_frames + start_frames + temp_start = total_concat * 2 - 1 + main_start = total_concat + (overlap_latent_frames if overlap_latent_frames > 2 else 0) + + images = torch.cat([temp_images[:, temp_start:].to(images), images[:, main_start:]], dim=1) + if len(images.shape) == 5: + images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + + return (images, ) + +class WanImageToVideoSVIPro(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WanImageToVideoSVIPro", + category="conditioning/video_models", + inputs=[ + io.Conditioning.Input("positive"), + io.Conditioning.Input("negative"), + io.Int.Input("length", default=81, min=1, max=MAX_RESOLUTION, step=4), + io.Latent.Input("anchor_samples"), + io.Latent.Input("prev_samples", optional=True), + io.Int.Input("motion_latent_count", default=1, min=0, max=128, step=1), + ], + outputs=[ + io.Conditioning.Output(display_name="positive"), + io.Conditioning.Output(display_name="negative"), + io.Latent.Output(display_name="latent"), + ], + ) + + @classmethod + def execute(cls, positive, negative, length, motion_latent_count, anchor_samples, prev_samples=None) -> io.NodeOutput: + anchor_latent = anchor_samples["samples"].clone() + + B, C, T, H, W = anchor_latent.shape + empty_latent = torch.zeros([B, 16, ((length - 1) // 4) + 1, H, W], device=model_management.intermediate_device()) + + total_latents = (length - 1) // 4 + 1 + device = anchor_latent.device + dtype = anchor_latent.dtype + + if prev_samples is None or motion_latent_count == 0: + padding_size = total_latents - anchor_latent.shape[2] + image_cond_latent = anchor_latent + else: + motion_latent = prev_samples["samples"][:, :, -motion_latent_count:].clone() + padding_size = total_latents - anchor_latent.shape[2] - motion_latent.shape[2] + image_cond_latent = torch.cat([anchor_latent, motion_latent], dim=2) + + padding = torch.zeros(1, C, padding_size, H, W, dtype=dtype, device=device) + padding = comfy.latent_formats.Wan21().process_out(padding) + image_cond_latent = torch.cat([image_cond_latent, padding], dim=2) + + mask = torch.ones((1, 1, empty_latent.shape[2], H, W), device=device, dtype=dtype) + mask[:, :, :1] = 0.0 + + positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": image_cond_latent, "concat_mask": mask}) + negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": image_cond_latent, "concat_mask": mask}) + + out_latent = {} + out_latent["samples"] = empty_latent + return io.NodeOutput(positive, negative, out_latent) + +class DeprecatedCompileNodeKJ: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": (IO.ANY,), + }, + } + RETURN_TYPES = (IO.ANY,) + FUNCTION = "passthrough" + CATEGORY = "KJNodes/deprecated" + DESCRIPTION = "This node has been replaced with TorchCompileModelAdvanced node, please use that instead." + def passthrough(self, model): + return (model,) + + +class VisualizeSigmasKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="VisualizeSigmasKJ", + category="KJNodes/misc", + inputs=[ + io.Sigmas.Input("sigmas"), + io.Int.Input("start_step", default=0, min=-1, max=1000, step=1, + tooltip="Step index to mark as the start of a range (inclusive). Set to -1 to disable."), + io.Int.Input("end_step", default=-1, min=-1, max=1000, step=1, + tooltip="Step index to mark as the end of a range (inclusive). Set to - 1 to disable."), + ], + outputs=[ + io.Sigmas.Output(display_name="sigmas_out"), + io.Image.Output(display_name="image"), + ], + ) + + @classmethod + def execute(cls, sigmas, start_step=0, end_step=-1) -> io.NodeOutput: + + start_idx = 0 + end_idx = len(sigmas) - 1 + + if isinstance(start_step, float): + idxs = (sigmas <= start_step).nonzero(as_tuple=True)[0] + if len(idxs) > 0: + start_idx = idxs[0].item() + elif isinstance(start_step, int): + if start_step > 0: + start_idx = start_step + + if isinstance(end_step, float): + idxs = (sigmas >= end_step).nonzero(as_tuple=True)[0] + if len(idxs) > 0: + end_idx = idxs[-1].item() + elif isinstance(end_step, int): + if end_step != -1: + end_idx = end_step - 1 + + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + sigmas_np = sigmas.cpu().numpy() + if not np.isclose(sigmas_np[-1], 0.0, atol=1e-6): + sigmas_np = np.append(sigmas_np, 0.0) + buf = BytesIO() + fig = plt.figure(facecolor='#353535') + ax = fig.add_subplot(111) + ax.set_facecolor('#353535') # Set axes background color + x_values = range(0, len(sigmas_np)) + ax.plot(x_values, sigmas_np) + # Annotate each sigma value + ax.scatter(x_values, sigmas_np, color='white', s=20, zorder=3) # Small dots at each sigma + for x, y in zip(x_values, sigmas_np): + # Show all annotations if few steps, or just show split step annotations + show_annotation = len(sigmas_np) <= 10 + is_split_step = (start_idx > 0 and x == start_idx) or (end_idx != -1 and x == end_idx + 1) + + if show_annotation or is_split_step: + color = 'orange' + if is_split_step: + color = 'yellow' + ax.annotate(f"{y:.3f}", (x, y), textcoords="offset points", xytext=(10, 1), ha='center', color=color, fontsize=12) + ax.set_xticks(x_values) + ax.set_title("Sigmas", color='white') # Title font color + ax.set_xlabel("Step", color='white') # X label font color + ax.set_ylabel("Sigma Value", color='white') # Y label font color + ax.tick_params(axis='x', colors='white', labelsize=10) # X tick color + ax.tick_params(axis='y', colors='white', labelsize=10) # Y tick color + # Add split point if end_step is defined + end_idx += 1 + if end_idx != -1 and 0 <= end_idx < len(sigmas_np) - 1: + ax.axvline(end_idx, color='red', linestyle='--', linewidth=2, label='end_step split') + # Add split point if start_step is defined + if start_idx > 0 and 0 <= start_idx < len(sigmas_np): + ax.axvline(start_idx, color='green', linestyle='--', linewidth=2, label='start_step split') + if (end_idx != -1 and 0 <= end_idx < len(sigmas_np)) or (start_idx > 0 and 0 <= start_idx < len(sigmas_np)): + handles, labels = ax.get_legend_handles_labels() + if labels: + ax.legend() + # Draw shaded range + range_start_idx = start_idx if start_idx > 0 else 0 + range_end_idx = end_idx if end_idx > 0 and end_idx < len(sigmas_np) else len(sigmas_np) - 1 + if range_start_idx < range_end_idx: + ax.axvspan(range_start_idx, range_end_idx, color='lightblue', alpha=0.1, label='Sampled Range') + + + plt.tight_layout() + fig.canvas.draw() + w, h = fig.canvas.get_width_height() + try: + buf = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8) + buf = buf.reshape(h, w, 4) + buf = buf[:, :, [1, 2, 3]] # Convert ARGB to RGB + except AttributeError: + buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) + buf = buf.reshape(h, w, 3).copy() + image = torch.from_numpy(buf).float() / 255.0 + image = image.unsqueeze(0) #(H, W, C) -> (1, H, W, C) + plt.close(fig) + + sigmas_out = sigmas[start_idx:end_idx + 1] if end_idx != -1 else sigmas[start_idx:] + + return io.NodeOutput(sigmas_out,image) + +class PreviewLatentNoiseMask(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PreviewLatentNoiseMask", + category="KJNodes/latents", + description="Previews the latent noise mask", + inputs=[ + io.Latent.Input("latent",), + ], + outputs=[ + io.Mask.Output(display_name="mask"), + ], + ) + + @classmethod + def execute(cls, latent) -> io.NodeOutput: + noise_mask = latent.get("noise_mask", None) + if noise_mask is None: + return io.NodeOutput(torch.zeros((1, 64, 64))) + noise_mask = noise_mask.clone() + + if noise_mask.ndim == 5: + noise_mask = noise_mask[0, 0] + + return io.NodeOutput(noise_mask) + +class PlaySoundKJ(io.ComfyNode): + """Plays audio in the browser when execution reaches this node.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PlaySoundKJ", + category="KJNodes/audio", + description="Plays the input audio in the browser. Modes: 'always' plays on every execution, 'on_empty_queue' plays only when the queue finishes, 'on_change' plays only when the audio content changes. Duration limits playback length (0 = full audio).", + inputs=[ + io.AnyType.Input("any_input", optional=True), + io.Audio.Input("audio", optional=True), + io.String.Input("audio_path", default="", tooltip="Path to an audio file. Used when audio input is not connected."), + io.Combo.Input("mode", options=["always", "on_empty_queue", "on_change"], default="always"), + io.Float.Input("volume", default=0.5, min=0.0, max=1.0, step=0.01), + io.Float.Input("duration", default=5.0, min=0.0, max=300.0, step=0.1, tooltip="Duration in seconds to play. 0 = play full audio."), + ], + outputs=[ + io.AnyType.Output("any_output", display_name="any_output"), + ], + is_output_node=True, + ) + + @classmethod + def fingerprint_inputs(cls, **kwargs): + if kwargs.get("mode") == "on_change": + return False + return float("NaN") + + @staticmethod + def _generate_chime(): + sr = 32000 + C, Ab, Bb = 523.26, 421.30, 466.16 # note frequencies + e = 0.14 # eighth note — adjust to change tempo + S, M, H = 16, 7, 2.5 # staccato / medium / held decay + melody = [ + (C,e,S), (C,e,S), (C,e,S), (C,e*3,M), + (Ab,e*3,M), (Bb,e*3,M), (C,e*2,S), (Bb,e,S), (C,e*5,H), + ] + k = torch.exp(-torch.linspace(-2, 2, 15) ** 2) + k = (k / k.sum()).reshape(1, 1, -1) + parts = [] + for freq, dur, decay in melody: + t = torch.linspace(0, dur, int(sr * dur)) + tone = torch.tanh(3 * torch.sin(2 * math.pi * freq * t)) + tone = torch.nn.functional.conv1d(tone.reshape(1, 1, -1), k, padding=7).squeeze() + parts.append(tone * torch.exp(-t * decay)) + wav = torch.cat(parts) * 0.45 + return {"waveform": wav.unsqueeze(0).unsqueeze(0), "sample_rate": sr} + + @classmethod + def execute(cls, audio=None, audio_path="", mode="always", volume=0.5, duration=5.0, any_input=None) -> io.NodeOutput: + if audio is None: + if audio_path: + import av + with av.open(audio_path) as af: + stream = af.streams.audio[0] + sr = stream.codec_context.sample_rate + frames = [] + for frame in af.decode(streams=stream.index): + buf = torch.from_numpy(frame.to_ndarray()) + if buf.shape[0] != stream.channels: + buf = buf.view(-1, stream.channels).t() + frames.append(buf) + wav = torch.cat(frames, dim=1).float() + audio = {"waveform": wav.unsqueeze(0), "sample_rate": sr} + else: + audio = cls._generate_chime() + + preview = ui.PreviewAudio(audio, cls=cls) + ui_dict = preview.as_dict() + ui_dict["audio_hash"] = [hash(audio["waveform"].sum().item())] + return io.NodeOutput( + any_input, + ui=ui_dict, + ) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/preview_override_node.py b/custom_nodes/ComfyUI-KJNodes/nodes/preview_override_node.py new file mode 100644 index 0000000000000000000000000000000000000000..ab0e1e76d7573dd2269dc1323bc1e0bfe2c70c5e --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/preview_override_node.py @@ -0,0 +1,765 @@ +import base64 +import io as pyio +import logging +import queue +import threading +import time + +import numpy as np +import torch + +import comfy.model_management +import comfy.patcher_extension +import latent_preview +from comfy_api.latest import io +from PIL import Image, ImageOps + +try: + from .ltxv_nodes import WrappedPreviewer as _LTXWrappedPreviewer, get_ltx_rgb_factors as _ltx_rgb_factors +except Exception as e: + logging.warning(f"[KJ PreviewOverride] LTX preview helpers unavailable ({e}); LTX previews disabled.") + _LTXWrappedPreviewer = None + _ltx_rgb_factors = None + + +try: + from server import PromptServer +except ImportError: + PromptServer = None + +def _suppressed_preview_image(self_, preview_format, x0): + return None + + +class _AsyncPreviewEncoder: + """Off-thread encoder. Bounded FIFO drops-on-full so the sampler never blocks on us.""" + + _STOP = object() + + def __init__(self, max_in_flight=2): + self.q = queue.Queue(maxsize=max_in_flight) + self.thread = threading.Thread(target=self._run, name="kj_preview_encoder", daemon=True) + self.thread.start() + + def submit(self, fn): + try: + self.q.put_nowait(fn) + return True + except queue.Full: + return False + + def _run(self): + while True: + item = self.q.get() + if item is self._STOP: + return + try: + item() + except Exception: + logging.exception("[KJ Preview Override] async encoder error") + + def shutdown(self, drain_timeout=5.0): + try: + self.q.put(self._STOP, timeout=drain_timeout) + except queue.Full: + pass + self.thread.join(timeout=drain_timeout) + + +def _get_core_previewer(load_device, latent_format): + # Walk past custom-node hooks on get_previewer to reach the unwrapped core function. + fn = latent_preview.get_previewer + seen = set() + while hasattr(fn, "__wrapped__") and id(fn) not in seen: + seen.add(id(fn)) + fn = fn.__wrapped__ + return fn(load_device, latent_format) + + +def _decode_video_frames_l2rgb(x0, latent_format, max_frames, stride=1): + # Bulk-blocking GPU→CPU copy (not per-frame non_blocking) avoids torn frames at high res. + if x0.ndim != 5: + return [] + rgb_factors = getattr(latent_format, "latent_rgb_factors", None) + if rgb_factors is None: + return [] + try: + reshape = getattr(latent_format, "latent_rgb_factors_reshape", None) + if reshape is not None: + x0 = reshape(x0) + bias = getattr(latent_format, "latent_rgb_factors_bias", None) + factors = torch.tensor(rgb_factors, device=x0.device, dtype=x0.dtype).transpose(0, 1) + bias_t = torch.tensor(bias, device=x0.device, dtype=x0.dtype) if bias is not None else None + x = x0[0] + if stride > 1: + x = x[:, ::stride] + t_total = x.shape[1] + if max_frames > 0 and max_frames < t_total: + indices = np.linspace(0, t_total - 1, max_frames).round().astype(int).tolist() + x = x[:, indices] + x = x.movedim(0, -1) + rgb = torch.nn.functional.linear(x, factors, bias=bias_t) + rgb.add_(1.0).mul_(127.5).clamp_(0, 255) + rgb_cpu = rgb.to(torch.uint8).cpu().numpy() + return [Image.fromarray(rgb_cpu[i]) for i in range(rgb_cpu.shape[0])] + except Exception: + return [] + + +# PyPI PyAV wheels typically lack NVENC; probe once at import. +def _probe_nvenc(): + try: + import av # noqa + av.Codec("h264_nvenc", "w") + return True + except Exception: + return False + +_NVENC_AVAILABLE = _probe_nvenc() + +# NVENC H.264 rejects sub-145×49 inputs at avcodec_open2 — fall back to WebP for small frames. +_NVENC_MIN_W = 145 +_NVENC_MIN_H = 49 + +_nvenc_warned = False + + +def _encode_mp4_nvenc(frames, fps, max_res): + # Fragmented MP4 so the browser can decode mid-download. Returns (None, 0, 0) on failure + # (including too-small-for-NVENC), so caller falls through to WebP. + global _nvenc_warned + if not frames: + return None, 0, 0 + try: + import av + except Exception: + return None, 0, 0 + pil_frames = [] + for f in frames: + pf = f if f.mode == "RGB" else f.convert("RGB") + if max_res and max_res > 0 and (pf.width > max_res or pf.height > max_res): + pf = ImageOps.contain(pf, (max_res, max_res), Image.LANCZOS) + pil_frames.append(pf) + # yuv420p requires even dimensions. + w0, h0 = pil_frames[0].width, pil_frames[0].height + out_w, out_h = w0 & ~1, h0 & ~1 + if (out_w, out_h) != (w0, h0): + pil_frames = [pf.resize((out_w, out_h), Image.LANCZOS) for pf in pil_frames] + if out_w < _NVENC_MIN_W or out_h < _NVENC_MIN_H: + return None, 0, 0 + # Driver/GPU varies what option combos are accepted; bare preset always works. + option_candidates = [ + {"preset": "p1", "rc": "vbr", "cq": "23"}, + {"preset": "p1"}, + ] + last_err = None + for opts in option_candidates: + buf = pyio.BytesIO() + try: + container = av.open( + buf, mode="w", format="mp4", + options={"movflags": "frag_keyframe+empty_moov+default_base_moof"}, + ) + stream = container.add_stream("h264_nvenc", rate=int(max(1, fps))) + stream.width = out_w + stream.height = out_h + stream.pix_fmt = "yuv420p" + stream.options = opts + for pf in pil_frames: + for pkt in stream.encode(av.VideoFrame.from_image(pf)): + container.mux(pkt) + for pkt in stream.encode(): + container.mux(pkt) + container.close() + return base64.b64encode(buf.getvalue()).decode("ascii"), out_w, out_h + except Exception as e: + last_err = e + continue + if not _nvenc_warned: + _nvenc_warned = True + logging.warning(f"[KJ PreviewOverride] NVENC MP4 encode failed, using WebP fallback: {last_err}") + return None, 0, 0 + + +def _encode_animated_webp(frames, fps, quality, max_res): + if not frames: + return None, 0, 0 + pil_frames = [] + for f in frames: + pf = f + if pf.mode != "RGB": + pf = pf.convert("RGB") + if max_res and max_res > 0 and (pf.width > max_res or pf.height > max_res): + pf = ImageOps.contain(pf, (max_res, max_res), Image.LANCZOS) + pil_frames.append(pf) + duration_ms = max(1, int(round(1000 / max(1, fps)))) + buf = pyio.BytesIO() + try: + pil_frames[0].save( + buf, + format="WEBP", + save_all=True, + append_images=pil_frames[1:], + duration=duration_ms, + loop=0, + quality=quality, + method=4, + ) + except Exception as e: + logging.warning(f"Animated WebP encode failed: {e}") + return None, 0, 0 + return base64.b64encode(buf.getvalue()).decode("ascii"), pil_frames[0].width, pil_frames[0].height + + +def _interp_db_curve(t, xs, ys): + # Mirrors sampler_nodes._interp_curve. + if t <= xs[0]: + return ys[0] + if t >= xs[-1]: + return ys[-1] + for i in range(len(xs) - 1): + if xs[i] <= t <= xs[i + 1]: + span = xs[i + 1] - xs[i] + if span <= 0: + return ys[i] + f = (t - xs[i]) / span + return ys[i] + f * (ys[i + 1] - ys[i]) + return 0.0 + + +def _detect_detail_boost_curve(sampler, model_patcher, sigmas_list): + # Amount is already baked into ys by the editor, so peak ys == user-set amount. + try: + extra = getattr(sampler, "extra_options", None) or {} + xs = extra.get("db_curve_xs") + ys = extra.get("db_curve_ys") + if "db_wrapped_sampler" not in extra or not xs or not ys or len(xs) != len(ys) or len(xs) < 2: + return None + ms = model_patcher.get_model_object("model_sampling") + start_sigma = float(ms.percent_to_sigma(extra.get("db_start_percent", 0.0))) + end_sigma = float(ms.percent_to_sigma(extra.get("db_end_percent", 1.0))) + # None outside the gate so JS can distinguish "inactive" from "active with value 0". + out = [] + for s in sigmas_list: + sig = float(s) + if sig <= 0 or start_sigma <= end_sigma or sig >= start_sigma or sig <= end_sigma: + out.append(None) + continue + t = (start_sigma - sig) / (start_sigma - end_sigma) + out.append(_interp_db_curve(t, xs, ys)) + return out + except Exception as e: + logging.warning(f"[KJ PreviewOverride] DB curve detection failed: {e}") + return None + + +def _ltx_decode_to_pil(ltx_previewer, x0_5d, max_frames=None, stride=1): + # Pre-shape (B, C, T, H, W) → (B*T, C, H, W); WrappedPreviewer adds the sequence-batch dim. + if ltx_previewer is None or x0_5d.ndim != 5: + return [] + if stride > 1: + x0_5d = x0_5d[:, :, ::stride] + x_moved = x0_5d.movedim(2, 1) # (B, T, C, H, W) — must take shape AFTER movedim + x_in = x_moved.reshape((-1,) + x_moved.shape[-3:]) + rgb = ltx_previewer.decode_latent_to_preview(x_in) + if rgb is None: + return [] + if rgb.ndim == 3: + rgb = rgb.unsqueeze(0) + if rgb.ndim != 4: + return [] + t_total = rgb.shape[0] + if max_frames is not None and 0 < max_frames < t_total: + indices = np.linspace(0, t_total - 1, max_frames).round().astype(int).tolist() + rgb = rgb[indices] + u8 = (rgb * 255).clamp(0, 255).to(torch.uint8).cpu().numpy() + return [Image.fromarray(u8[i]) for i in range(u8.shape[0])] + + +def _ltx_full_vae_decode_to_pil(vae, x0_5d, max_frames=None, stride=1): + # vae.decode handles device + tiling. Slow vs TAEHV but full quality. Output shape + # varies by VAE; we accept (B, T, H, W, C) or (T, H, W, C) and normalize. + if vae is None or x0_5d.ndim != 5: + return [] + if stride > 1: + x0_5d = x0_5d[:, :, ::stride] + try: + images = vae.decode(x0_5d) + except Exception as e: + logging.warning(f"[KJ PreviewOverride] LTX VAE decode failed: {e}") + return [] + if images.ndim == 5: + images = images[0] + if images.ndim != 4: + return [] + t_total = images.shape[0] + if max_frames is not None and 0 < max_frames < t_total: + indices = np.linspace(0, t_total - 1, max_frames).round().astype(int).tolist() + images = images[indices] + u8 = (images.float() * 255).clamp(0, 255).to(torch.uint8).cpu().numpy() + return [Image.fromarray(u8[i]) for i in range(u8.shape[0])] + + +def _is_ltx_latent_format(latent_format): + return "LTX" in type(latent_format).__name__ + + +def _is_ltx2_diffusion_model(model_patcher): + # Same probe as ltxv_nodes.OuterSampleCallbackWrapper. + try: + dm = model_patcher.model.diffusion_model + return not getattr(dm, "caption_projection_first_linear", True) + except Exception: + return False + + +def _ltx_num_keyframes(guider): + try: + positive = guider.conds.get("positive") if hasattr(guider, "conds") else None + if positive and len(positive) > 0: + kf = positive[0].get("keyframe_idxs") + if kf is not None: + return int(torch.unique(kf[0, 0, :, 0]).numel()) + except Exception: + pass + return 0 + + +def _normalize_ltx_x0(x0, latent_shapes, num_keyframes): + # LTX flattens spatial+temporal into a token sequence and may append keyframe latents + # at the tail. Restore 5D and trim so downstream previewers see standard video latents. + if latent_shapes and len(latent_shapes) > 0: + target = latent_shapes[0] + if x0.ndim == 3 and len(target) >= 3: + cut = 1 + for d in target[1:]: + cut *= int(d) + x0 = x0[:, :, :cut].reshape([x0.shape[0]] + list(target)[1:]) + if num_keyframes > 0 and x0.ndim == 5: + x0 = x0[:, :, :-num_keyframes] + return x0 + + +class _PreviewOverrideWrapper: + def __init__(self, max_resolution, node_id, jpeg_quality, suppress_default, preview_frames=1, preview_fps=12, vae=None): + self.max_resolution = max_resolution + self.node_id = str(node_id) if node_id is not None else None + self.jpeg_quality = jpeg_quality + self.suppress_default = suppress_default + self.preview_frames = preview_frames + self.preview_fps = preview_fps + self.vae = vae + self.frames = [] + + def __call__(self, executor, noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes): + guider = executor.class_obj + model_patcher = guider.model_patcher + + is_ltx = _is_ltx_latent_format(model_patcher.model.latent_format) + is_ltx2 = is_ltx and _is_ltx2_diffusion_model(model_patcher) + num_keyframes = _ltx_num_keyframes(guider) if is_ltx else 0 + + # LTX reuses the LTX-specific node's WrappedPreviewer; we call decode_latent_to_preview + # directly per step, bypassing its decode_latent_to_preview_image rate-limiting. + # If a non-TAEHV VAE is supplied, decode via vae.decode() for full quality (slower). + ltx_previewer = None + ltx_full_vae = None + vae_restore_device = None + if is_ltx: + try: + factors, bias = _ltx_rgb_factors(is_ltx2) + taeltx = None + if self.vae is not None: + if self.vae.first_stage_model.__class__.__name__ == "TAEHV": + # TAEHV-LTX decode needs the VAE on GPU; restored at end of __call__. + target_device = comfy.model_management.get_torch_device() + try: + for p in self.vae.first_stage_model.parameters(): + vae_restore_device = p.device + break + self.vae.first_stage_model.to(target_device) + taeltx = self.vae + except Exception as e: + logging.warning(f"[KJ PreviewOverride] Could not move TAEHV-LTX to GPU, skipping: {e}") + else: + # Comfy VAE.decode manages its own device — no pin-to-GPU needed. + ltx_full_vae = self.vae + ltx_previewer = _LTXWrappedPreviewer(factors, bias, rate=8, taeltx=taeltx) + except Exception as e: + logging.warning(f"[KJ PreviewOverride] LTX previewer setup failed: {e}") + + previewer = _get_core_previewer(model_patcher.load_device, model_patcher.model.latent_format) + # Latent2RGB fallback — used when the active previewer returns a non-PIL result + # (e.g. TAEHV/TAESD on a 5D latent). LTX skips this and goes through ltx_previewer. + fallback_previewer = None + try: + lf = model_patcher.model.latent_format + rgb_factors = getattr(lf, "latent_rgb_factors", None) + if rgb_factors is not None: + fallback_previewer = latent_preview.Latent2RGBPreviewer( + rgb_factors, + getattr(lf, "latent_rgb_factors_bias", None), + getattr(lf, "latent_rgb_factors_reshape", None), + ) + except Exception: + pass + + original_callback = callback + node_id = self.node_id + max_res = self.max_resolution + quality = self.jpeg_quality + self.frames = [] + + # N+1 boundaries for N steps: keep them all so the step marker advances through each. + sigmas_list = sigmas.detach().cpu().tolist() if sigmas is not None else [] + # Pre-seed so step 1 has a measurable Δ (model's first transformation from noise → x0). + initial_seed_cpu = None + try: + if sigmas is not None and len(sigmas) > 0: + # sigmas often lives on CPU while noise is on CUDA — align before the multiply. + s0 = sigmas[0].to(noise.device) if hasattr(sigmas[0], "to") else sigmas[0] + seeded = noise * s0 + if is_ltx: + seeded = _normalize_ltx_x0(seeded, latent_shapes, num_keyframes) + initial_seed_cpu = seeded.detach().float().cpu() + except Exception as e: + logging.warning(f"[KJ PreviewOverride] initial seed Δ pre-fill failed: {e}") + state = {"last_x0_cpu": initial_seed_cpu, "last_time": None, "step_ms_window": []} + total_steps_init = max(0, len(sigmas_list) - 1) + + # Boundary-0 message: sigmas (required by JS hover handler) plus optional noise preview. + if node_id is not None and PromptServer is not None: + init_payload = { + "node_id": node_id, + "step": 0, + "total": total_steps_init, + "sigma": sigmas_list[0] if sigmas_list else None, + "sigmas": sigmas_list, + } + db_curve = _detect_detail_boost_curve(sampler, model_patcher, sigmas_list) + if db_curve is not None: + init_payload["db_curve"] = db_curve + # Use Latent2RGB (or LTX previewer) directly — the model's default previewer (TAEHV) + # slices to one temporal frame and returns a shape PIL can't render on raw noise. + try: + lf = model_patcher.model.latent_format + rgb_factors = getattr(lf, "latent_rgb_factors", None) + if sigmas is not None and len(sigmas) > 0: + s0 = sigmas[0].to(noise.device) if hasattr(sigmas[0], "to") else sigmas[0] + init_latent = noise * s0 + else: + init_latent = noise + if is_ltx: + init_latent = _normalize_ltx_x0(init_latent, latent_shapes, num_keyframes) + pil_init = None + if ltx_previewer is not None and init_latent.ndim == 5: + pil_frames = _ltx_decode_to_pil(ltx_previewer, init_latent, max_frames=1) + pil_init = pil_frames[0] if pil_frames else None + elif rgb_factors is not None: + noise_previewer = latent_preview.Latent2RGBPreviewer( + rgb_factors, + getattr(lf, "latent_rgb_factors_bias", None), + getattr(lf, "latent_rgb_factors_reshape", None), + ) + out = noise_previewer.decode_latent_to_preview(init_latent) + if isinstance(out, Image.Image): + pil_init = out + if pil_init is not None: + if pil_init.mode != "RGB": + pil_init = pil_init.convert("RGB") + if max_res and max_res > 0 and (pil_init.width > max_res or pil_init.height > max_res): + pil_init = ImageOps.contain(pil_init, (max_res, max_res), Image.LANCZOS) + ibuf = pyio.BytesIO() + pil_init.save(ibuf, format="JPEG", quality=quality) + init_payload["image"] = base64.b64encode(ibuf.getvalue()).decode("ascii") + init_payload["w"] = pil_init.width + init_payload["h"] = pil_init.height + except Exception as e: + logging.warning(f"Initial noise preview failed (sigmas still sent): {e}") + PromptServer.instance.send_sync("kj_preview_override", init_payload, PromptServer.instance.client_id) + + encoder = _AsyncPreviewEncoder() + animate_video = self.preview_frames > 1 + anim_frames = self.preview_frames + anim_fps = self.preview_fps + + + def new_callback(step, x0, x, total_steps_): + if previewer is not None or fallback_previewer is not None or ltx_previewer is not None: + try: + # NEVER rebind x0 — the sampler reuses the same tensor downstream + # (unpack_latents reshapes it). Preview mutations stay on x0_view. + x0_view = x0 + if is_ltx: + x0_view = _normalize_ltx_x0(x0_view, latent_shapes, num_keyframes) + + pil_frames = [] + max_pil = anim_frames if animate_video else 1 + if ltx_full_vae is not None and x0_view.ndim == 5: + pil_frames = _ltx_full_vae_decode_to_pil(ltx_full_vae, x0_view, max_frames=max_pil) + if not pil_frames and ltx_previewer is not None and x0_view.ndim == 5: + try: + pil_frames = _ltx_decode_to_pil(ltx_previewer, x0_view, max_frames=max_pil) + except Exception as e: + logging.warning(f"LTX preview decode failed: {e}") + if not pil_frames and animate_video and x0_view.ndim == 5 and ltx_previewer is None: + pil_frames = _decode_video_frames_l2rgb( + x0_view, model_patcher.model.latent_format, anim_frames, + ) + + if not pil_frames: + for prev in (previewer, fallback_previewer): + if prev is None: + continue + try: + out = prev.decode_latent_to_preview(x0_view) + except Exception as e: + if prev is previewer: + logging.warning(f"Active previewer raised, trying Latent2RGB fallback: {e}") + continue + if isinstance(out, Image.Image): + pil_frames = [out] + break + elif prev is previewer: + logging.warning( + f"Preview override: {type(previewer).__name__} returned " + f"{type(out).__name__} instead of PIL.Image — falling back to Latent2RGB." + ) + + if not pil_frames: + if original_callback is not None: + original_callback(step, x0, x, total_steps_) + return + + pil_first = pil_frames[0] + if pil_first.mode != "RGB": + pil_first = pil_first.convert("RGB") + pil_frames[0] = pil_first + # Consumed by GetPreviewOverrideFramesKJ. + self.frames.append(pil_first) + + if node_id is not None and PromptServer is not None: + # x0_view (not x0) so LTX keyframe padding doesn't dampen the Δ norm. + x0_cpu_now = x0_view.detach().float().cpu() + prev_x0_cpu = state["last_x0_cpu"] + state["last_x0_cpu"] = x0_cpu_now + + now = time.perf_counter() + step_ms = None + if state["last_time"] is not None: + step_ms = (now - state["last_time"]) * 1000.0 + w = state["step_ms_window"] + w.append(step_ms) + if len(w) > 8: + w.pop(0) + state["last_time"] = now + avg_step_ms = (sum(state["step_ms_window"]) / len(state["step_ms_window"])) if state["step_ms_window"] else None + sigma_val = sigmas_list[step] if 0 <= step < len(sigmas_list) else None + sent_step = step + 1 + + def _encode_and_send( + pil_frames=pil_frames, x0_cpu_now=x0_cpu_now, prev_x0_cpu=prev_x0_cpu, + step_ms=step_ms, avg_step_ms=avg_step_ms, sigma_val=sigma_val, + sent_step=sent_step, total_steps_=total_steps_, + ): + if len(pil_frames) > 1: + # NVENC ~8x faster + ~5x smaller than PIL WebP when available. + b64, w_, h_, mime = None, 0, 0, None + if _NVENC_AVAILABLE: + b64, w_, h_ = _encode_mp4_nvenc(pil_frames, anim_fps, max_res) + if b64: + mime = "video/mp4" + if not b64: + b64, w_, h_ = _encode_animated_webp(pil_frames, anim_fps, quality, max_res) + mime = "image/webp" + else: + pil_send = pil_frames[0] + if max_res and max_res > 0 and (pil_send.width > max_res or pil_send.height > max_res): + pil_send = ImageOps.contain(pil_send, (max_res, max_res), Image.LANCZOS) + buf = pyio.BytesIO() + pil_send.save(buf, format="JPEG", quality=quality) + b64 = base64.b64encode(buf.getvalue()).decode("ascii") + w_, h_ = pil_send.width, pil_send.height + mime = "image/jpeg" + + if not b64: + return + + delta_v = None + if prev_x0_cpu is not None and prev_x0_cpu.shape == x0_cpu_now.shape: + diff = x0_cpu_now - prev_x0_cpu + delta_v = (diff.norm() / max(1, diff.numel()) ** 0.5).item() + + PromptServer.instance.send_sync( + "kj_preview_override", + { + "node_id": node_id, + "image": b64, + "mime": mime, + "w": w_, + "h": h_, + "step": sent_step, + "total": total_steps_, + "sigma": sigma_val, + "sigmas": None, + "delta": delta_v, + "step_ms": step_ms, + "avg_step_ms": avg_step_ms, + "fps": anim_fps if mime in ("video/mp4", "image/webp") else None, + }, + PromptServer.instance.client_id, + ) + + encoder.submit(_encode_and_send) + except Exception as e: + logging.warning(f"Preview override failed: {e}") + if original_callback is not None: + original_callback(step, x0, x, total_steps_) + + # Patch every concrete decode_latent_to_preview_image — subclasses like VHS's + # WrappedPreviewer override it and would otherwise still emit previews of their own. + prev_methods = [] + if self.suppress_default: + targets = [latent_preview.LatentPreviewer] + stack = list(latent_preview.LatentPreviewer.__subclasses__()) + while stack: + cls = stack.pop() + targets.append(cls) + stack.extend(cls.__subclasses__()) + for cls in targets: + if "decode_latent_to_preview_image" in cls.__dict__: + prev_methods.append((cls, cls.__dict__["decode_latent_to_preview_image"])) + cls.decode_latent_to_preview_image = _suppressed_preview_image + try: + # Seeds step 1's duration measurement (sampling-start → end of step 1). + state["last_time"] = time.perf_counter() + return executor(noise, latent_image, sampler, sigmas, denoise_mask, new_callback, disable_pbar, seed, latent_shapes=latent_shapes) + finally: + encoder.shutdown(drain_timeout=5.0) + for cls, prev in prev_methods: + cls.decode_latent_to_preview_image = prev + if vae_restore_device is not None and self.vae is not None: + try: + self.vae.first_stage_model.to(vae_restore_device) + except Exception: + pass + + +class ModelPreviewOverrideKJ(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="ModelPreviewOverrideKJ", + display_name="Model Preview Override", + category="KJNodes/sampling", + description=( + "Adds a dedicated live-preview frame on this node, with overridable max resolution. " + "Default ComfyUI preview caps at 512px; this node sends its own preview straight to a " + "DOM widget on the node so pixel-space models (Chroma Radiance, ZImage, HiDream-O1, …) " + "can be previewed at full sampler resolution." + ), + inputs=[ + io.Model.Input("model", tooltip="Model to attach the preview override to."), + io.Int.Input( + "max_resolution", + default=1024, + min=0, + max=8192, + step=8, + tooltip="Max preview side in pixels for the live widget. 0 = full sampler resolution (no downscale).", + ), + io.Int.Input( + "jpeg_quality", + default=80, + min=30, + max=100, + step=1, + tooltip="JPEG quality for the live preview transport.", + ), + io.Boolean.Input( + "suppress_default_preview", + default=True, + tooltip="Suppress the standard sampler-node preview overlay while sampling, so only this node's frame updates. Progress bar still advances normally.", + ), + io.Int.Input( + "preview_frames", + default=1, + min=1, + max=1024, + step=1, + tooltip="Frames to sample from each video step's latent for animated preview. " + "1 = single frame (current behavior, fastest). >1 = animated WebP playing back at preview_fps. " + "Only applies to video models (5D latents); ignored for image models.", + ), + io.Int.Input( + "preview_fps", + default=12, + min=1, + max=60, + step=1, + tooltip="Playback FPS for the animated WebP preview. Ignored when preview_frames=1.", + ), + io.Vae.Input( + "vae", + optional=True, + tooltip="Optional LTX VAE for true-RGB previews. TAEHV-LTX = fast tiny decode " + "(VAE pinned to GPU). Any other LTX VAE = full-quality decode via " + "vae.decode() — MUCH slower per step.", + ), + ], + outputs=[io.Model.Output(tooltip="Model with preview override attached.")], + hidden=[io.Hidden.unique_id], + is_experimental=True, + ) + + @classmethod + def execute(cls, model, max_resolution, jpeg_quality, suppress_default_preview, preview_frames, preview_fps, vae=None) -> io.NodeOutput: + m = model.clone() + m.add_wrapper_with_key( + comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, + "kj_preview_override", + _PreviewOverrideWrapper( + max_resolution, cls.hidden.unique_id, jpeg_quality, suppress_default_preview, + preview_frames, preview_fps, vae, + ), + ) + return io.NodeOutput(m) + + +class GetPreviewOverrideFramesKJ(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="GetPreviewOverrideFramesKJ", + display_name="Get Preview Override Frames", + category="KJNodes/sampling", + description=( + "Returns the frames captured by Model Preview Override during the most recent sampling. " + "Wire 'model' from Model Preview Override (the same one feeding the sampler) and 'after_sample' " + "from after the sampler (LATENT/IMAGE) to enforce correct execution order." + ), + inputs=[ + io.Model.Input("model", tooltip="The model output by Model Preview Override (used to locate the captured frames)."), + io.MultiType.Input( + "after_sample", + [io.Latent, io.Image], + tooltip="Anything from after the sampler (LATENT or IMAGE). The value is ignored — it just forces this node to run after sampling.", + ), + ], + outputs=[io.Image.Output(display_name="frames")], + is_experimental=True, + ) + + @classmethod + def execute(cls, model, after_sample) -> io.NodeOutput: + wrappers = model.get_wrappers(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, "kj_preview_override") + if not wrappers: + raise RuntimeError("Get Preview Override Frames: no Model Preview Override wrapper found on this model.") + frames = wrappers[-1].frames + if not frames: + raise RuntimeError("Get Preview Override Frames: no frames captured. Ensure the sampler ran with this model.") + tensors = [] + for pil in frames: + arr = np.asarray(pil, dtype=np.float32) / 255.0 + tensors.append(torch.from_numpy(arr)) + return io.NodeOutput(torch.stack(tensors, dim=0)) diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/sharpen_nodes.py b/custom_nodes/ComfyUI-KJNodes/nodes/sharpen_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..dad2e8ce6c0bdae0511648eeb3dfeab20aa53bdc --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/sharpen_nodes.py @@ -0,0 +1,259 @@ +import torch +import torch.nn.functional as F + +from comfy import model_management +from comfy_api.latest import io + +class ImageSharpenKJ(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ImageSharpenKJ", + category="KJNodes/image", + display_name="Image Sharpen KJ", + search_aliases=["sharpen", "unsharp mask", "deconvolution", "cas", "rcas", "high-pass", "postprocessing"], + description="""GPU-accelerated image sharpening with multiple methods. + +**RCAS** — AMD's Robust Contrast-Adaptive Sharpening (from FSR). +Single 5-tap cross filter that adapts to local contrast. +Minimal artifacts, good for general use with little tuning. + +**Adaptive USM** — Unsharp mask with local variance modulation. +Sharpens detail-rich areas more, flat/noisy areas less. +More controllable than RCAS via radius and threshold parameters. + +**High-Pass** — Extracts high-frequency detail and blends it back. +Gives a "clarity" enhancement feel. Uses radius to control detail scale. + +**Deconvolution** — Richardson-Lucy iterative deconvolution. +Can recover actual lost detail from blur, not just enhance edges. +Uses radius as the estimated blur kernel and iterations to control convergence.""", + inputs=[ + io.MatchType.Input("image", io.MatchType.Template("img_or_mask", [io.Image, io.Mask])), + io.DynamicCombo.Input("method", options=[ + io.DynamicCombo.Option(key="rcas", inputs=[ + io.Float.Input("strength", default=0.8, min=0.0, max=1.0, step=0.01, + tooltip="0 = no sharpening, 1 = full RCAS sharpening."), + ]), + io.DynamicCombo.Option(key="adaptive_usm", inputs=[ + io.Float.Input("strength", default=0.5, min=0.0, max=3.0, step=0.01, + tooltip="Sharpening multiplier. Values above 1.0 give aggressive sharpening."), + io.Float.Input("radius", default=1.0, min=0.5, max=5.0, step=0.1, + tooltip="Gaussian blur sigma for the unsharp mask. Larger = enhances coarser detail."), + io.Float.Input("threshold", default=0.05, min=0.0, max=1.0, step=0.01, + tooltip="Noise gate. Higher = only sharpen areas with more texture/detail. 0 = sharpen everything."), + ]), + io.DynamicCombo.Option(key="high_pass", inputs=[ + io.Float.Input("strength", default=0.5, min=0.0, max=3.0, step=0.01, + tooltip="Blend factor for high-frequency detail. Values above 1.0 give a punchier effect."), + io.Float.Input("radius", default=1.0, min=0.5, max=5.0, step=0.1, + tooltip="Gaussian blur sigma defining the frequency cutoff. Larger = enhances coarser detail."), + ]), + io.DynamicCombo.Option(key="deconvolution", inputs=[ + io.Float.Input("strength", default=0.5, min=0.0, max=1.0, step=0.01, + tooltip="Blend between original (0) and fully deconvolved (1)."), + io.Float.Input("radius", default=1.0, min=0.5, max=5.0, step=0.1, + tooltip="Sigma of the assumed Gaussian blur to reverse."), + io.Int.Input("iterations", default=10, min=1, max=100, step=1, + tooltip="Richardson-Lucy iterations. More = sharper but slower, diminishing returns past ~20."), + ]), + ]), + ], + outputs=[ + io.MatchType.Output(io.MatchType.Template("img_or_mask", [io.Image, io.Mask]), display_name="output"), + ], + ) + + @classmethod + def execute(cls, image, method) -> io.NodeOutput: + selected = method["method"] + strength = method.get("strength", 0.5) + if strength == 0: + return io.NodeOutput(image) + + is_mask = image.ndim == 3 # BHW + if is_mask: + image = image.unsqueeze(-1) # BHW -> BHWC with C=1 + + radius = method.get("radius", 1.0) + threshold = method.get("threshold", 0.05) + iterations = method.get("iterations", 10) + if selected == "rcas": + result = _rcas(image, strength) + elif selected == "adaptive_usm": + result = _adaptive_usm(image, strength, radius, threshold) + elif selected == "high_pass": + result = _high_pass(image, strength, radius) + else: + result = _deconvolution(image, strength, radius, iterations) + + if is_mask: + result = result.squeeze(-1) # BHWC -> BHW + + return io.NodeOutput(result) + + +def _rcas(image: torch.Tensor, strength: float) -> torch.Tensor: + """AMD FidelityFX RCAS — 5-tap cross filter with contrast-adaptive lobe.""" + device = model_management.get_torch_device() + intermediate_device = model_management.intermediate_device() + dtype = model_management.intermediate_dtype() + + B, H, W, C = image.shape + out = torch.empty(B, H, W, C, device=intermediate_device, dtype=dtype) + attenuation = strength + + for i in range(B): + img = image[i:i+1].to(device=device, dtype=dtype).permute(0, 3, 1, 2) + padded = F.pad(img, (1, 1, 1, 1), mode='reflect') + + n = padded[:, :, 0:H, 1:W+1] + s = padded[:, :, 2:H+2, 1:W+1] + w = padded[:, :, 1:H+1, 0:W] + e = padded[:, :, 1:H+1, 2:W+2] + center = img + + mn = torch.min(torch.min(torch.min(torch.min(n, s), w), e), center) + mx = torch.max(torch.max(torch.max(torch.max(n, s), w), e), center) + + # hitMin = -mn / (4*mx), hitMax = -(1-mx) / (4*(1-mn)) + hit_min = mn.div(mx * 4.0 + 1e-6).neg_() + hit_max = mx.neg().add_(1.0).div_(mn.neg().add_(1.0).mul_(4.0).add_(1e-6)).neg_() + + torch.max(hit_min, hit_max, out=hit_min) + lobe = hit_min.min(dim=1, keepdim=True).values + + lobe.mul_(attenuation) + lobe.clamp_(-0.1875, 0.0) + + # (center + lobe*(n+s+e+w)) / (1 + 4*lobe) + norm = lobe.mul(4.0).add_(1.0).reciprocal_() + neighbors = n + s + neighbors.add_(e).add_(w) + neighbors.mul_(lobe).add_(center).mul_(norm) + neighbors.clamp_(0.0, 1.0) + + out[i:i+1] = neighbors.permute(0, 2, 3, 1).to(device=intermediate_device) + + return out + + +def _gaussian_kernel_1d(sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + radius = max(int(3.0 * sigma + 0.5), 1) + x = torch.arange(-radius, radius + 1, device=device, dtype=dtype) + kernel = torch.exp(-0.5 * (x / sigma) ** 2) + kernel.div_(kernel.sum()) + return kernel + + +def _gaussian_blur(img: torch.Tensor, sigma: float) -> torch.Tensor: + """Separable Gaussian blur on BCHW tensor.""" + kernel_1d = _gaussian_kernel_1d(sigma, img.device, img.dtype) + k = kernel_1d.shape[0] + C = img.shape[1] + pad = k // 2 + + kh = kernel_1d.view(1, 1, 1, k).expand(C, -1, -1, -1) + blurred = F.conv2d(F.pad(img, (pad, pad, 0, 0), mode='reflect'), kh, groups=C) + + kv = kernel_1d.view(1, 1, k, 1).expand(C, -1, -1, -1) + blurred = F.conv2d(F.pad(blurred, (0, 0, pad, pad), mode='reflect'), kv, groups=C) + + return blurred + + +def _adaptive_usm(image: torch.Tensor, strength: float, radius: float, threshold: float) -> torch.Tensor: + """Unsharp mask modulated by local variance — sharpens texture, skips flat areas.""" + device = model_management.get_torch_device() + intermediate_device = model_management.intermediate_device() + dtype = model_management.intermediate_dtype() + + B, H, W, C = image.shape + out = torch.empty(B, H, W, C, device=intermediate_device, dtype=dtype) + var_sigma = radius * 1.5 + + for i in range(B): + img = image[i:i+1].to(device=device, dtype=dtype, copy=True).permute(0, 3, 1, 2) + + blurred = _gaussian_blur(img, radius) + detail = img - blurred + + # Local variance: E[x^2] - E[x]^2 + local_mean = _gaussian_blur(img, var_sigma) + local_mean_sq = _gaussian_blur(img.square(), var_sigma) + local_mean_sq.sub_(local_mean * local_mean).clamp_(min=0.0) + + # Convert to std deviation so threshold is in pixel-intensity scale + modulation = local_mean_sq.mean(dim=1, keepdim=True).sqrt_() + + if threshold > 0: + modulation.div_(threshold).clamp_(0.0, 1.0) + else: + modulation.fill_(1.0) + + img.addcmul_(detail, modulation, value=strength).clamp_(0.0, 1.0) + + out[i:i+1] = img.permute(0, 2, 3, 1).to(device=intermediate_device) + + return out + + +def _high_pass(image: torch.Tensor, strength: float, radius: float) -> torch.Tensor: + """High-pass sharpening: original + strength * (original - blur).""" + device = model_management.get_torch_device() + intermediate_device = model_management.intermediate_device() + dtype = model_management.intermediate_dtype() + + B, H, W, C = image.shape + out = torch.empty(B, H, W, C, device=intermediate_device, dtype=dtype) + + for i in range(B): + img = image[i:i+1].to(device=device, dtype=dtype).permute(0, 3, 1, 2) + + low_pass = _gaussian_blur(img, radius) + low_pass.neg_().add_(img).mul_(strength).add_(img).clamp_(0.0, 1.0) + + out[i:i+1] = low_pass.permute(0, 2, 3, 1).to(device=intermediate_device) + + return out + + +def _deconvolution(image: torch.Tensor, strength: float, radius: float, iterations: int) -> torch.Tensor: + """Richardson-Lucy deconvolution with Gaussian PSF.""" + device = model_management.get_torch_device() + intermediate_device = model_management.intermediate_device() + dtype = model_management.intermediate_dtype() + + B, H, W, C = image.shape + out = torch.empty(B, H, W, C, device=intermediate_device, dtype=dtype) + + # Gaussian PSF kernel (symmetric, so PSF^T = PSF) + kernel_1d = _gaussian_kernel_1d(radius, device, dtype) + k = kernel_1d.shape[0] + pad = k // 2 + kh = kernel_1d.view(1, 1, 1, k).expand(C, -1, -1, -1) + kv = kernel_1d.view(1, 1, k, 1).expand(C, -1, -1, -1) + + for i in range(B): + img = image[i:i+1].to(device=device, dtype=dtype).permute(0, 3, 1, 2) + estimate = img.clone() + + for _ in range(iterations): + # estimate *= (observed / (estimate * PSF)) * PSF^T + blurred = F.conv2d(F.pad(estimate, (pad, pad, 0, 0), mode='reflect'), kh, groups=C) + blurred = F.conv2d(F.pad(blurred, (0, 0, pad, pad), mode='reflect'), kv, groups=C) + + blurred.add_(1e-6) + torch.div(img, blurred, out=blurred) + + correction = F.conv2d(F.pad(blurred, (pad, pad, 0, 0), mode='reflect'), kh, groups=C) + correction = F.conv2d(F.pad(correction, (0, 0, pad, pad), mode='reflect'), kv, groups=C) + + estimate.mul_(correction) + + torch.lerp(img, estimate, strength, out=estimate) + estimate.clamp_(0.0, 1.0) + + out[i:i+1] = estimate.permute(0, 2, 3, 1).to(device=intermediate_device) + + return out diff --git a/custom_nodes/ComfyUI-KJNodes/nodes/triton_vae.py b/custom_nodes/ComfyUI-KJNodes/nodes/triton_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..0459d95c88659acf5caeca2e0b852d51dd4f75c3 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/nodes/triton_vae.py @@ -0,0 +1,413 @@ +"""Fused norm+SiLU Triton kernels for VAE decode/encode. + +Two architectures are covered, auto-detected by build_object_patches: +- Wan video VAEs: channel-dim RMSNorm (F.normalize(x, dim=1) * sqrt(C) * gamma) + followed by SiLU, wired as Sequential pairs. +- KL image VAEs (Flux2, SDXL, SD1.5, ...): GroupNorm followed by a shared swish + module inside ResnetBlock, plus the GN-only norm_out heads. + +The eager chains make 4-5 full memory passes per call; the fused kernels do one +read+write with fp32 accumulation. The GroupNorm kernel requires channels_last activations, which is +also what makes channels_last conv layout viable for KL VAEs at all: PyTorch's +own channels_last GroupNorm kernel is ~4x slower than NCHW, ours is ~4x faster. +""" +import torch +import torch.nn as nn +import torch.nn.functional as F +try: + import triton + import triton.language as tl +except ImportError as e: + raise ImportError("PatchTritonVAE requires triton (pip install triton, or triton-windows)") from e + +from comfy.ldm.wan.vae import RMS_norm +from comfy.ldm.modules.diffusionmodules.model import ResnetBlock +from comfy.ldm.lightricks.vae.pixel_norm import PixelNorm +from comfy.ldm.lightricks.vae.causal_video_autoencoder import ResnetBlock3D as LTXResnetBlock3D, Encoder as LTXEncoder, Decoder as LTXDecoder + +import copy +import logging + +import comfy.model_patcher +from comfy.patcher_extension import CallbacksMP +from comfy_api.latest import io + +@triton.jit +def _rms_silu_kernel(x_ptr, g_ptr, out_ptr, C, S, stride_b, eps, scale, rcount, BLOCK_S: tl.constexpr, BLOCK_C: tl.constexpr, + HAS_GAMMA: tl.constexpr = True, SILU: tl.constexpr = True, EPS_INSIDE: tl.constexpr = False): + pid_s = tl.program_id(0) + pid_b = tl.program_id(1) + offs_s = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + mask_s = offs_s < S + # int64 offsets: C*S per batch can exceed 2^31 at high resolutions + base = x_ptr + pid_b.to(tl.int64) * stride_b + acc = tl.zeros([BLOCK_S], dtype=tl.float32) + for c0 in range(0, C, BLOCK_C): + offs_c = c0 + tl.arange(0, BLOCK_C) + mask_c = offs_c < C + ptrs = base + offs_c[:, None].to(tl.int64) * S + offs_s[None, :] + v = tl.load(ptrs, mask=mask_c[:, None] & mask_s[None, :], other=0.0).to(tl.float32) + acc += tl.sum(v * v, axis=0) + if EPS_INSIDE: # pixel_norm: 1/sqrt(mean(x^2) + eps) + inv = scale / tl.sqrt(acc * rcount + eps) + else: + inv = scale / tl.maximum(tl.sqrt(acc), eps) + out_base = out_ptr + pid_b.to(tl.int64) * stride_b + for c0 in range(0, C, BLOCK_C): + offs_c = c0 + tl.arange(0, BLOCK_C) + mask_c = offs_c < C + m = mask_c[:, None] & mask_s[None, :] + ptrs = base + offs_c[:, None].to(tl.int64) * S + offs_s[None, :] + v = tl.load(ptrs, mask=m, other=0.0).to(tl.float32) + y = v * inv[None, :] + if HAS_GAMMA: + g = tl.load(g_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + y = y * g[:, None] + if SILU: + y = y * tl.sigmoid(y) + tl.store(out_base + offs_c[:, None].to(tl.int64) * S + offs_s[None, :], y.to(out_ptr.dtype.element_ty), mask=m) + + +@triton.jit +def _rms_silu_cl_kernel(x_ptr, g_ptr, out_ptr, C, ROWS, eps, scale, rcount, BLOCK_S: tl.constexpr, BLOCK_C: tl.constexpr, + HAS_GAMMA: tl.constexpr = True, SILU: tl.constexpr = True, EPS_INSIDE: tl.constexpr = False): + pid = tl.program_id(0) + offs_s = pid * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = tl.arange(0, BLOCK_C) + mask_s = offs_s < ROWS + mask_c = offs_c < C + m = mask_s[:, None] & mask_c[None, :] + # int64 offsets: ROWS*C (= numel) can exceed 2^31 at high resolutions + ptrs = x_ptr + offs_s[:, None].to(tl.int64) * C + offs_c[None, :] + v = tl.load(ptrs, mask=m, other=0.0).to(tl.float32) + acc = tl.sum(v * v, axis=1) + if EPS_INSIDE: + inv = scale / tl.sqrt(acc * rcount + eps) + else: + inv = scale / tl.maximum(tl.sqrt(acc), eps) + y = v * inv[:, None] + if HAS_GAMMA: + g = tl.load(g_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + y = y * g[None, :] + if SILU: + y = y * tl.sigmoid(y) + tl.store(out_ptr + offs_s[:, None].to(tl.int64) * C + offs_c[None, :], y.to(out_ptr.dtype.element_ty), mask=m) + + +@triton.jit +def _gn_stats_cl(x_ptr, sum_ptr, sumsq_ptr, S, + G: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_C: tl.constexpr): + pid = tl.program_id(0) + b = tl.program_id(1) + offs_c = tl.arange(0, BLOCK_C) + acc1 = tl.zeros([G], dtype=tl.float32) + acc2 = tl.zeros([G], dtype=tl.float32) + s0 = pid * BLOCK_S + step = tl.num_programs(0) * BLOCK_S + while s0 < S: + offs_s = s0 + tl.arange(0, BLOCK_S) + m = (offs_s < S)[:, None] + # int64 offsets: (B*S)*C can exceed 2^31 at high resolutions + ptr = x_ptr + (b.to(tl.int64) * S + offs_s[:, None]) * BLOCK_C + offs_c[None, :] + v = tl.load(ptr, mask=m, other=0.0).to(tl.float32) + v3 = tl.reshape(v, (BLOCK_S, G, BLOCK_C // G)) + acc1 += tl.sum(tl.sum(v3, axis=2), axis=0) + acc2 += tl.sum(tl.sum(v3 * v3, axis=2), axis=0) + s0 += step + tl.atomic_add(sum_ptr + b * G + tl.arange(0, G), acc1) + tl.atomic_add(sumsq_ptr + b * G + tl.arange(0, G), acc2) + + +@triton.jit +def _gn_silu_apply_cl(x_ptr, sum_ptr, sumsq_ptr, w_ptr, bias_ptr, out_ptr, S, count, eps, + G: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_C: tl.constexpr, SILU: tl.constexpr): + pid = tl.program_id(0) + b = tl.program_id(1) + CS: tl.constexpr = BLOCK_C // G + offs_s = pid * BLOCK_S + tl.arange(0, BLOCK_S) + offs_c = tl.arange(0, BLOCK_C) + m = (offs_s < S)[:, None] + ptr = x_ptr + (b.to(tl.int64) * S + offs_s[:, None]) * BLOCK_C + offs_c[None, :] + v = tl.load(ptr, mask=m, other=0.0).to(tl.float32) + mean_g = tl.load(sum_ptr + b * G + tl.arange(0, G)) / count + var_g = tl.load(sumsq_ptr + b * G + tl.arange(0, G)) / count - mean_g * mean_g + var_g = tl.maximum(var_g, 0.0) # E[x²]-mean² can go slightly negative from fp32 cancellation + rstd_g = 1.0 / tl.sqrt(var_g + eps) + mean_c = tl.reshape(tl.broadcast_to(mean_g[:, None], (G, CS)), (BLOCK_C,)) + rstd_c = tl.reshape(tl.broadcast_to(rstd_g[:, None], (G, CS)), (BLOCK_C,)) + w = tl.load(w_ptr + offs_c).to(tl.float32) + bias = tl.load(bias_ptr + offs_c).to(tl.float32) + y = (v - mean_c[None, :]) * rstd_c[None, :] * w[None, :] + bias[None, :] + if SILU: + y = y * tl.sigmoid(y) + tl.store(out_ptr + (b.to(tl.int64) * S + offs_s[:, None]) * BLOCK_C + offs_c[None, :], + y.to(out_ptr.dtype.element_ty), mask=m) + + +# autotuned variants: benchmark configs once per shape key, then cache the fastest. +# BLOCK_C of the cl/GN kernels is structural (must hold all of C) so only BLOCK_S/num_warps are tuned. +_rms_silu_kernel_tuned = triton.autotune( + configs=[triton.Config({"BLOCK_S": bs, "BLOCK_C": bc}, num_warps=w) + for bs in (128, 256, 512) for bc in (32, 64) for w in (4, 8)], + key=["C", "S"])(_rms_silu_kernel) + +_rms_silu_cl_kernel_tuned = triton.autotune( + configs=[triton.Config({"BLOCK_S": bs}, num_warps=w) for bs in (4, 8, 16, 32, 64) for w in (4, 8)], + key=["C", "ROWS"])(_rms_silu_cl_kernel) + +_gn_stats_cl_tuned = triton.autotune( + configs=[triton.Config({"BLOCK_S": bs}, num_warps=w) for bs in (8, 16, 32, 64) for w in (4, 8)], + key=["S"], reset_to_zero=["sum_ptr", "sumsq_ptr"])(_gn_stats_cl) + +_gn_silu_apply_cl_tuned = triton.autotune( + configs=[triton.Config({"BLOCK_S": bs}, num_warps=w) for bs in (8, 16, 32, 64) for w in (4, 8)], + key=["S", "count"])(_gn_silu_apply_cl) + + +def fused_rms_silu(x, gamma, scale, eps=1e-12, autotune=False, silu=True, eps_inside=False): + B, C = x.shape[0], x.shape[1] + S = x.numel() // (B * C) + x = x.contiguous() + out = torch.empty_like(x) + flags = {"HAS_GAMMA": gamma is not None, "SILU": silu, "EPS_INSIDE": eps_inside} + g = gamma if gamma is not None else x + if autotune: + grid = lambda meta: (triton.cdiv(S, meta["BLOCK_S"]), B) + _rms_silu_kernel_tuned[grid](x, g, out, C, S, C * S, eps, scale, 1.0 / C, **flags) + else: + _rms_silu_kernel[(triton.cdiv(S, 256), B)](x, g, out, C, S, C * S, eps, scale, 1.0 / C, + BLOCK_S=256, BLOCK_C=32, num_warps=8, **flags) + return out + + +def fused_rms_silu_cl(x, gamma, scale, eps=1e-12, autotune=False, silu=True, eps_inside=False): + C = x.shape[1] + rows = x.numel() // C + out = torch.empty_like(x) + BLOCK_C = triton.next_power_of_2(C) + flags = {"HAS_GAMMA": gamma is not None, "SILU": silu, "EPS_INSIDE": eps_inside} + g = gamma if gamma is not None else x + if autotune: + grid = lambda meta: (triton.cdiv(rows, meta["BLOCK_S"]),) + _rms_silu_cl_kernel_tuned[grid](x, g, out, C, rows, eps, scale, 1.0 / C, BLOCK_C=BLOCK_C, **flags) + else: + BLOCK_S = max(1, 4096 // BLOCK_C) + _rms_silu_cl_kernel[(triton.cdiv(rows, BLOCK_S),)](x, g, out, C, rows, eps, scale, 1.0 / C, + BLOCK_S=BLOCK_S, BLOCK_C=BLOCK_C, num_warps=4, **flags) + return out + + +def fused_gn_silu_cl(x, weight, bias, groups, eps, silu=True, autotune=False): + B, C, H, W = x.shape + S = H * W + sums = torch.zeros(2, B * groups, device=x.device, dtype=torch.float32) + out = torch.empty_like(x) + count = S * (C // groups) + if autotune: + # program count beyond ~1024 only adds same-address atomic contention + grid_stats = lambda meta: (min(1024, triton.cdiv(S, meta["BLOCK_S"])), B) + _gn_stats_cl_tuned[grid_stats](x, sums[0], sums[1], S, G=groups, BLOCK_C=C) + grid_apply = lambda meta: (triton.cdiv(S, meta["BLOCK_S"]), B) + _gn_silu_apply_cl_tuned[grid_apply](x, sums[0], sums[1], weight, bias, out, S, count, eps, + G=groups, BLOCK_C=C, SILU=silu) + else: + BLOCK_S = max(1, 8192 // C) + nprog = min(1024, triton.cdiv(S, BLOCK_S)) + _gn_stats_cl[(nprog, B)](x, sums[0], sums[1], S, G=groups, BLOCK_S=BLOCK_S, BLOCK_C=C, num_warps=8) + _gn_silu_apply_cl[(triton.cdiv(S, BLOCK_S), B)](x, sums[0], sums[1], weight, bias, out, S, count, eps, + G=groups, BLOCK_S=BLOCK_S, BLOCK_C=C, SILU=silu, num_warps=8) + return out + + +class FusedRMSSiLU(nn.Module): + def __init__(self, rms, autotune=False): + super().__init__() + self.gamma = rms.gamma + self.scale = rms.scale + self.autotune = autotune + + def forward(self, x): + if not x.is_cuda: + return F.silu(F.normalize(x, dim=1) * self.scale * self.gamma.to(x)) + gamma = self.gamma + if gamma.device != x.device: # lowvram partial load keeps params offloaded on cpu + gamma = gamma.to(x.device) + gamma = gamma.reshape(-1) + if not gamma.is_contiguous(): + gamma = gamma.contiguous() + if x.ndim == 5 and x.is_contiguous(memory_format=torch.channels_last_3d) and not x.is_contiguous(): + return fused_rms_silu_cl(x, gamma, self.scale, autotune=self.autotune) + return fused_rms_silu(x, gamma, self.scale, autotune=self.autotune) + + +class FusedPixelNorm(nn.Module): + def __init__(self, pn, silu=True, autotune=False): + super().__init__() + self.eps = pn.eps + self.silu = silu + self.autotune = autotune + + def forward(self, x): + if x.is_cuda: + if x.ndim == 5 and x.is_contiguous(memory_format=torch.channels_last_3d) and not x.is_contiguous(): + return fused_rms_silu_cl(x, None, 1.0, self.eps, autotune=self.autotune, silu=self.silu, eps_inside=True) + return fused_rms_silu(x, None, 1.0, self.eps, autotune=self.autotune, silu=self.silu, eps_inside=True) + out = x / torch.sqrt(torch.mean(x ** 2, dim=1, keepdim=True) + self.eps) + if self.silu: + out = F.silu(out) + return out + + +class FusedGNSiLU(nn.Module): + def __init__(self, gn, silu=True, autotune=False): + super().__init__() + self.weight = gn.weight + self.bias = gn.bias + self.num_groups = gn.num_groups + self.eps = gn.eps + self.silu = silu + self.autotune = autotune + + def forward(self, x): + C = x.shape[1] + G = self.num_groups + weight, bias = self.weight, self.bias + if weight.device != x.device: # lowvram partial load keeps params offloaded on cpu + weight = weight.to(x.device) + bias = bias.to(x.device) + if x.is_cuda and x.ndim == 4 and (C & (C - 1)) == 0 and (G & (G - 1)) == 0 and C % G == 0 \ + and x.is_contiguous(memory_format=torch.channels_last) and not x.is_contiguous(): + return fused_gn_silu_cl(x, weight, bias, self.num_groups, self.eps, self.silu, + autotune=self.autotune) + out = F.group_norm(x, self.num_groups, weight, bias, self.eps) + if self.silu: + out = F.silu(out) + return out + + +def convert_conv_layout(model, channels_last=True): + for mod in model.modules(): + if isinstance(mod, nn.Conv3d): + mod.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format) + elif isinstance(mod, nn.Conv2d): + mod.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format) + + +def build_object_patches(model, autotune=False): + """Object patches for vae.patcher: applied at model load, reverted at unload. + Also matches already-fused modules so a loaded+patched model rebuilds cleanly.""" + patches = {} + for name, mod in model.named_modules(): + if isinstance(mod, nn.Sequential): # Wan-style RMS_norm+SiLU pairs + for i in range(len(mod) - 1): + if isinstance(mod[i], RMS_norm) and isinstance(mod[i + 1], nn.SiLU) \ + and mod[i].gamma.ndim == 4 and mod[i].bias is None: + patches[f"{name}.{i}"] = FusedRMSSiLU(mod[i], autotune=autotune) + patches[f"{name}.{i + 1}"] = nn.Identity() + elif isinstance(mod[i], FusedRMSSiLU) and isinstance(mod[i + 1], nn.Identity): + mod[i].autotune = autotune + patches[f"{name}.{i}"] = mod[i] + patches[f"{name}.{i + 1}"] = mod[i + 1] + elif isinstance(mod, ResnetBlock) and not hasattr(mod, "temb_proj"): # KL-style + for norm_name in ("norm1", "norm2"): + norm = getattr(mod, norm_name) + if isinstance(norm, nn.GroupNorm): + patches[f"{name}.{norm_name}"] = FusedGNSiLU(norm, autotune=autotune) + elif isinstance(norm, FusedGNSiLU): + norm.autotune = autotune + patches[f"{name}.{norm_name}"] = norm + patches[f"{name}.swish"] = nn.Identity() + elif name.endswith("norm_out"): # KL encoder/decoder head, SiLU applied separately + if isinstance(mod, nn.GroupNorm): + patches[name] = FusedGNSiLU(mod, silu=False, autotune=autotune) + elif isinstance(mod, FusedGNSiLU): + mod.autotune = autotune + patches[name] = mod + elif isinstance(mod, LTXResnetBlock3D): + # timestep-conditioned decoder blocks apply a scale-shift between norm and + # SiLU, so only the norm itself can be fused there + fuse_silu = not mod.timestep_conditioning + for norm_name in ("norm1", "norm2"): + norm = getattr(mod, norm_name) + if isinstance(norm, PixelNorm) and norm.dim == 1: + patches[f"{name}.{norm_name}"] = FusedPixelNorm(norm, silu=fuse_silu, autotune=autotune) + elif isinstance(norm, FusedPixelNorm): + norm.autotune = autotune + patches[f"{name}.{norm_name}"] = norm + if fuse_silu and f"{name}.norm1" in patches: + patches[f"{name}.non_linearity"] = nn.Identity() + elif isinstance(mod, (LTXEncoder, LTXDecoder)): # LTX head: conv_norm_out (+ scale-shift on conditioned decoder) + conv_act + fuse_silu = isinstance(mod, LTXEncoder) or not mod.timestep_conditioning + if isinstance(mod.conv_norm_out, PixelNorm) and mod.conv_norm_out.dim == 1: + patches[f"{name}.conv_norm_out"] = FusedPixelNorm(mod.conv_norm_out, silu=fuse_silu, autotune=autotune) + elif isinstance(mod.conv_norm_out, FusedPixelNorm): + mod.conv_norm_out.autotune = autotune + patches[f"{name}.conv_norm_out"] = mod.conv_norm_out + if fuse_silu and f"{name}.conv_norm_out" in patches: + patches[f"{name}.conv_act"] = nn.Identity() + return patches + + +class PatchTritonVAE(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PatchTritonVAE", + display_name="Patch Triton VAE", + category="KJNodes/experimental", + is_experimental=True, + description="Speeds up VAE decode/encode with fused Triton norm+SiLU kernels and channels_last conv layout. " + "Supported VAEs (auto-detected): Wan 2.1/2.2 video VAEs incl. Qwen-Image (RMSNorm, ~1.4x/1.15x), KL image VAEs " + "such as Flux/Flux2, SDXL and SD1.5 (GroupNorm, ~1.6-1.8x at 2048px), and LTXV/LTX2 video VAEs (PixelNorm; " + "timestep-conditioned decoder blocks get norm-only fusion). Other architectures are not supported. " + "Applied as object patches on a cloned patcher, so it only exists while this VAE is loaded.", + inputs=[ + io.Vae.Input("vae"), + io.Boolean.Input("fuse_norm_silu", default=True, tooltip="Replace norm+SiLU chains (RMSNorm for Wan, GroupNorm for KL VAEs) with fused Triton kernels (single pass, fp32 accumulation). Requires triton."), + io.Boolean.Input("channels_last", default=True, tooltip="Convert conv weights to channels_last memory format, removing cuDNN layout transposes around every conv. Required for the fused GroupNorm kernel to engage on KL VAEs."), + io.Boolean.Input("autotune", default=False, tooltip="Benchmark several kernel block-size configs on first use of each tensor shape and cache the fastest. Brief stutter per new resolution, usually a few percent faster after warmup."), + ], + outputs=[ + io.Vae.Output(display_name="vae"), + ], + ) + + @classmethod + def execute(cls, vae, fuse_norm_silu=True, channels_last=True, autotune=False) -> io.NodeOutput: + vae = copy.copy(vae) + model = vae.first_stage_model + if vae.patcher.is_dynamic(): + # the dynamic-vram patcher rematerializes weights from pinned host buffers per forward, which would discard channels_last layout + model.to(vae.vae_dtype) + new_patcher = comfy.model_patcher.ModelPatcher( + model, load_device=vae.patcher.load_device, offload_device=vae.patcher.offload_device) + # without a parent, LoadedModel flags is_dead when this patcher is GC'd (leak warning + full gc) + new_patcher.parent = vae.patcher + vae.patcher = new_patcher + + def clear_dynamic_cast_flags(patcher, device_to, lowvram_model_memory, force_patch_weights, full_load): + # a dynamic load of the shared model sets these without ever resetting them; + for mod in patcher.model.modules(): + if hasattr(mod, "comfy_cast_weights") and not hasattr(mod, "prev_comfy_cast_weights"): + mod.comfy_cast_weights = False + mod._v_signature = None + vae.patcher.add_callback_with_key(CallbacksMP.ON_LOAD, "wan_vae_fused_clear_cast", clear_dynamic_cast_flags) + else: + vae.patcher = vae.patcher.clone() + + if fuse_norm_silu: + patches = build_object_patches(model, autotune=autotune) + if not patches: + raise RuntimeError("No fusable norm layers found, this node supports Wan video VAEs, KL image VAEs (Flux2/SDXL/SD1.5) and LTXV/LTX2 video VAEs") + for name, obj in patches.items(): + vae.patcher.add_object_patch(name, obj) + logging.info(f"PatchTritonVAE: registered {len(patches)} fused norm object patches") + + if channels_last: + # reapply on every load: pinned-host/dynamic-vram weight staging restores the original contiguous layout + def reapply_channels_last(patcher, device_to, lowvram_model_memory, force_patch_weights, full_load): + convert_conv_layout(patcher.model, channels_last=True) + vae.patcher.add_callback_with_key(CallbacksMP.ON_LOAD, "wan_vae_fused_channels_last", reapply_channels_last) + else: + convert_conv_layout(model, channels_last=False) + return io.NodeOutput(vae) diff --git a/custom_nodes/ComfyUI-KJNodes/utility/fluid.py b/custom_nodes/ComfyUI-KJNodes/utility/fluid.py new file mode 100644 index 0000000000000000000000000000000000000000..c0691987f5249a031ecbb74329ba513d5788b691 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/utility/fluid.py @@ -0,0 +1,67 @@ +import numpy as np +from scipy.ndimage import map_coordinates, spline_filter +from scipy.sparse.linalg import factorized + +from .numerical import difference, operator + + +class Fluid: + def __init__(self, shape, *quantities, pressure_order=1, advect_order=3): + self.shape = shape + self.dimensions = len(shape) + + # Prototyping is simplified by dynamically + # creating advected quantities as needed. + self.quantities = quantities + for q in quantities: + setattr(self, q, np.zeros(shape)) + + self.indices = np.indices(shape) + self.velocity = np.zeros((self.dimensions, *shape)) + + laplacian = operator(shape, difference(2, pressure_order)) + self.pressure_solver = factorized(laplacian) + + self.advect_order = advect_order + + def step(self): + # Advection is computed backwards in time as described in Stable Fluids. + advection_map = self.indices - self.velocity + + # SciPy's spline filter introduces checkerboard divergence. + # A linear blend of the filtered and unfiltered fields based + # on some value epsilon eliminates this error. + def advect(field, filter_epsilon=10e-2, mode='constant'): + filtered = spline_filter(field, order=self.advect_order, mode=mode) + field = filtered * (1 - filter_epsilon) + field * filter_epsilon + return map_coordinates(field, advection_map, prefilter=False, order=self.advect_order, mode=mode) + + # Apply advection to each axis of the + # velocity field and each user-defined quantity. + for d in range(self.dimensions): + self.velocity[d] = advect(self.velocity[d]) + + for q in self.quantities: + setattr(self, q, advect(getattr(self, q))) + + # Compute the jacobian at each point in the + # velocity field to extract curl and divergence. + jacobian_shape = (self.dimensions,) * 2 + partials = tuple(np.gradient(d) for d in self.velocity) + jacobian = np.stack(partials).reshape(*jacobian_shape, *self.shape) + + divergence = jacobian.trace() + + # If this curl calculation is extended to 3D, the y-axis value must be negated. + # This corresponds to the coefficients of the levi-civita symbol in that dimension. + # Higher dimensions do not have a vector -> scalar, or vector -> vector, + # correspondence between velocity and curl due to differing isomorphisms + # between exterior powers in dimensions != 2 or 3 respectively. + curl_mask = np.triu(np.ones(jacobian_shape, dtype=bool), k=1) + curl = (jacobian[curl_mask] - jacobian[curl_mask.T]).squeeze() + + # Apply the pressure correction to the fluid's velocity field. + pressure = self.pressure_solver(divergence.flatten()).reshape(self.shape) + self.velocity -= np.gradient(pressure) + + return divergence, curl, pressure \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/utility/magictex.py b/custom_nodes/ComfyUI-KJNodes/utility/magictex.py new file mode 100644 index 0000000000000000000000000000000000000000..e6d426f7deb3deb977604dd37581eb4e9fe9e6a9 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/utility/magictex.py @@ -0,0 +1,95 @@ +"""Generates psychedelic color textures in the spirit of Blender's magic texture shader using Python/Numpy + +https://github.com/cheind/magic-texture +""" +from typing import Tuple, Optional +import numpy as np + + +def coordinate_grid(shape: Tuple[int, int], dtype=np.float32): + """Returns a three-dimensional coordinate grid of given shape for use in `magic`.""" + x = np.linspace(-1, 1, shape[1], endpoint=True, dtype=dtype) + y = np.linspace(-1, 1, shape[0], endpoint=True, dtype=dtype) + X, Y = np.meshgrid(x, y) + XYZ = np.stack((X, Y, np.ones_like(X)), -1) + return XYZ + + +def random_transform(coords: np.ndarray, rng: np.random.Generator = None): + """Returns randomly transformed coordinates""" + H, W = coords.shape[:2] + rng = rng or np.random.default_rng() + m = rng.uniform(-1.0, 1.0, size=(3, 3)).astype(coords.dtype) + return (coords.reshape(-1, 3) @ m.T).reshape(H, W, 3) + + +def magic( + coords: np.ndarray, + depth: Optional[int] = None, + distortion: Optional[int] = None, + rng: np.random.Generator = None, +): + """Returns color magic color texture. + + The implementation is based on Blender's (https://www.blender.org/) magic + texture shader. The following adaptions have been made: + - we exchange the nested if-cascade by a probabilistic iterative approach + + Kwargs + ------ + coords: HxWx3 array + Coordinates transformed into colors by this method. See + `magictex.coordinate_grid` to generate the default. + depth: int (optional) + Number of transformations applied. Higher numbers lead to more + nested patterns. If not specified, randomly sampled. + distortion: float (optional) + Distortion of patterns. Larger values indicate more distortion, + lower values tend to generate smoother patterns. If not specified, + randomly sampled. + rng: np.random.Generator + Optional random generator to draw samples from. + + Returns + ------- + colors: HxWx3 array + Three channel color image in range [0,1] + """ + rng = rng or np.random.default_rng() + if distortion is None: + distortion = rng.uniform(1, 4) + if depth is None: + depth = rng.integers(1, 5) + + H, W = coords.shape[:2] + XYZ = coords + x = np.sin((XYZ[..., 0] + XYZ[..., 1] + XYZ[..., 2]) * distortion) + y = np.cos((-XYZ[..., 0] + XYZ[..., 1] - XYZ[..., 2]) * distortion) + z = -np.cos((-XYZ[..., 0] - XYZ[..., 1] + XYZ[..., 2]) * distortion) + + if depth > 0: + x *= distortion + y *= distortion + z *= distortion + y = -np.cos(x - y + z) + y *= distortion + + xyz = [x, y, z] + fns = [np.cos, np.sin] + for _ in range(1, depth): + axis = rng.choice(3) + fn = fns[rng.choice(2)] + signs = rng.binomial(n=1, p=0.5, size=4) * 2 - 1 + + xyz[axis] = signs[-1] * fn( + signs[0] * xyz[0] + signs[1] * xyz[1] + signs[2] * xyz[2] + ) + xyz[axis] *= distortion + + x, y, z = xyz + x /= 2 * distortion + y /= 2 * distortion + z /= 2 * distortion + c = 0.5 - np.stack((x, y, z), -1) + np.clip(c, 0, 1.0) + return c \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/utility/numerical.py b/custom_nodes/ComfyUI-KJNodes/utility/numerical.py new file mode 100644 index 0000000000000000000000000000000000000000..b5b88bc63c45d63d8913e56cbd06eb7ab413fe4f --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/utility/numerical.py @@ -0,0 +1,25 @@ +from functools import reduce +from itertools import cycle +from math import factorial + +import numpy as np +import scipy.sparse as sp + + +def difference(derivative, accuracy=1): + # Central differences implemented based on the article here: + # http://web.media.mit.edu/~crtaylor/calculator.html + derivative += 1 + radius = accuracy + derivative // 2 - 1 + points = range(-radius, radius + 1) + coefficients = np.linalg.inv(np.vander(points)) + return coefficients[-derivative] * factorial(derivative - 1), points + + +def operator(shape, *differences): + # Credit to Philip Zucker for figuring out + # that kronsum's argument order is reversed. + # Without that bit of wisdom I'd have lost it. + differences = zip(shape, cycle(differences)) + factors = (sp.diags(*diff, shape=(dim,) * 2) for dim, diff in differences) + return reduce(lambda a, f: sp.kronsum(f, a, format='csc'), factors) \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/utility/utility.py b/custom_nodes/ComfyUI-KJNodes/utility/utility.py new file mode 100644 index 0000000000000000000000000000000000000000..fac22950a7fcf07abf904f925cd06b8d13f67cde --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/utility/utility.py @@ -0,0 +1,88 @@ +import torch +import numpy as np +from PIL import Image, ImageColor +from typing import Union, List +import logging + +# Utility functions from mtb nodes: https://github.com/melMass/comfy_mtb +def pil2tensor(image: Union[Image.Image, List[Image.Image]]) -> torch.Tensor: + if isinstance(image, list): + return torch.cat([pil2tensor(img) for img in image], dim=0) + + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + +def np2tensor(img_np: Union[np.ndarray, List[np.ndarray]]) -> torch.Tensor: + if isinstance(img_np, list): + return torch.cat([np2tensor(img) for img in img_np], dim=0) + + return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0) + + +def tensor2np(tensor: torch.Tensor): + if len(tensor.shape) == 3: # Single image + return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8) + else: # Batch of images + return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor] + +def tensor2pil(image: torch.Tensor) -> List[Image.Image]: + batch_count = image.size(0) if len(image.shape) > 3 else 1 + if batch_count > 1: + out = [] + for i in range(batch_count): + out.extend(tensor2pil(image[i])) + return out + + return [ + Image.fromarray( + np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) + ) + ] + +def string_to_color(color_string: str) -> List[int]: + color_list = [0, 0, 0] # Default fallback (black) + + if ',' in color_string: + # Handle CSV format (e.g., "255, 0, 0" or "255, 0, 0, 128" or "1.0, 0.5, 0.0") + try: + values = [float(channel.strip()) for channel in color_string.split(',')] + # Convert to 0-255 range if values are in 0-1 range + if all(0 <= v <= 1 for v in values): + color_list = [int(v * 255) for v in values] + else: + color_list = [int(v) for v in values] + except ValueError: + logging.warning(f"Invalid color format: {color_string}. Using default black.") + elif color_string.startswith('#') or (color_string.lstrip('#').isalnum() and not color_string.lstrip('#').replace('.', '', 1).isdigit()): + # Could be Hex format or color name + color_string_stripped = color_string.lstrip('#') + # Try hex first + if len(color_string_stripped) in [6, 8] and all(c in '0123456789ABCDEFabcdef' for c in color_string_stripped): + if len(color_string_stripped) == 6: # #RRGGBB + color_list = [int(color_string_stripped[i:i+2], 16) for i in (0, 2, 4)] + elif len(color_string_stripped) == 8: # #RRGGBBAA + color_list = [int(color_string_stripped[i:i+2], 16) for i in (0, 2, 4, 6)] + else: + # Try color name (e.g., "red", "blue", "cyan") + try: + rgb = ImageColor.getrgb(color_string) + color_list = list(rgb) + except ValueError: + logging.warning(f"Invalid color name or hex format: {color_string}. Using default black.") + else: + # Handle single value (grayscale) - can be int or float + try: + value = float(color_string.strip()) + # Convert to 0-255 range if it's a float between 0-1 + if 0 <= value <= 1: + value = int(value * 255) + else: + value = int(value) + color_list = [value, value, value] + except ValueError: + logging.warning(f"Invalid color format: {color_string}. Using default black.") + + # Clip values to valid range + color_list = np.clip(color_list, 0, 255).tolist() + + return color_list diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/appearance.js b/custom_nodes/ComfyUI-KJNodes/web/js/appearance.js new file mode 100644 index 0000000000000000000000000000000000000000..9069207e2ca079363e757da234318c76d24806ee --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/appearance.js @@ -0,0 +1,23 @@ +const { app } = window.comfyAPI.app; + +app.registerExtension({ + name: "KJNodes.appearance", + nodeCreated(node) { + switch (node.comfyClass) { + case "INTConstant": + node.setSize([200, 58]); + node.color = "#1b4669"; + node.bgcolor = "#29699c"; + break; + case "FloatConstant": + node.setSize([200, 58]); + node.color = LGraphCanvas.node_colors.green.color; + node.bgcolor = LGraphCanvas.node_colors.green.bgcolor; + break; + case "ConditioningMultiCombine": + node.color = LGraphCanvas.node_colors.brown.color; + node.bgcolor = LGraphCanvas.node_colors.brown.bgcolor; + break; + } + } +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/canvas_background.js b/custom_nodes/ComfyUI-KJNodes/web/js/canvas_background.js new file mode 100644 index 0000000000000000000000000000000000000000..530a189628cff5a218effb20914a82d308e8d0a4 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/canvas_background.js @@ -0,0 +1,476 @@ +const { app } = window.comfyAPI.app; + +// ── Helpers ───────────────────────────────────────────────────────────────── + +/** Linearly mix two hex colors: 0 = pure a, 1 = pure b. */ +function mixColors(a, b, t) { + const parse = (hex) => { + if (hex.length === 4) hex = `#${hex[1]}${hex[1]}${hex[2]}${hex[2]}${hex[3]}${hex[3]}`; + return [parseInt(hex.slice(1, 3), 16), parseInt(hex.slice(3, 5), 16), parseInt(hex.slice(5, 7), 16)]; + }; + const ca = parse(a), cb = parse(b); + const r = Math.round(ca[0] + (cb[0] - ca[0]) * t); + const g = Math.round(ca[1] + (cb[1] - ca[1]) * t); + const bl = Math.round(ca[2] + (cb[2] - ca[2]) * t); + return `#${((1 << 24) | (r << 16) | (g << 8) | bl).toString(16).slice(1)}`; +} + +/** Ensure a color string has a leading '#' (PrimeVue ColorPicker omits it). */ +function asHex(color, fallback = "#000000") { + if (!color) return fallback; + return color.startsWith("#") ? color : `#${color}`; +} + +function getSetting(id, fallback) { + return app.ui.settings.getSettingValue(id) ?? fallback; +} + +// ── Pattern generators ────────────────────────────────────────────────────── +// Each returns an OffscreenCanvas to be used with createPattern("repeat"). +// `size` = tile size, `fg`/`bg` = hex colors, `t` = feature thickness. + +function makeDots(size, fg, bg, t) { + const c = new OffscreenCanvas(size, size); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, size, size); + ctx.fillStyle = fg; + ctx.beginPath(); + ctx.arc(size / 2, size / 2, t * 2, 0, Math.PI * 2); + ctx.fill(); + return c; +} + +function makeGrid(size, fg, bg, t) { + const c = new OffscreenCanvas(size, size); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, size, size); + ctx.fillStyle = fg; + ctx.fillRect(size - t, 0, t, size - t); + ctx.fillRect(0, size - t, size, t); + return c; +} + +function makeCrossDots(size, fg, bg, t) { + const c = new OffscreenCanvas(size, size); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, size, size); + const half = size / 2; + const arm = Math.max(t + 1, size * 0.08); + const ht = t / 2; + ctx.fillStyle = fg; + ctx.fillRect(half - arm, half - ht, arm * 2, t); + ctx.fillRect(half - ht, half - arm, t, arm * 2); + return c; +} + +function makeBlueprint(size, fg, bg, t) { + const c = new OffscreenCanvas(size, size); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, size, size); + const sub = size / 4; + const subColor = mixColors(bg, fg, 0.3); + ctx.fillStyle = subColor; + for (let i = sub; i < size; i += sub) { + ctx.fillRect(i, 0, t, size - t); + ctx.fillRect(0, i, size, t); + } + ctx.fillStyle = fg; + ctx.fillRect(size - t, 0, t, size - t); + ctx.fillRect(0, size - t, size, t); + return c; +} + +function makeIsometric(size, fg, bg, t) { + const w = size; + const h = Math.round(size / 2); + const c = new OffscreenCanvas(w, h); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, w, h); + ctx.strokeStyle = fg; + ctx.lineWidth = t; + ctx.beginPath(); + ctx.moveTo(0, h / 2); + ctx.lineTo(w / 2, 0); + ctx.lineTo(w, h / 2); + ctx.moveTo(0, h / 2); + ctx.lineTo(w / 2, h); + ctx.lineTo(w, h / 2); + ctx.stroke(); + return c; +} + +function makeHexagons(size, fg, bg, t) { + const r = size / 2; + const w = r * 3; + const h = Math.round(r * Math.sqrt(3)); + const c = new OffscreenCanvas(w, h); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, w, h); + ctx.strokeStyle = fg; + ctx.lineWidth = t; + + function hex(cx, cy) { + ctx.beginPath(); + for (let i = 0; i < 6; i++) { + const a = Math.PI / 3 * i; + const px = cx + r * Math.cos(a); + const py = cy + r * Math.sin(a); + if (i === 0) ctx.moveTo(px, py); + else ctx.lineTo(px, py); + } + ctx.closePath(); + ctx.stroke(); + } + + hex(r, h / 2); + hex(r * 2.5, 0); + hex(r * 2.5, h); + return c; +} + +function makeOctagons(size, fg, bg, t) { + const c = new OffscreenCanvas(size, size); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, size, size); + ctx.strokeStyle = fg; + ctx.lineWidth = t; + + const s = size / (1 + Math.SQRT2); + const d = (size - s) / 2; + + ctx.beginPath(); + ctx.moveTo(d, 0); + ctx.lineTo(size - d, 0); + ctx.lineTo(size, d); + ctx.lineTo(size, size - d); + ctx.lineTo(size - d, size); + ctx.lineTo(d, size); + ctx.lineTo(0, size - d); + ctx.lineTo(0, d); + ctx.closePath(); + ctx.stroke(); + return c; +} + +function makeWaves(size, fg, bg, t) { + const c = new OffscreenCanvas(size, size); + const ctx = c.getContext("2d"); + ctx.fillStyle = bg; + ctx.fillRect(0, 0, size, size); + ctx.lineWidth = t; + + const lines = Math.max(2, Math.round(size / 20)); + const spacing = size / lines; + + for (let i = 0; i < lines; i++) { + const blend = 0.1 + (i % 3) * 0.07; + ctx.strokeStyle = mixColors(bg, fg, blend); + const baseY = spacing * i + spacing / 2; + const amp = spacing * (0.15 + (i % 4) * 0.05); + + ctx.beginPath(); + for (let x = 0; x <= size; x++) { + const y = baseY + Math.sin((x / size) * Math.PI * 2 + i * 0.8) * amp; + if (x === 0) ctx.moveTo(x, y); + else ctx.lineTo(x, y); + } + ctx.stroke(); + } + return c; +} + +function makeCarbonFiber(size, fg, bg, t) { + const cell = Math.max(2, Math.round(size / 4)); + const s = cell * 4; + const c = new OffscreenCanvas(s, s); + const ctx = c.getContext("2d"); + + const dark = bg; + const light = mixColors(bg, fg, 0.2); + const groove = mixColors(bg, fg, 0.07); + + for (let row = 0; row < 4; row++) { + for (let col = 0; col < 4; col++) { + const phase = (col + row) % 4; + ctx.fillStyle = (phase < 2) ? light : dark; + ctx.fillRect(col * cell, row * cell, cell, cell); + } + } + + ctx.fillStyle = groove; + const gt = Math.max(1, Math.round(t)); + for (let i = 1; i < 4; i++) { + ctx.fillRect(0, i * cell, s, gt); + ctx.fillRect(i * cell, 0, gt, s); + } + return c; +} + + +const PATTERNS = { + default: null, + none: null, + dots: makeDots, + grid: makeGrid, + "cross dots": makeCrossDots, + blueprint: makeBlueprint, + isometric: makeIsometric, + hexagons: makeHexagons, + octagons: makeOctagons, + waves: makeWaves, + "carbon fiber": makeCarbonFiber, +}; + +// ── State ─────────────────────────────────────────────────────────────────── + +let currentType = "default"; +let currentFg = "#444444"; +let currentBg = "#212121"; +let currentSize = 32; +let currentThickness = 2; +let dirty = true; + +function invalidateCache() { + dirty = true; + app.canvas?.setDirty(true, true); +} + + + +function readSettings() { + currentType = getSetting("KJNodes.canvasBg.pattern", "default"); + currentFg = asHex(getSetting("KJNodes.canvasBg.patternColor", "#444444")); + currentBg = asHex(getSetting("KJNodes.canvasBg.bgColor", "#212121")); + currentSize = getSetting("KJNodes.canvasBg.scale", 32); + currentThickness = getSetting("KJNodes.canvasBg.thickness", 2); +} + +// ── Commands ──────────────────────────────────────────────────────────────── + +const patternCommands = Object.keys(PATTERNS).map((key) => ({ + id: `KJNodes.canvasBg.setPattern.${key}`, + label: key.charAt(0).toUpperCase() + key.slice(1), + menubarLabel: key.charAt(0).toUpperCase() + key.slice(1), + function: () => { + app.ui.settings.setSettingValue("KJNodes.canvasBg.pattern", key); + }, + active: () => getSetting("KJNodes.canvasBg.pattern", "default") === key, +})); + +// Single source of truth for slider bounds — commands reference these same objects. +const SCALE_ATTRS = { min: 8, max: 512, step: 4 }; +const THICKNESS_ATTRS = { min: 0.5, max: 20, step: 0.5 }; + +function stepSetting(id, fallback, attrs, direction) { + const cur = getSetting(id, fallback); + const next = cur + attrs.step * direction; + if (next >= attrs.min && next <= attrs.max) { + app.ui.settings.setSettingValue(id, next); + } +} + +const scaleCommands = [ + { + id: "KJNodes.canvasBg.increaseScale", + label: "Increase pattern scale", + menubarLabel: "Increase scale", + function: () => stepSetting("KJNodes.canvasBg.scale", 32, SCALE_ATTRS, 1), + active: () => false, + }, + { + id: "KJNodes.canvasBg.decreaseScale", + label: "Decrease pattern scale", + menubarLabel: "Decrease scale", + function: () => stepSetting("KJNodes.canvasBg.scale", 32, SCALE_ATTRS, -1), + active: () => false, + }, +]; + +const thicknessCommands = [ + { + id: "KJNodes.canvasBg.increaseThickness", + label: "Increase pattern thickness", + menubarLabel: "Increase thickness", + function: () => stepSetting("KJNodes.canvasBg.thickness", 2, THICKNESS_ATTRS, 1), + active: () => false, + }, + { + id: "KJNodes.canvasBg.decreaseThickness", + label: "Decrease pattern thickness", + menubarLabel: "Decrease thickness", + function: () => stepSetting("KJNodes.canvasBg.thickness", 2, THICKNESS_ATTRS, -1), + active: () => false, + }, +]; + +// ── Extension ─────────────────────────────────────────────────────────────── + +app.registerExtension({ + name: "KJNodes.CanvasBackground", + + commands: [...patternCommands, ...scaleCommands, ...thicknessCommands], + + menuCommands: [ + { + path: ["KJNodes", "Canvas Background"], + commands: patternCommands.map((c) => c.id), + }, + { + path: ["KJNodes", "Canvas Background"], + commands: scaleCommands.map((c) => c.id), + }, + { + path: ["KJNodes", "Canvas Background"], + commands: thicknessCommands.map((c) => c.id), + }, + ], + + settings: [ + { + id: "KJNodes.canvasBg.pattern", + name: "Background pattern", + category: ["KJNodes", "Canvas Background", "Pattern"], + tooltip: "Choose a background pattern for the node graph canvas. 'default' uses LiteGraph's built-in dot grid. 'none' disables all background drawing.", + type: "combo", + defaultValue: "default", + options: Object.keys(PATTERNS), + onChange: invalidateCache, + }, + { + id: "KJNodes.canvasBg.patternColor", + name: "Pattern color", + category: ["KJNodes", "Canvas Background", "Pattern color"], + tooltip: "Foreground color of the pattern", + type: "color", + defaultValue: "#444444", + onChange: invalidateCache, + }, + { + id: "KJNodes.canvasBg.bgColor", + name: "Background color", + category: ["KJNodes", "Canvas Background", "Background color"], + tooltip: "Background fill color behind the pattern", + type: "color", + defaultValue: "#212121", + onChange: invalidateCache, + }, + { + id: "KJNodes.canvasBg.scale", + name: "Pattern scale", + category: ["KJNodes", "Canvas Background", "Pattern scale"], + tooltip: "Size of one pattern tile in pixels", + type: "slider", + defaultValue: 32, + attrs: SCALE_ATTRS, + onChange: invalidateCache, + }, + { + id: "KJNodes.canvasBg.thickness", + name: "Pattern thickness", + category: ["KJNodes", "Canvas Background", "Pattern thickness"], + tooltip: "Thickness of lines and size of dots in pixels", + type: "slider", + defaultValue: 2, + attrs: THICKNESS_ATTRS, + onChange: invalidateCache, + }, + ], + + setup() { + const canvas = app.canvas; + if (!canvas) return; + + let patternToken = null; // null = default, "" = none, string = custom + // Snapshot palette values lazily on first override so we capture post-palette state + let origBackgroundImage = null; + let origClearBgColor = null; + let origsCaptured = false; + + function captureOriginals() { + if (origsCaptured) return; + origsCaptured = true; + origBackgroundImage = canvas.background_image; + origClearBgColor = canvas.clear_background_color; + } + + function applyToLiteGraph() { + if (!dirty) return; + dirty = false; + readSettings(); + + // If ComfyUI's own background image setting is active, don't interfere + if (getSetting("Comfy.Canvas.BackgroundImage", "")) { + patternToken = null; + return; + } + + if (currentType === "default") { + if (origsCaptured) { + canvas.background_image = origBackgroundImage; + canvas.clear_background_color = origClearBgColor; + canvas._pattern = undefined; + canvas._bg_img = undefined; + } + patternToken = null; + canvas.setDirty(true, true); + return; + } + + // Capture originals before we overwrite them for the first time + captureOriginals(); + + if (currentType === "none") { + canvas.background_image = ""; + canvas.clear_background_color = currentBg; + canvas._pattern = undefined; + canvas._bg_img = undefined; + patternToken = ""; + canvas.setDirty(true, true); + return; + } + + // Generate tile as OffscreenCanvas — passed directly to LiteGraph's + // createPattern which accepts any CanvasImageSource. + const gen = PATTERNS[currentType]; + if (!gen) return; + + const tile = gen(currentSize, currentFg, currentBg, currentThickness); + patternToken = `kjbg_${currentType}_${currentSize}_${currentThickness}_${currentFg}_${currentBg}`; + tile.name = patternToken; + + canvas.background_image = patternToken; + canvas.clear_background_color = currentBg; + canvas._bg_img = tile; + canvas._pattern = undefined; + canvas.setDirty(true, true); + } + + applyToLiteGraph(); + + // Enforce our values if palette service overwrites them. + // Returns false — LiteGraph's native code does all drawing. + const origCallback = canvas.onRenderBackground; + canvas.onRenderBackground = function (cvs, ctx) { + if (origCallback) { + const result = origCallback.call(this, cvs, ctx); + if (result) return true; + } + + applyToLiteGraph(); + + // If palette service overwrote our values, mark dirty for next frame. + if (!dirty && patternToken !== null && this.background_image !== patternToken) { + dirty = true; + } + + return false; + }; + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/contextmenu.js b/custom_nodes/ComfyUI-KJNodes/web/js/contextmenu.js new file mode 100644 index 0000000000000000000000000000000000000000..9abccc68afa731ce54bb85a772a81eabc6ec6ca2 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/contextmenu.js @@ -0,0 +1,182 @@ +const { app } = window.comfyAPI.app; + +function addNode(name, nextTo, options) { + options = { side: "left", select: true, shiftY: 0, shiftX: 0, ...(options || {}) }; + const node = LiteGraph.createNode(name); + const graph = app.canvas?.graph || app.graph; + graph.add(node); + + node.pos = [ + options.side === "left" ? nextTo.pos[0] - (node.size[0] + options.offset): nextTo.pos[0] + nextTo.size[0] + options.offset, + nextTo.pos[1] + options.shiftY, + ]; + + // Automatically connect nodes + if (options.side === "left") { + // New node on left: connect new node's first output to nextTo's first free input + if (node.outputs && node.outputs.length > 0 && nextTo.inputs && nextTo.inputs.length > 0) { + for (let i = 0; i < nextTo.inputs.length; i++) { + if (!nextTo.inputs[i].link) { + node.pos[1] += i * (node.size[1] + 32); + node.connect(0, nextTo, i); + break; + } + } + } + } else { + // New node on right: connect nextTo's first free output to new node's first free input + if (nextTo.outputs && nextTo.outputs.length > 0 && node.inputs && node.inputs.length > 0) { + for (let o = 0; o < nextTo.outputs.length; o++) { + if (!nextTo.outputs[o].links || nextTo.outputs[o].links.length === 0) { + // Offset vertically by slot index so multiple Set nodes don't overlap + node.pos[1] += o * (node.size[1] + 32); + for (let i = 0; i < node.inputs.length; i++) { + if (!node.inputs[i].link) { + nextTo.connect(o, node, i); + break; + } + } + break; + } + } + } + } + + if (options.select) { + app.canvas.selectNode(node, false); + } + return node; +} + +// Expose addNode for use in setgetnodes.js +window.kjNodes = window.kjNodes || {}; +window.kjNodes.addNode = addNode; + +app.registerExtension({ + name: "KJNodes.Contextmenu", + settings: [ + { + id: "KJNodes.helpPopup", + name: "Help popups", + category: ["KJNodes", "General", "Help popups"], + tooltip: "Show help popups when hovering over KJNodes", + defaultValue: true, + type: "boolean", + }, + { + id: "KJNodes.showSetGetInConnectionMenu", + name: "Show Set/Get in connection menu", + category: ["KJNodes", "Set & Get", "Show Set/Get in connection menu"], + tooltip: "Add Set/Get entries to the slot connection menu (next to Add Reroute)", + type: "boolean", + defaultValue: true, + }, + ], + getNodeMenuItems(node) { + try { + const items = []; + + if (node.inputs && node.inputs.length > 0) { + const selectedForConvert = window.kjNodes.snapshotSelectedNodes(node); + items.push( + { + content: "Add GetNode", + callback: () => { addNode("GetNode", node, { side: "left", offset: 30 }); } + }, + { + content: "Add SetNode", + callback: () => { addNode("SetNode", node, { side: "right", offset: 30 }); } + }, + { + content: "Add PreviewAsTextNode", + callback: () => { addNode("PreviewAny", node, { side: "right", offset: 30 }); } + }, + { + content: "Convert all outputs to Set/Get", + callback: () => { + for (const n of selectedForConvert) window.kjNodes.convertOutputsToSetGet(node.graph, n); + } + }, + ); + } + + const cls = node.constructor?.comfyClass || node.comfyClass; + if (cls && window.kjNodes?.recreateNode) { + items.push({ + content: "Recreate node", + has_submenu: true, + submenu: { + options: [ + { content: "Keep widget values", callback: () => window.kjNodes.recreateNode(node, cls, { resetValues: false }) }, + { content: "Reset widget values", callback: () => window.kjNodes.recreateNode(node, cls, { resetValues: true }) }, + ], + }, + }); + } + + return items; + } catch (err) { + console.error("[KJNodes.Contextmenu] getNodeMenuItems failed:", err); + return []; + } + }, + async setup(app) { + const origShowConnectionMenu = app.canvas.showConnectionMenu.bind(app.canvas); + + // Inject "Add SetNode" / "Add GetNode" into the slot connection dropdown menu. + app.canvas.showConnectionMenu = function(optPass) { + const showSetGet = app.ui.settings.getSettingValue("KJNodes.showSetGetInConnectionMenu") ?? true; + if (!showSetGet) return origShowConnectionMenu(optPass); + + const isFrom = optPass.nodeFrom && optPass.slotFrom != null; + const nodeType = isFrom ? "SetNode" : "GetNode"; + const label = isFrom ? "Add SetNode" : "Add GetNode"; + + const OrigCM = LiteGraph.ContextMenu; + let interceptActive = true; + LiteGraph.ContextMenu = function(options, menuOpts) { + LiteGraph.ContextMenu = OrigCM; + if (!interceptActive) return new OrigCM(options, menuOpts); + interceptActive = false; + + const idx = options.indexOf("Add Reroute"); + if (idx !== -1) options.splice(idx + 1, 0, label); + + const origCb = menuOpts.callback; + menuOpts.callback = function(v, cbOpts, e) { + if (v === label) { + const node = LiteGraph.createNode(nodeType); + if (node) { + const graph = app.canvas?.graph || app.graph; + node.pos = [optPass.e?.canvasX ?? 0, optPass.e?.canvasY ?? 0]; + graph.add(node); + if (isFrom && optPass.nodeFrom) { + const slotIdx = typeof optPass.slotFrom === 'number' + ? optPass.slotFrom + : optPass.nodeFrom.findOutputSlot(optPass.slotFrom.name); + optPass.nodeFrom.connect(slotIdx, node, 0); + } else if (optPass.nodeTo) { + const slotIdx = typeof optPass.slotTo === 'number' + ? optPass.slotTo + : optPass.nodeTo.findInputSlot(optPass.slotTo.name); + node.connect(0, optPass.nodeTo, slotIdx); + } + graph.change(); + app.canvas.setDirty(true, true); + } + return; + } + return origCb?.call(this, v, cbOpts, e); + }; + + return new OrigCM(options, menuOpts); + }; + LiteGraph.ContextMenu.prototype = OrigCM.prototype; + + const result = origShowConnectionMenu(optPass); + LiteGraph.ContextMenu = OrigCM; + interceptActive = false; + return result; + }; + } +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/editors/editor_base.js b/custom_nodes/ComfyUI-KJNodes/web/js/editors/editor_base.js new file mode 100644 index 0000000000000000000000000000000000000000..963d2f7cfaea49efe520fabf7dff1581e9a9a705 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/editors/editor_base.js @@ -0,0 +1,770 @@ +import { addMiddleClickPan, captureVideoFrame, chainCallback, makeUUID, watchImageInputs } from '../utility.js'; + +export function createEditorStylesheet(id, className) { + let styleTag = document.head.querySelector(`#${id}`) + if (!styleTag) { + styleTag = document.createElement('style') + styleTag.type = 'text/css' + styleTag.id = id + styleTag.innerHTML = ` + .${className} { + position: absolute; + font: 12px monospace; + line-height: 1.5em; + padding: 10px; + z-index: 0; + overflow: hidden; + } + .${className} canvas { + position: relative; + z-index: 2; + } + ` + document.head.appendChild(styleTag) + } +} + +function styleMenuItem(menuItem) { + menuItem.style.display = "block"; + menuItem.style.padding = "5px"; + menuItem.style.color = "#FFF"; + menuItem.style.fontFamily = "Arial, sans-serif"; + menuItem.style.fontSize = "16px"; + menuItem.style.textDecoration = "none"; + menuItem.style.marginBottom = "5px"; +} + +function createMenuItem(id, textContent) { + let menuItem = document.createElement("a"); + menuItem.href = "#"; + menuItem.dataset.menuId = id; + menuItem.textContent = textContent; + styleMenuItem(menuItem); + return menuItem; +} + +function setupMenuItems(contextMenu, menuItems) { + menuItems.forEach(mi => { + mi.addEventListener('mouseover', function () { this.style.backgroundColor = "gray"; }); + mi.addEventListener('mouseout', function () { this.style.backgroundColor = "#202020"; }); + contextMenu.appendChild(mi); + }); +} + +export function createContextMenuElement(className) { + const menu = document.createElement("div"); + if (className) menu.className = className; + menu.id = `context-menu-${Math.random().toString(36).slice(2, 10)}`; + menu.style.display = "none"; + menu.style.position = "absolute"; + menu.style.backgroundColor = "#202020"; + menu.style.minWidth = "100px"; + menu.style.boxShadow = "0px 8px 16px 0px rgba(0,0,0,0.2)"; + menu.style.zIndex = "100"; + menu.style.padding = "5px"; + return menu; +} + +// ─── Base Editor Canvas ─── + +const maxDisplayDim = 1024; +const buttonRowHeight = 28; + +export class BaseEditorCanvas { + constructor(context, reset = false) { + this.node = context; + this.reset = reset; + this.bgImage = null; + this.margin = 14; + this.dragIndex = -1; + this.dragType = null; + this.dragOffset = null; + + this._uploadGeneration = 0; + } + + // ─── Shared Methods ─── + + setNodeWidth(width) { + this.node.setSize([width, this.node.size[1]]); + const nodeEl = document.querySelector(`[data-node-id="${this.node.id}"]`); + if (nodeEl) nodeEl.style.setProperty('--node-width', `${width}px`); + } + + // Scale factors: coord space → canvas space + get scaleX() { return this.width / this.coordWidth; } + get scaleY() { return this.height / this.coordHeight; } + + // Returns mouse position in coord space + getLocalMouse(e) { + const rect = this.canvas.getBoundingClientRect(); + const canvasScaleX = this.canvas.width / rect.width; + const canvasScaleY = this.canvas.height / rect.height; + const canvasX = (e.clientX - rect.left) * canvasScaleX - this.margin; + const canvasY = (e.clientY - rect.top) * canvasScaleY - this.margin; + return { + x: canvasX / this.scaleX, + y: canvasY / this.scaleY + }; + } + + // Clamp to coord space + clamp(x, y) { + return { + x: Math.max(0, Math.min(this.coordWidth, x)), + y: Math.max(0, Math.min(this.coordHeight, y)) + }; + } + + // ─── Canvas Setup ─── + + createCanvas(parentElement) { + this.canvas = document.createElement('canvas'); + this.canvas.width = this.width + this.margin * 2; + this.canvas.height = this.height + this.margin * 2; + this.ctx = this.canvas.getContext('2d'); + parentElement.appendChild(this.canvas); + } + + resizeCanvas() { + this.canvas.width = this.width + this.margin * 2; + this.canvas.height = this.height + this.margin * 2; + } + + // ─── Event Listeners ─── + + setupEventListeners() { + this._onDragMove = (e) => this.onMouseMove(e); + this._onDragEnd = (e) => this.onMouseUp(e); + this._removeMiddleClickPan = addMiddleClickPan(this.canvas); + this._onCanvasPointerDown = (e) => { e.stopPropagation(); this.onMouseDown(e); }; + this._onCanvasPointerMove = (e) => this.onMouseMove(e); + this._onCanvasContextMenu = (e) => { e.preventDefault(); e.stopPropagation(); }; + this.canvas.addEventListener('pointerdown', this._onCanvasPointerDown); + this.canvas.addEventListener('pointermove', this._onCanvasPointerMove); + this.canvas.addEventListener('contextmenu', this._onCanvasContextMenu); + } + + removeEventListeners() { + if (this.canvas) { + this.canvas.removeEventListener('pointerdown', this._onCanvasPointerDown); + this.canvas.removeEventListener('pointermove', this._onCanvasPointerMove); + this.canvas.removeEventListener('pointermove', this._onDragMove); + this.canvas.removeEventListener('pointerup', this._onDragEnd); + this.canvas.removeEventListener('contextmenu', this._onCanvasContextMenu); + } + document.removeEventListener('pointermove', this._onDragMove); + document.removeEventListener('pointerup', this._onDragEnd); + if (this._removeMiddleClickPan) this._removeMiddleClickPan(); + } + + startDocumentDrag(e) { + if (e && e.pointerId !== undefined) { + this.canvas.setPointerCapture(e.pointerId); + this._capturedPointerId = e.pointerId; + // With pointer capture, events fire on the canvas, not document + this.canvas.addEventListener('pointermove', this._onDragMove); + this.canvas.addEventListener('pointerup', this._onDragEnd); + } else { + document.addEventListener('pointermove', this._onDragMove); + document.addEventListener('pointerup', this._onDragEnd); + } + } + + endDrag() { + this.dragIndex = -1; + this.dragType = null; + this.dragOffset = null; + this.canvas.style.cursor = 'default'; + if (this._capturedPointerId !== undefined) { + try { this.canvas.releasePointerCapture(this._capturedPointerId); } catch (_) {} + this._capturedPointerId = undefined; + this.canvas.removeEventListener('pointermove', this._onDragMove); + this.canvas.removeEventListener('pointerup', this._onDragEnd); + } else { + document.removeEventListener('pointermove', this._onDragMove); + document.removeEventListener('pointerup', this._onDragEnd); + } + } + + // ─── Render Batching ─── + + render() { + if (this._renderPending) return; + this._renderPending = true; + requestAnimationFrame(() => { + this._renderPending = false; + this._render(); + }); + } + + // Subclasses must implement _render() + + // ─── Common Render Helpers ─── + + // Convert coord-space point to canvas-space pixel + toCanvas(x, y) { + return { x: x * this.scaleX, y: y * this.scaleY }; + } + + // Call at start of _render() — clears and draws background + beginRender() { + const ctx = this.ctx; + ctx.clearRect(0, 0, this.canvas.width, this.canvas.height); + ctx.save(); + ctx.translate(this.margin, this.margin); + ctx.fillStyle = '#222'; + ctx.fillRect(0, 0, this.width, this.height); + ctx.strokeStyle = 'gray'; + ctx.lineWidth = 2; + ctx.strokeRect(0, 0, this.width, this.height); + if (this.bgImage) { + ctx.drawImage(this.bgImage, 0, 0, this.width, this.height); + } + } + + endRender() { + this.ctx.restore(); + } + + // ─── Image Handling ─── + + handleImageLoad = (img, downscaledImg) => { + // Set coord space to image dimensions, rescaling existing points if the space changed + const oldCoordW = this.coordWidth, oldCoordH = this.coordHeight; + this.coordWidth = img.width; + this.coordHeight = img.height; + this.widthWidget.value = img.width; + this.heightWidget.value = img.height; + if (oldCoordW && oldCoordH && (oldCoordW !== img.width || oldCoordH !== img.height)) { + this.onCoordSpaceResized?.(oldCoordW, oldCoordH); + } + this.onImageResize?.(img); + + // Cap display size to the current node width if the user has already resized it, + // otherwise fall back to maxDisplayDim. This prevents the node from expanding to + // fill the image — instead the image scales to fit the node. + const nodeCanvasW = Math.max(64, Math.round(this.node.size[0] - 45)); + const fitDim = Math.min(nodeCanvasW, maxDisplayDim); + let displayW = img.width, displayH = img.height; + if (displayW > fitDim || displayH > fitDim) { + const scale = fitDim / Math.max(displayW, displayH); + displayW = Math.round(displayW * scale); + displayH = Math.round(displayH * scale); + } + + if (displayW !== this.width || displayH !== this.height) { + this.width = displayW; + this.height = displayH; + this.resizeCanvas(); + + // Only expand the node width if it's narrower than the image — never shrink it + if (displayW + 45 > this.node.size[0]) this.setNodeWidth(displayW + 45); + this.onSizeChanged(); + if (this.node.graph) { + try { this.node.arrange?.(); } catch (_) {} + this.node.graph.setDirtyCanvas(true, true); + } + } + + // Use downscaled image if available to avoid holding full-res in memory + this.bgImage = downscaledImg || img; + this.render(); + this.onDataChanged(); + }; + + // resizeCanvas: if true (default), resize editor canvas to match image. + // If false, just store the image and update the background without resizing. + processImage = (img, { resize = true } = {}) => { + const canvas = document.createElement('canvas'); + const ctx = canvas.getContext('2d'); + let width = img.width, height = img.height; + if (width > maxDisplayDim || height > maxDisplayDim) { + const scale = maxDisplayDim / Math.max(width, height); + width = Math.round(width * scale); + height = Math.round(height * scale); + } + canvas.width = width; + canvas.height = height; + ctx.drawImage(img, 0, 0, width, height); + + const { app } = window.comfyAPI.app; + const embed = app.ui.settings.getSettingValue("KJNodes.editors.embedBackgroundImage") ?? false; + + // Use the downscaled canvas directly as bgImage — drawImage accepts canvas elements, + // avoids a data URL round-trip, and is immediately available (no async decode). + const onStored = () => { + if (resize) { + this.handleImageLoad(img, canvas); + } else { + this.bgImage = canvas; + this.render(); + } + }; + + const gen = ++this._uploadGeneration; + if (embed) { + const base64String = canvas.toDataURL('image/webp', 0.5).replace(/^data:.+?,/, ''); + if (gen !== this._uploadGeneration) return; + this.node.properties.imgData = { type: 'image/webp', base64: base64String }; + onStored(); + } else { + canvas.toBlob((blob) => { + if (gen !== this._uploadGeneration) return; + const filename = `editor_bg_${this.node.id}_${Date.now()}.webp`; + const formData = new FormData(); + formData.append('image', blob, filename); + formData.append('type', 'temp'); + formData.append('overwrite', 'true'); + fetch('/upload/image', { method: 'POST', body: formData }) + .then(r => r.json()) + .then(result => { + if (gen !== this._uploadGeneration) return; + this.node.properties.imgData = { type: 'temp', filename: result.name }; + onStored(); + }) + .catch(e => console.error("Failed to upload editor background:", e)); + }, 'image/webp', 0.5); + } + }; + + handleImageFile = (file) => { + const url = URL.createObjectURL(file); + const img = new Image(); + img.onload = () => { URL.revokeObjectURL(url); this.processImage(img); }; + img.onerror = () => URL.revokeObjectURL(url); + img.src = url; + }; + + refreshBackgroundImage = () => { + const imgData = this.node.properties.imgData; + if (!imgData) return; + + const img = new Image(); + img.onerror = (e) => console.error("Background image failed to load:", e); + img.onload = () => { + // Just set the background — don't resize canvas, widget values are already correct from serialization + this.bgImage = img; + this.render(); + }; + + if (imgData.base64) { + const mimeType = imgData.type || 'image/png'; + img.src = `data:${mimeType};base64,${imgData.base64}`; + } else if (imgData.filename) { + img.src = `/view?filename=${encodeURIComponent(imgData.filename)}&type=temp&no-cache=${Date.now()}`; + } + }; + + // ─── Context Menu Helpers ─── + + // Set up document-level listeners for context menu behavior + setupContextMenuListeners(editorIdPrefix) { + this._onContextMenu = (e) => { + if (e.target.closest(`#${editorIdPrefix}-${this.node.uuid}`) || + this.node.contextMenu.contains(e.target)) { + e.preventDefault(); + } + }; + this._onDocClick = (e) => { + if (!this.node.contextMenu.contains(e.target)) { + this.node.contextMenu.style.display = 'none'; + } + }; + document.addEventListener('contextmenu', this._onContextMenu); + document.addEventListener('click', this._onDocClick); + } + + // Clean up previous editor instance + cleanupPreviousEditor(context) { + if (context.editor) { + context.editor.destroy(); + } + } + + // Full cleanup — override in subclass to add additional cleanup + destroy() { + this._uploadGeneration++; // Invalidate any pending async image uploads + this.removeEventListeners(); + if (this._onContextMenu) document.removeEventListener('contextmenu', this._onContextMenu); + if (this._onDocClick) document.removeEventListener('click', this._onDocClick); + if (this._onKeyUp) document.removeEventListener('keyup', this._onKeyUp); + } + + // Common menu actions + openImageFilePicker() { + const fileInput = document.createElement('input'); + fileInput.type = 'file'; + fileInput.accept = 'image/*'; + fileInput.addEventListener('change', (event) => { + const file = event.target.files[0]; + if (file) this.handleImageFile(file); + }); + fileInput.click(); + } + + clearBackgroundImage() { + this.bgImage = null; + this.node.properties.imgData = null; + this.render(); + } + + // Find a widget by name on the node + findWidget(name) { + const w = this.node.widgets.find(w => w.name === name); + if (!w) console.warn(`${this.constructor.name}: widget "${name}" not found`); + return w; + } + + // Show context menu at mouse position + showContextMenu(e) { + this._updateMenuToggleStates(); + const menu = this.node.contextMenu; + menu.style.display = 'block'; + menu.style.left = `${e.clientX}px`; + menu.style.top = `${e.clientY}px`; + // Adjust if menu overflows viewport + const rect = menu.getBoundingClientRect(); + if (rect.right > window.innerWidth) menu.style.left = `${Math.max(0, e.clientX - rect.width)}px`; + if (rect.bottom > window.innerHeight) menu.style.top = `${Math.max(0, e.clientY - rect.height)}px`; + } + + // ─── Width/Height Resize ─── + + setupSizeCallbacks() { + this.widthWidget.callback = () => { + const oldCoordW = this.coordWidth; + this.coordWidth = this.widthWidget.value; + this.onCoordSpaceResized(oldCoordW, this.coordHeight); + this.render(); + this.onDataChanged(); + }; + this.heightWidget.callback = () => { + const oldCoordH = this.coordHeight; + this.coordHeight = this.heightWidget.value; + this.onCoordSpaceResized(this.coordWidth, oldCoordH); + this.render(); + this.onDataChanged(); + }; + } + + // ─── imgData persistence ─── + // Call once from onNodeCreated to ensure imgData survives save/reload + static setupImgDataPersistence(node) { + chainCallback(node, "onSerialize", function (o) { + if (this.properties.imgData && o.properties) { + o.properties.imgData = this.properties.imgData; + } + }); + chainCallback(node, "onConfigure", function (info) { + if (info?.properties?.imgData) { + this.properties.imgData = info.properties.imgData; + } + }); + } + + // ─── Node registration helper ─── + // Call from onNodeCreated to set up the common editor scaffolding. + // config: { editorClass, editorKey, heightKey, className, menuItems, hiddenWidgets, initialSize, extraProperties } + // editorClass: the editor constructor (e.g. SplineEditor) + // editorKey: node property name for the DOM widget (e.g. 'splineEditor') + // heightKey: node property name for editor height (e.g. 'splineEditorHeight') + // className: CSS class for the editor container + // menuItems: { id: { label, action(editor), toggle?(editor) }, ... } — context menu definition + // menuClassName: optional CSS class for the context menu + // hiddenWidgets: widget names to hide + // initialSize: [width, height] for the node + // extraProperties: array of [name, default, type] to register via addProperty on first create + static setupNode(node, nodeData, config) { + const { editorClass, editorKey, heightKey, className, menuItems, + menuClassName, hiddenWidgets, initialSize, extraProperties } = config; + + if (!node.properties) node.properties = {}; + if (node.properties.imgData === undefined) node.properties.imgData = null; + + // File handling — set up once per node, delegate to current editor + node.pasteFile = (file) => { + if (node.editor && file.type.startsWith("image/")) { + node.editor.handleImageFile(file); + return true; + } + return false; + }; + node.onDragOver = (e) => { + if (node.editor && e.dataTransfer && e.dataTransfer.items) { + return [...e.dataTransfer.items].some(f => f.kind === "file" && f.type.startsWith("image/")); + } + return false; + }; + node.onDragDrop = (e) => { + if (!node.editor) return false; + let handled = false; + for (const file of e.dataTransfer.files) { + if (file.type.startsWith("image/")) { + node.editor.handleImageFile(file); + handled = true; + } + } + return handled; + }; + + for (const name of (hiddenWidgets || [])) { + const w = node.widgets.find(w => w.name === name); + if (w) w.hidden = true; + } + + const element = document.createElement("div"); + node.uuid = makeUUID(); + element.id = `${className}-${node.uuid}`; + node.previewMediaType = 'image'; + + node[editorKey] = node.addDOMWidget(nodeData.name, `${editorClass.name}Widget`, element, { + serialize: false, hideOnZoom: false, + getMinHeight: () => node[heightKey] || 550, + getMaxHeight: () => node[heightKey] || 550, + getHeight: () => node[heightKey] || 550, + }); + node[heightKey] = 550; + + node.contextMenu = createContextMenuElement(menuClassName); + node._menuDef = menuItems; + const menuEls = Object.entries(menuItems).map(([id, def]) => createMenuItem(id, def.label || id)); + setupMenuItems(node.contextMenu, menuEls); + document.body.appendChild(node.contextMenu); + + // Shared helper — loads the stored background image into the editor. + // alignToImage: if true, resizes the node to fit the image first (old behaviour). + // if false, scales the image to fit the current node size (new behaviour). + const _reloadBgImage = (alignToImage) => { + const imgData = node.properties.imgData; + if (!imgData) return; + const img = new Image(); + img.onload = () => { + if (!node.editor) return; + if (alignToImage) { + // Temporarily lift the node-size cap so processImage can expand the node + const savedSize = [node.size[0], node.size[1]]; + node.setSize([Math.min(img.width, maxDisplayDim) + 45, savedSize[1]]); + } + node.editor.processImage(img); + }; + if (imgData.base64) { + img.src = `data:${imgData.type || 'image/png'};base64,${imgData.base64}`; + } else if (imgData.filename) { + img.src = `/view?filename=${encodeURIComponent(imgData.filename)}&type=temp&no-cache=${Date.now()}`; + } + }; + + // Two side-by-side buttons in a single DOM widget row. + // "Reset canvas" — clears points and re-fits the image to the current node size. + // "Align to image" — resizes the node to match the image dimensions (legacy behaviour). + const buttonRow = document.createElement("div"); + buttonRow.style.cssText = "display:flex;gap:4px;width:100%;box-sizing:border-box;"; + const makeRowBtn = (label, onClick) => { + const btn = document.createElement("button"); + btn.textContent = label; + btn.style.cssText = "flex:1;height:24px;padding:0 6px;font:12px sans-serif;cursor:pointer;"; + btn.addEventListener("click", onClick); + return btn; + }; + buttonRow.appendChild(makeRowBtn("Reset canvas", () => { + try { + node.editor = new editorClass(node, true); + _reloadBgImage(false); + } catch (error) { console.error(`Error creating ${editorClass.name}:`, error); } + })); + buttonRow.appendChild(makeRowBtn("Align to image", () => { + try { + if (!node.editor) node.editor = new editorClass(node); + _reloadBgImage(true); + } catch (error) { console.error(`Error aligning ${editorClass.name}:`, error); } + })); + node.addDOMWidget("editor_buttons", "editorButtonRow", buttonRow, { + serialize: false, hideOnZoom: false, + getMinHeight: () => buttonRowHeight, getMaxHeight: () => buttonRowHeight, getHeight: () => buttonRowHeight, + }); + + node.setSize(initialSize); + node[editorKey].parentEl = document.createElement("div"); + node[editorKey].parentEl.className = className; + node[editorKey].parentEl.id = `${className}-${node.uuid}`; + element.appendChild(node[editorKey].parentEl); + + // Auto-create for new nodes; onConfigure cancels this for saved workflows + node._autoCreatePending = setTimeout(() => { + if (!node.editor) { + try { + node.editor = new editorClass(node); + for (const [name, value, type] of (extraProperties || [])) { + node.addProperty(name, value, type); + } + } catch (error) { console.error(`Error creating ${editorClass.name}:`, error); } + } + }, 0); + + BaseEditorCanvas.setupImgDataPersistence(node); + + chainCallback(node, "onResize", function () { + const editor = this.editor; + if (!editor) return; + const newWidth = Math.max(64, Math.round(this.size[0] - 45)); + if (newWidth === editor.width) return; + + // Only change display size — coord space stays the same + editor.width = newWidth; + editor.height = Math.round(newWidth * (editor.coordHeight / editor.coordWidth)); + + editor.resizeCanvas(); + editor.onSizeChanged(); + editor.render(); + }); + + chainCallback(node, "onConfigure", function () { + clearTimeout(this._autoCreatePending); + try { this.editor = new editorClass(this); } + catch (error) { console.error(`Error configuring ${editorClass.name}:`, error); } + // Allow watchImageInputs to operate only after configure + a delay, + // so refreshBackgroundImage has time to restore the saved image first. + node._bgWatchReady = false; + setTimeout(() => { node._bgWatchReady = true; }, 500); + }); + + chainCallback(node, "onExecuted", function (message) { + let bg_image = message["bg_image"]; + if (Array.isArray(bg_image)) bg_image = bg_image[0]; + if (bg_image) { + const img = new Image(); + img.src = `data:image/jpeg;base64,${bg_image}`; + img.onload = () => { + if (this.editor) this.editor.processImage(img); + }; + } + }); + + // Load background image from connected source node (LoadImage, VHS, etc.) + // _bgWatchReady is false during reload until the editor has restored its saved background. + // _bgFromConnectedSource tracks whether the current bg came from a connection (vs execution/paste). + node._bgFromConnectedSource = false; + node._bgWatchReady = true; // true for new nodes, set false during onConfigure + watchImageInputs(node, "bg_image", (sources) => { + if (!node.editor || !node._bgWatchReady) return; + const source = sources[0]; + if (source?.isVideo && source.videoEl) { + // Capture the first/current frame from a VHS-style video preview + node._bgFromConnectedSource = true; + captureVideoFrame(source.videoEl, (canvas) => { + if (node.editor) node.editor.processImage(canvas); + }); + } else if (source && !source.isVideo) { + node._bgFromConnectedSource = true; + const img = new Image(); + img.crossOrigin = "anonymous"; + img.onerror = (e) => console.error("[Editor] source image load error:", e); + img.onload = () => { if (node.editor) node.editor.processImage(img); }; + img.src = source.url; + } else if (!source && node._bgFromConnectedSource) { + node._bgFromConnectedSource = false; + node.editor.clearBackgroundImage(); + } + }); + + chainCallback(node, "onRemoved", function () { + if (this.editor) this.editor.destroy(); + if (this.contextMenu?.parentNode) this.contextMenu.parentNode.removeChild(this.contextMenu); + }); + } + + // ─── Shared constructor flow ─── + // Call from subclass constructor after setting up widgets and data. + // editorKey: e.g. 'pointsEditor' or 'splineEditor' + // heightKey: e.g. 'pointsEditorHeight' or 'splineEditorHeight' + // heightOffset: pixels added to canvas height for full node height (e.g. 310 or 460) + initEditor(editorKey, heightKey, heightOffset) { + this._editorKey = editorKey; + this._heightKey = heightKey; + this._heightOffset = heightOffset; + + this.createCanvas(this.node[editorKey].element); + + if (this.width > 256) this.setNodeWidth(this.width + 45); + this.node[heightKey] = this.height + 40; + this.node.setSize([this.node.size[0], this.height + heightOffset + buttonRowHeight]); + + this.setupEventListeners(); + this.render(); + } + + // Shared onSizeChanged — uses stored heightKey/heightOffset + onSizeChanged() { + this.node[this._heightKey] = this.height + 40; + this.node.setSize([this.node.size[0], this.height + this._heightOffset + buttonRowHeight]); + if (this.node.graph) this.node.graph.setDirtyCanvas(true, true); + } + + // Shared constructor preamble — cleanup, reset, context menu, coord/display init + initEditorPreamble(editorKey, className) { + this._className = className; + this.cleanupPreviousEditor(this.node); + if (this.reset && this.node[editorKey].element) { + this.node[editorKey].element.innerHTML = ''; + } + this.createContextMenu(); + } + + // Shared coord/display size init from widgets + saved node size + initDisplaySize() { + this.coordWidth = this.widthWidget.value; + this.coordHeight = this.heightWidget.value; + const savedWidth = Math.max(64, Math.round(this.node.size[0] - 45)); + this.width = Math.min(savedWidth, maxDisplayDim); + this.height = Math.round(this.width * (this.coordHeight / this.coordWidth)); + } + + // Shared context menu creation — clone to clear stale listeners, wire up action handlers + createContextMenu() { + const oldMenu = this.node.contextMenu; + const newMenu = oldMenu.cloneNode(true); + oldMenu.parentNode.replaceChild(newMenu, oldMenu); + this.node.contextMenu = newMenu; + this.setupContextMenuListeners(this._className); + + const self = this; + newMenu.addEventListener('click', (e) => { + e.preventDefault(); + if (e.target.tagName !== 'A') return; + const id = e.target.dataset.menuId; + const def = this.node._menuDef[id]; + if (def?.action) { + def.action(self); + self._updateMenuToggleStates(); + } + newMenu.style.display = 'none'; + }); + } + + // Update toggle item styling from menu definitions + _updateMenuToggleStates() { + const menuDef = this.node._menuDef; + this.node.contextMenu.querySelectorAll('a').forEach(item => { + const def = menuDef[item.dataset.menuId]; + if (def?.toggle) { + const on = def.toggle(this); + item.style.color = on ? '#4fc3f7' : '#FFF'; + item.style.borderLeft = on ? '3px solid #4fc3f7' : '3px solid transparent'; + item.style.paddingLeft = '8px'; + } + }); + } + + // ─── Hooks for subclasses ─── + // Override these instead of duplicating logic: + + // Called after data changes (render + update widgets) + onDataChanged() {} + + // Called on image load for editor-specific state (e.g., drawRuler = false) + onImageResize() {} + + // Called when coord space changes — override to rescale coordinates + onCoordSpaceResized(_oldWidth, _oldHeight) {} +} diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/editors/interpolation.js b/custom_nodes/ComfyUI-KJNodes/web/js/editors/interpolation.js new file mode 100644 index 0000000000000000000000000000000000000000..728f923109856572ba9f6474e92fd842eb21ad13 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/editors/interpolation.js @@ -0,0 +1,211 @@ +// ─── Interpolation: SVG path d-string builders ─── +// Ported from protovis basis/cardinal/monotone/hermite algorithms + +const Interpolation = { + _basisPoint(t0, p0, t1, p1, t2, p2, t3, p3) { + return t0 * p0 + t1 * p1 + t2 * p2 + t3 * p3; + }, + + _pathBasis(p0, p1, p2, p3) { + const bp = Interpolation._basisPoint; + const x1 = bp(0, p0.x, 2/3, p1.x, 1/3, p2.x, 0, p3.x); + const y1 = bp(0, p0.y, 2/3, p1.y, 1/3, p2.y, 0, p3.y); + const x2 = bp(0, p0.x, 1/3, p1.x, 2/3, p2.x, 0, p3.x); + const y2 = bp(0, p0.y, 1/3, p1.y, 2/3, p2.y, 0, p3.y); + const x = bp(0, p0.x, 1/6, p1.x, 2/3, p2.x, 1/6, p3.x); + const y = bp(0, p0.y, 1/6, p1.y, 2/3, p2.y, 1/6, p3.y); + return `C${x1},${y1},${x2},${y2},${x},${y}`; + }, + + basis(points) { + if (points.length <= 2) return Interpolation.linear(points); + let d = ''; + let b0 = points[0], b1 = points[0], b2 = points[0], b3 = points[1]; + d += Interpolation._pathBasis(b0, b1, b2, b3); + for (let i = 2; i < points.length; i++) { + b0 = b1; b1 = b2; b2 = b3; b3 = points[i]; + d += Interpolation._pathBasis(b0, b1, b2, b3); + } + b0 = b1; b1 = b2; b2 = b3; + d += Interpolation._pathBasis(b0, b1, b2, b3); + b0 = b1; b1 = b2; + d += Interpolation._pathBasis(b0, b1, b2, b3); + return d; + }, + + _cardinalTangents(points, tension) { + const alpha = (1 - tension) / 2; + const tangents = []; + let f = points[0], g = points[1], h = points[2]; + for (let i = 3; i < points.length; i++) { + tangents.push({ x: alpha * (h.x - f.x), y: alpha * (h.y - f.y) }); + f = g; g = h; h = points[i]; + } + tangents.push({ x: alpha * (h.x - f.x), y: alpha * (h.y - f.y) }); + return tangents; + }, + + _hermite(points, tangents) { + if (tangents.length < 1 || + (points.length !== tangents.length && points.length !== tangents.length + 2)) + return ''; + const quad = points.length !== tangents.length; + let d = ''; + let g = points[0], h = points[1], i = tangents[0], j = i, k = 1; + + if (quad) { + d += `Q${h.x - i.x * 2/3},${h.y - i.y * 2/3},${h.x},${h.y}`; + g = points[1]; k = 2; + } + if (tangents.length > 1) { + j = tangents[1]; h = points[k]; k++; + d += `C${g.x + i.x},${g.y + i.y},${h.x - j.x},${h.y - j.y},${h.x},${h.y}`; + for (let idx = 2; idx < tangents.length; idx++, k++) { + h = points[k]; j = tangents[idx]; + d += `S${h.x - j.x},${h.y - j.y},${h.x},${h.y}`; + } + } + if (quad) { + const last = points[k]; + d += `Q${h.x + j.x * 2/3},${h.y + j.y * 2/3},${last.x},${last.y}`; + } + return d; + }, + + cardinal(points, tension) { + if (points.length <= 2) return Interpolation.linear(points); + return Interpolation._hermite(points, Interpolation._cardinalTangents(points, tension)); + }, + + _monotoneTangents(points) { + const n = points.length; + const d = [], m = [], dx = []; + for (let i = 0; i < n - 1; i++) d.push((points[i + 1].y - points[i].y) / (points[i + 1].x - points[i].x)); + m.push(d[0]); + for (let i = 1; i < n - 1; i++) m.push((d[i - 1] + d[i]) / 2); + m.push(d[n - 2]); + dx.push(points[1].x - points[0].x); + for (let i = 1; i < n - 1; i++) dx.push((points[i + 1].x - points[i - 1].x) / 2); + dx.push(points[n - 1].x - points[n - 2].x); + for (let i = 0; i < n - 1; i++) { if (Math.abs(d[i]) < 1e-7) { m[i] = 0; m[i + 1] = 0; } } + for (let i = 0; i < n - 1; i++) { + if (Math.abs(m[i]) >= 1e-5 && Math.abs(m[i + 1]) >= 1e-5) { + const alpha = m[i] / d[i], beta = m[i + 1] / d[i], sigma = alpha * alpha + beta * beta; + if (sigma > 9) { const k = 3 / Math.sqrt(sigma); m[i] = k * alpha * d[i]; m[i + 1] = k * beta * d[i]; } + } + } + const tangents = []; + for (let i = 0; i < n; i++) { const denom = 1 + m[i] * m[i]; tangents.push({ x: dx[i] / 3 / denom, y: m[i] * dx[i] / 3 / denom }); } + return tangents; + }, + + monotone(points) { + if (points.length <= 2) return Interpolation.linear(points); + return Interpolation._hermite(points, Interpolation._monotoneTangents(points)); + }, + + linear(points) { let d = ''; for (let i = 1; i < points.length; i++) d += `L${points[i].x},${points[i].y}`; return d; }, + stepBefore(points) { let d = ''; for (let i = 1; i < points.length; i++) d += `V${points[i].y}H${points[i].x}`; return d; }, + stepAfter(points) { let d = ''; for (let i = 1; i < points.length; i++) d += `H${points[i].x}V${points[i].y}`; return d; }, + + bezier(points) { + let d = ''; + for (let i = 1; i < points.length; i++) { + const prev = points[i - 1], cur = points[i]; + d += `C${prev.h2x ?? prev.x},${prev.h2y ?? prev.y},${cur.h1x ?? cur.x},${cur.h1y ?? cur.y},${cur.x},${cur.y}`; + } + return d; + }, + + ensureBezierHandles(points) { + for (let i = 0; i < points.length; i++) { + if (points[i].h1x !== undefined) continue; + const prev = points[Math.max(0, i - 1)], next = points[Math.min(points.length - 1, i + 1)]; + const dx = (next.x - prev.x) * 0.25, dy = (next.y - prev.y) * 0.25; + points[i].h1x = points[i].x - dx; points[i].h1y = points[i].y - dy; + points[i].h2x = points[i].x + dx; points[i].h2y = points[i].y + dy; + } + }, + + buildPathD(points, interpolation, tension) { + if (!points || points.length === 0) return ''; + let d = `M${points[0].x},${points[0].y}`; + if (points.length === 1) return d; + switch (interpolation) { + case 'basis': d += Interpolation.basis(points); break; + case 'cardinal': d += Interpolation.cardinal(points, tension); break; + case 'monotone': d += Interpolation.monotone(points); break; + case 'step-before': d += Interpolation.stepBefore(points); break; + case 'step-after': d += Interpolation.stepAfter(points); break; + case 'bezier': Interpolation.ensureBezierHandles(points); d += Interpolation.bezier(points); break; + case 'linear': default: d += Interpolation.linear(points); break; + } + return d; + }, + + drawOnCanvas(ctx, points, interpolation, tension) { + if (!points || points.length < 2) return; + ctx.beginPath(); + ctx.moveTo(points[0].x, points[0].y); + switch (interpolation) { + case 'basis': + if (points.length < 3) { ctx.lineTo(points[1].x, points[1].y); break; } + Interpolation._drawBasisCanvas(ctx, points); break; + case 'cardinal': + if (points.length < 3) { ctx.lineTo(points[1].x, points[1].y); break; } + Interpolation._drawHermiteCanvas(ctx, points, Interpolation._cardinalTangents(points, tension)); break; + case 'monotone': + if (points.length < 3) { ctx.lineTo(points[1].x, points[1].y); break; } + Interpolation._drawHermiteCanvas(ctx, points, Interpolation._monotoneTangents(points)); break; + case 'step-before': + for (let i = 1; i < points.length; i++) { ctx.lineTo(points[i - 1].x, points[i].y); ctx.lineTo(points[i].x, points[i].y); } break; + case 'step-after': + for (let i = 1; i < points.length; i++) { ctx.lineTo(points[i].x, points[i - 1].y); ctx.lineTo(points[i].x, points[i].y); } break; + case 'bezier': + Interpolation.ensureBezierHandles(points); + for (let i = 1; i < points.length; i++) { + const prev = points[i - 1], cur = points[i]; + ctx.bezierCurveTo(prev.h2x ?? prev.x, prev.h2y ?? prev.y, cur.h1x ?? cur.x, cur.h1y ?? cur.y, cur.x, cur.y); + } break; + case 'linear': default: + for (let i = 1; i < points.length; i++) ctx.lineTo(points[i].x, points[i].y); break; + } + }, + + _drawBasisCanvas(ctx, points) { + const bp = Interpolation._basisPoint; + const drawSeg = (p0, p1, p2, p3) => { + ctx.bezierCurveTo( + bp(0,p0.x,2/3,p1.x,1/3,p2.x,0,p3.x), bp(0,p0.y,2/3,p1.y,1/3,p2.y,0,p3.y), + bp(0,p0.x,1/3,p1.x,2/3,p2.x,0,p3.x), bp(0,p0.y,1/3,p1.y,2/3,p2.y,0,p3.y), + bp(0,p0.x,1/6,p1.x,2/3,p2.x,1/6,p3.x), bp(0,p0.y,1/6,p1.y,2/3,p2.y,1/6,p3.y)); + }; + let b0 = points[0], b1 = points[0], b2 = points[0], b3 = points[1]; + drawSeg(b0, b1, b2, b3); + for (let i = 2; i < points.length; i++) { b0 = b1; b1 = b2; b2 = b3; b3 = points[i]; drawSeg(b0, b1, b2, b3); } + b0 = b1; b1 = b2; b2 = b3; drawSeg(b0, b1, b2, b3); + b0 = b1; b1 = b2; drawSeg(b0, b1, b2, b3); + }, + + _drawHermiteCanvas(ctx, points, tangents) { + if (tangents.length < 1 || (points.length !== tangents.length && points.length !== tangents.length + 2)) return; + const quad = points.length !== tangents.length; + let g = points[0], h = points[1], i = tangents[0], j = i, k = 1; + let prevCp2x, prevCp2y; + if (quad) { ctx.quadraticCurveTo(h.x - i.x * 2/3, h.y - i.y * 2/3, h.x, h.y); g = points[1]; k = 2; } + if (tangents.length > 1) { + j = tangents[1]; h = points[k]; k++; + prevCp2x = h.x - j.x; prevCp2y = h.y - j.y; + ctx.bezierCurveTo(g.x + i.x, g.y + i.y, prevCp2x, prevCp2y, h.x, h.y); + for (let idx = 2; idx < tangents.length; idx++, k++) { + const prevH = h; h = points[k]; j = tangents[idx]; + const rcp1x = 2 * prevH.x - prevCp2x, rcp1y = 2 * prevH.y - prevCp2y; + prevCp2x = h.x - j.x; prevCp2y = h.y - j.y; + ctx.bezierCurveTo(rcp1x, rcp1y, prevCp2x, prevCp2y, h.x, h.y); + } + } + if (quad) { const last = points[k]; ctx.quadraticCurveTo(h.x + j.x * 2/3, h.y + j.y * 2/3, last.x, last.y); } + } +}; + +export { Interpolation }; diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/editors/point_editor_canvas.js b/custom_nodes/ComfyUI-KJNodes/web/js/editors/point_editor_canvas.js new file mode 100644 index 0000000000000000000000000000000000000000..395825992b8599c83bf8c6c71dccbd18e7b3568a --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/editors/point_editor_canvas.js @@ -0,0 +1,370 @@ +import { chainCallback, rectHitTest, cursorForBboxMode } from '../utility.js'; +import { BaseEditorCanvas, createEditorStylesheet } from './editor_base.js'; +const { app } = window.comfyAPI.app; + +createEditorStylesheet('kj-pointseditor-stylesheet', 'points-editor'); + +app.registerExtension({ + name: 'KJNodes.PointEditor', + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name === 'PointsEditor') { + chainCallback(nodeType.prototype, "onNodeCreated", function () { + BaseEditorCanvas.setupNode(this, nodeData, { + editorClass: PointsEditor, + editorKey: 'pointsEditor', + heightKey: 'pointsEditorHeight', + className: 'points-editor', + menuItems: { + "Load Image": { action: (ed) => ed.openImageFilePicker() }, + "Clear Image": { action: (ed) => ed.clearBackgroundImage() }, + }, + hiddenWidgets: ["coordinates", "neg_coordinates", "bboxes"], + initialSize: [550, 550], + extraProperties: [ + ["points", this.constructor.type, "string"], + ["neg_points", this.constructor.type, "string"], + ], + }); + }); + } + } +}); + +class PointsEditor extends BaseEditorCanvas { + constructor(context, reset = false) { + super(context, reset); + this.initEditorPreamble('pointsEditor', 'points-editor'); + + this.pos_coordWidget = this.findWidget("coordinates"); + this.neg_coordWidget = this.findWidget("neg_coordinates"); + this.pointsStoreWidget = this.findWidget("points_store"); + this.widthWidget = this.findWidget("width"); + this.heightWidget = this.findWidget("height"); + this.bboxStoreWidget = this.findWidget("bbox_store"); + this.bboxWidget = this.findWidget("bboxes"); + + this.setupSizeCallbacks(); + this.pointsStoreWidget.callback = () => { + try { + const parsed = JSON.parse(this.pointsStoreWidget.value); + this.points = parsed.positive || []; + this.neg_points = parsed.negative || []; + } catch (e) { console.error("Error parsing points data:", e); } + this.render(); + this.updateData(); + }; + this.bboxStoreWidget.callback = () => { + try { this.bbox = JSON.parse(this.bboxStoreWidget.value); } + catch (e) { console.error("Error parsing bbox data:", e); } + this.render(); + this.updateData(); + }; + + this.initDisplaySize(); + this.points = []; + this.neg_points = []; + this.bbox = [{}]; + this.drawing = false; + this.selectedIndex = -1; + this.pointRadius = Math.max(4, Math.log(Math.max(1, Math.min(this.coordWidth, this.coordHeight))) * 4); + + if (!reset && this.pointsStoreWidget.value !== "") { + try { + const parsed = JSON.parse(this.pointsStoreWidget.value); + this.points = parsed.positive || []; + this.neg_points = parsed.negative || []; + this.bbox = JSON.parse(this.bboxStoreWidget.value); + } catch (e) { + console.error("Error parsing stored points:", e); + this.points = [{ x: this.coordWidth / 2, y: this.coordHeight / 2 }]; + this.neg_points = []; + this.bbox = [{}]; + } + } else { + this.points = [{ x: this.coordWidth / 2, y: this.coordHeight / 2 }]; + this.neg_points = []; + this.pointsStoreWidget.value = JSON.stringify({ positive: this.points, negative: this.neg_points }); + this.bboxStoreWidget.value = JSON.stringify(this.bbox); + } + + this.initEditor('pointsEditor', 'pointsEditorHeight', 310); + this.updateData(); + this.refreshBackgroundImage(); + } + + onDataChanged() { this.updateData(); } + + onCoordSpaceResized(oldWidth, oldHeight) { + const sx = this.coordWidth / oldWidth, sy = this.coordHeight / oldHeight; + const clampPt = (p) => { p.x = Math.max(0, Math.min(this.coordWidth, p.x * sx)); p.y = Math.max(0, Math.min(this.coordHeight, p.y * sy)); }; + for (const p of this.points) clampPt(p); + for (const p of this.neg_points) clampPt(p); + for (const b of this.bbox) { + if (b.startX != null) { + b.startX = Math.max(0, Math.min(this.coordWidth, b.startX * sx)); + b.startY = Math.max(0, Math.min(this.coordHeight, b.startY * sy)); + b.endX = Math.max(0, Math.min(this.coordWidth, b.endX * sx)); + b.endY = Math.max(0, Math.min(this.coordHeight, b.endY * sy)); + } + } + this.pointRadius = Math.max(4, Math.log(Math.max(1, Math.min(this.coordWidth, this.coordHeight))) * 4); + } + + // ─── Hit Testing ─── + + findPointAt(x, y, pointsArray) { + // Convert pixel radius to coord space + const r = this.pointRadius / Math.min(this.scaleX, this.scaleY); + for (let i = pointsArray.length - 1; i >= 0; i--) { + const p = pointsArray[i]; + const dx = p.x - x, dy = p.y - y; + if (dx * dx + dy * dy <= r * r) return i; + } + return -1; + } + + normBbox() { + const b = this.bbox[0]; + if (!b || Object.keys(b).length === 0) return null; + return { + x1: Math.min(b.startX, b.endX), y1: Math.min(b.startY, b.endY), + x2: Math.max(b.startX, b.endX), y2: Math.max(b.startY, b.endY), + }; + } + + hitTestBbox(mx, my) { + const n = this.normBbox(); + if (!n) return null; + return rectHitTest(mx, my, n.x1, n.y1, n.x2, n.y2, 10 / Math.min(this.scaleX, this.scaleY)); + } + + // ─── Mouse Handlers ─── + + onMouseDown(e) { + const mouse = this.getLocalMouse(e); + const clamped = this.clamp(mouse.x, mouse.y); + + if (e.shiftKey && e.button === 2) { + this.neg_points.push({ x: clamped.x, y: clamped.y }); + this.dragType = 'negative'; this.dragIndex = this.neg_points.length - 1; + this.selectedIndex = this.dragIndex; + this.dragOffset = { x: 0, y: 0 }; + this.render(); this.updateData(); this.startDocumentDrag(e); return; + } + if (e.shiftKey && e.button === 0) { + this.points.push({ x: clamped.x, y: clamped.y }); + this.dragType = 'positive'; this.dragIndex = this.points.length - 1; + this.selectedIndex = this.dragIndex; + this.dragOffset = { x: 0, y: 0 }; + this.render(); this.updateData(); this.startDocumentDrag(e); return; + } + + if (e.button === 2) { + const posIdx = this.findPointAt(mouse.x, mouse.y, this.points); + if (posIdx >= 0) { this.points.splice(posIdx, 1); this.render(); this.updateData(); return; } + const negIdx = this.findPointAt(mouse.x, mouse.y, this.neg_points); + if (negIdx >= 0) { this.neg_points.splice(negIdx, 1); this.render(); this.updateData(); return; } + const bboxHit = this.hitTestBbox(mouse.x, mouse.y); + if (bboxHit) { this.bbox = [{}]; this.render(); this.updateData(); return; } + this.showContextMenu(e); + return; + } + + if (e.button !== 0) return; + + if (e.ctrlKey) { + this.drawing = true; + this.dragType = 'bbox-draw'; + this.bbox[0] = { startX: clamped.x, startY: clamped.y, endX: clamped.x, endY: clamped.y }; + this.render(); this.startDocumentDrag(e); return; + } + + const bboxHit = this.hitTestBbox(mouse.x, mouse.y); + if (bboxHit) { + const n = this.normBbox(); + this.bboxAtDragStart = { startX: n.x1, startY: n.y1, endX: n.x2, endY: n.y2 }; + this.dragStart = { x: mouse.x, y: mouse.y }; + this.dragType = bboxHit === "move" ? 'bbox-move' : bboxHit; + this.dragIndex = 0; + this.startDocumentDrag(e); return; + } + + const posIdx = this.findPointAt(mouse.x, mouse.y, this.points); + if (posIdx >= 0) { + this.dragType = 'positive'; this.dragIndex = posIdx; this.selectedIndex = posIdx; + this.dragOffset = { x: mouse.x - this.points[posIdx].x, y: mouse.y - this.points[posIdx].y }; + this.render(); this.startDocumentDrag(e); return; + } + + const negIdx = this.findPointAt(mouse.x, mouse.y, this.neg_points); + if (negIdx >= 0) { + this.dragType = 'negative'; this.dragIndex = negIdx; this.selectedIndex = negIdx; + this.dragOffset = { x: mouse.x - this.neg_points[negIdx].x, y: mouse.y - this.neg_points[negIdx].y }; + this.render(); this.startDocumentDrag(e); return; + } + } + + onMouseMove(e) { + const mouse = this.getLocalMouse(e); + const clamped = this.clamp(mouse.x, mouse.y); + + if (!this.drawing && this.dragIndex < 0 && !this.dragType) { + const bboxHit = this.hitTestBbox(mouse.x, mouse.y); + const bboxCursor = cursorForBboxMode(bboxHit); + if (bboxCursor) { this.canvas.style.cursor = bboxCursor; } + else if (this.findPointAt(mouse.x, mouse.y, this.points) >= 0 || + this.findPointAt(mouse.x, mouse.y, this.neg_points) >= 0) { this.canvas.style.cursor = "move"; } + else { this.canvas.style.cursor = "default"; } + return; + } + + if (this.dragType === 'bbox-draw') { + this.bbox[0].endX = clamped.x; this.bbox[0].endY = clamped.y; + this.render(); return; + } + + if (this.dragType === 'bbox-move' && this.bboxAtDragStart) { + const dx = mouse.x - this.dragStart.x, dy = mouse.y - this.dragStart.y; + const bs = this.bboxAtDragStart; + const bw = bs.endX - bs.startX, bh = bs.endY - bs.startY; + const nx = Math.max(0, Math.min(this.coordWidth - bw, bs.startX + dx)); + const ny = Math.max(0, Math.min(this.coordHeight - bh, bs.startY + dy)); + this.bbox[0] = { startX: nx, startY: ny, endX: nx + bw, endY: ny + bh }; + this.render(); return; + } + + if (this.dragType?.startsWith('resize-') && this.bboxAtDragStart) { + const bs = this.bboxAtDragStart; + const offX = this.dragStart.x - (this.dragType === 'resize-tl' || this.dragType === 'resize-bl' ? bs.startX : bs.endX); + const offY = this.dragStart.y - (this.dragType === 'resize-tl' || this.dragType === 'resize-tr' ? bs.startY : bs.endY); + const cx = Math.max(0, Math.min(this.coordWidth, mouse.x - offX)); + const cy = Math.max(0, Math.min(this.coordHeight, mouse.y - offY)); + if (this.dragType === 'resize-tl') this.bbox[0] = { startX: bs.endX, startY: bs.endY, endX: cx, endY: cy }; + else if (this.dragType === 'resize-tr') this.bbox[0] = { startX: bs.startX, startY: bs.endY, endX: cx, endY: cy }; + else if (this.dragType === 'resize-bl') this.bbox[0] = { startX: bs.endX, startY: bs.startY, endX: cx, endY: cy }; + else this.bbox[0] = { startX: bs.startX, startY: bs.startY, endX: cx, endY: cy }; + this.render(); return; + } + + if (this.dragType === 'positive' && this.dragIndex >= 0 && this.dragIndex < this.points.length) { + const ox = this.dragOffset?.x || 0, oy = this.dragOffset?.y || 0; + const target = this.clamp(mouse.x - ox, mouse.y - oy); + this.points[this.dragIndex] = { x: target.x, y: target.y }; + this.render(); + } else if (this.dragType === 'negative' && this.dragIndex >= 0 && this.dragIndex < this.neg_points.length) { + const ox = this.dragOffset?.x || 0, oy = this.dragOffset?.y || 0; + const target = this.clamp(mouse.x - ox, mouse.y - oy); + this.neg_points[this.dragIndex] = { x: target.x, y: target.y }; + this.render(); + } + } + + onMouseUp() { + if (this.drawing || this.dragType) { + this.drawing = false; + this.bboxAtDragStart = null; + this.dragStart = null; + this.endDrag(); + this.updateData(); + } + } + + // ─── Rendering ─── + + _render() { + const ctx = this.ctx; + this.beginRender(); + + // Bounding box + const nb = this.normBbox(); + if (nb) { + const tl = this.toCanvas(nb.x1, nb.y1), br = this.toCanvas(nb.x2, nb.y2); + const bw = br.x - tl.x, bh = br.y - tl.y; + ctx.fillStyle = 'rgba(70, 130, 180, 0.3)'; + ctx.fillRect(tl.x, tl.y, bw, bh); + ctx.strokeStyle = 'rgba(0, 0, 0, 0.5)'; ctx.lineWidth = 4; + ctx.strokeRect(tl.x, tl.y, bw, bh); + ctx.strokeStyle = 'steelblue'; ctx.lineWidth = 2; + ctx.strokeRect(tl.x, tl.y, bw, bh); + // Center cross + const cx = (tl.x + br.x) / 2, cy = (tl.y + br.y) / 2; + ctx.strokeStyle = 'steelblue'; ctx.lineWidth = 1; + ctx.beginPath(); ctx.moveTo(cx - 5, cy); ctx.lineTo(cx + 5, cy); + ctx.moveTo(cx, cy - 5); ctx.lineTo(cx, cy + 5); ctx.stroke(); + // Corner handles + const hs = 5; + ctx.fillStyle = 'steelblue'; + for (const [hx, hy] of [[tl.x, tl.y], [br.x, tl.y], [tl.x, br.y], [br.x, br.y]]) { + ctx.fillRect(hx - hs, hy - hs, hs * 2, hs * 2); + } + } + + // Points + this._badgePositions = []; + this._drawPoints(ctx, this.points, '#139613', '#07f907', 'positive'); + // Negative points + this._drawPoints(ctx, this.neg_points, '#891616', '#f91111', 'negative'); + + this.endRender(); + } + + _drawPoints(ctx, points, color, selectedColor, type) { + const r = this.pointRadius; + + for (let i = 0; i < points.length; i++) { + const p = points[i]; + const cp = this.toCanvas(p.x, p.y); + const isSelected = (this.dragType === type && this.dragIndex === i); + // Outer circle + ctx.beginPath(); ctx.arc(cp.x, cp.y, r, 0, Math.PI * 2); + ctx.fillStyle = 'rgba(100, 100, 100, 0.25)'; ctx.fill(); + ctx.strokeStyle = isSelected ? selectedColor : color; ctx.lineWidth = 2; ctx.stroke(); + // Center dot + ctx.beginPath(); ctx.arc(cp.x, cp.y, 1.5, 0, Math.PI * 2); + ctx.fillStyle = 'red'; ctx.fill(); + // Index badge + const label = i.toString(); + const badgeR = label.length > 1 ? 8 : 7; + const off = r * 0.7; + const candidates = [ + { x: cp.x + off, y: cp.y - off }, + { x: cp.x - off, y: cp.y - off }, + { x: cp.x + off, y: cp.y + off }, + { x: cp.x - off, y: cp.y + off }, + ]; + const valid = candidates.filter(c => c.x - badgeR >= 0 && c.x + badgeR <= this.width && c.y - badgeR >= 0 && c.y + badgeR <= this.height); + const pool = valid.length > 0 ? valid : candidates; + let badge = pool[0]; + const minDist = badgeR * 2.2; + for (const c of pool) { + const overlaps = this._badgePositions.some(b => Math.abs(c.x - b.x) < minDist && Math.abs(c.y - b.y) < minDist); + if (!overlaps) { badge = c; break; } + } + this._badgePositions.push(badge); + + ctx.font = 'bold 10px sans-serif'; + ctx.textAlign = 'center'; ctx.textBaseline = 'middle'; + ctx.beginPath(); ctx.arc(badge.x, badge.y, badgeR, 0, Math.PI * 2); + ctx.fillStyle = isSelected ? selectedColor : color; ctx.fill(); + ctx.fillStyle = '#fff'; + ctx.fillText(label, badge.x, badge.y + 0.5); + ctx.textAlign = 'start'; ctx.textBaseline = 'alphabetic'; + } + } + + // ─── Data ─── + + updateData() { + const combinedPoints = { positive: this.points || [], negative: this.neg_points || [] }; + this.pointsStoreWidget.value = JSON.stringify(combinedPoints); + this.pos_coordWidget.value = JSON.stringify(this.points || []); + this.neg_coordWidget.value = JSON.stringify(this.neg_points || []); + if (this.bbox.length !== 0) { + let bboxString = JSON.stringify(this.bbox); + this.bboxStoreWidget.value = bboxString; + this.bboxWidget.value = bboxString; + } + } + +} diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/editors/spline_editor_canvas.js b/custom_nodes/ComfyUI-KJNodes/web/js/editors/spline_editor_canvas.js new file mode 100644 index 0000000000000000000000000000000000000000..00c943de6c965a5416c3efc4ffcda50493935f69 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/editors/spline_editor_canvas.js @@ -0,0 +1,720 @@ +import { chainCallback } from '../utility.js'; +import { BaseEditorCanvas, createEditorStylesheet } from './editor_base.js'; +import { Interpolation } from './interpolation.js'; +const { app } = window.comfyAPI.app; + +createEditorStylesheet('kj-splineeditor-stylesheet', 'spline-editor'); + +// ─── Hidden SVG for path sampling (singleton) ─── +let _sharedSampler = null; +let _samplerRefCount = 0; +class PathSampler { + constructor() { + const ns = 'http://www.w3.org/2000/svg'; + this._svg = document.createElementNS(ns, 'svg'); + this._svg.style.cssText = 'position:absolute;width:0;height:0;overflow:hidden;pointer-events:none;'; + this._svg.setAttribute('aria-hidden', 'true'); + document.body.appendChild(this._svg); + this._path = document.createElementNS(ns, 'path'); + this._svg.appendChild(this._path); + } + static acquire() { + if (_sharedSampler) { _samplerRefCount++; return _sharedSampler; } + _sharedSampler = new PathSampler(); + _samplerRefCount = 1; + return _sharedSampler; + } + setPath(d) { this._path.setAttribute('d', d); } + getTotalLength() { return this._path.getTotalLength(); } + getPointAtLength(len) { return this._path.getPointAtLength(len); } + release() { + _samplerRefCount--; + if (_samplerRefCount <= 0 && this._svg && this._svg.parentNode) { + this._svg.parentNode.removeChild(this._svg); + _sharedSampler = null; + } + } + findPointAtX(targetX, pathLength) { + let low = 0, high = pathLength, bestPoint = this.getPointAtLength(0); + const epsilon = 0.5; + while (high - low > epsilon) { + const mid = (low + high) / 2, point = this.getPointAtLength(mid); + if (Math.abs(point.x - targetX) < Math.abs(bestPoint.x - targetX)) bestPoint = point; + if (Math.abs(point.x - targetX) < 0.5) return point; + if (point.x < targetX) low = mid; else high = mid; + } + return bestPoint; + } +} + +// ─── Extension Registration ─── +app.registerExtension({ + name: 'KJNodes.SplineEditor', + + init() { + app.ui.settings.addSetting({ + id: "KJNodes.splineEditor.showControlLines", + name: "Show control lines", + category: ["KJNodes", "Editors", "Show control lines"], + tooltip: "Display straight lines between control points", + type: "boolean", + defaultValue: false, + }); + app.ui.settings.addSetting({ + id: "KJNodes.splineEditor.showSamplePoints", + name: "Show sample points", + category: ["KJNodes", "Editors", "Show sample points"], + tooltip: "Display the sampled output points as red dots", + type: "boolean", + defaultValue: false, + }); + app.ui.settings.addSetting({ + id: "KJNodes.splineEditor.showArrows", + name: "Show direction arrows", + category: ["KJNodes", "Editors", "Show direction arrows"], + tooltip: "Display chevron arrows on control points showing curve direction", + type: "boolean", + defaultValue: true, + }); + app.ui.settings.addSetting({ + id: "KJNodes.editors.embedBackgroundImage", + name: "Embed background image in workflow", + category: ["KJNodes", "Editors", "Embed background image"], + tooltip: "When enabled, editor background images are embedded as base64 in the workflow file (portable, larger files). When disabled, images are stored as temp files on the server (small workflows, survives refresh, but not portable).", + type: "boolean", + defaultValue: false, + }); + }, + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name === 'SplineEditor') { + chainCallback(nodeType.prototype, "onNodeCreated", function () { + BaseEditorCanvas.setupNode(this, nodeData, { + editorClass: SplineEditor, + editorKey: 'splineEditor', + heightKey: 'splineEditorHeight', + className: 'spline-editor', + menuClassName: 'spline-editor-context-menu', + menuItems: { + "Display control lines": { toggle: (ed) => ed.drawHandles, action: (ed) => { ed.drawHandles = !ed.drawHandles; app.ui.settings.setSettingValue("KJNodes.splineEditor.showControlLines", ed.drawHandles); ed.render(); } }, + "Display sample points": { toggle: (ed) => ed.drawSamplePoints, action: (ed) => { ed.drawSamplePoints = !ed.drawSamplePoints; app.ui.settings.setSettingValue("KJNodes.splineEditor.showSamplePoints", ed.drawSamplePoints); ed.updatePath(); } }, + "Display arrows": { toggle: (ed) => ed.showChevrons, action: (ed) => { ed.showChevrons = !ed.showChevrons; app.ui.settings.setSettingValue("KJNodes.splineEditor.showArrows", ed.showChevrons); ed.render(); } }, + "Background image": { action: (ed) => ed.openImageFilePicker() }, + "Invert point order": { action: (ed) => { ed.splines[ed.activeSplineIndex].points.reverse(); ed.render(); ed.updatePath(); } }, + "Clear Image": { action: (ed) => { ed.clearBackgroundImage(); ed.drawRuler = true; } }, + "Add new spline": { action: (ed) => { const idx = ed.splines.length; ed.splines.push({ points: [{ x: 0, y: ed.coordHeight }, { x: ed.coordWidth / 2, y: ed.coordHeight / 2 }, { x: ed.coordWidth, y: 0 }], color: ed.getSplineColor(idx), name: `Spline ${idx + 1}` }); ed.activeSplineIndex = idx; ed.render(); ed.updatePath(); } }, + "Add new single point": { action: (ed) => { const idx = ed.splines.length; ed.splines.push({ points: [{ x: ed.lastContextMenuPos.x, y: ed.lastContextMenuPos.y }], color: ed.getSplineColor(idx), name: `Spline ${idx + 1}`, isSinglePoint: true }); ed.activeSplineIndex = idx; ed.render(); ed.updatePath(); } }, + "Delete current spline": { action: (ed) => { if (ed.splines.length > 1) { ed.splines.splice(ed.activeSplineIndex, 1); ed.activeSplineIndex = Math.min(ed.activeSplineIndex, ed.splines.length - 1); ed.render(); ed.updatePath(); } } }, + "Next spline": { action: (ed) => { ed.activeSplineIndex = (ed.activeSplineIndex + 1) % ed.splines.length; ed.render(); ed.updatePath(); } }, + }, + hiddenWidgets: ["coordinates"], + initialSize: [550, 1000], + extraProperties: [ + ["points", this.constructor.type, "string"], + ], + }); + }); + } + } +}); + +// ─── SplineEditor class ─── +class SplineEditor extends BaseEditorCanvas { + constructor(context, reset = false) { + super(context, reset); + this.initEditorPreamble('splineEditor', 'spline-editor'); + + this.sampler = PathSampler.acquire(); + this.drawSamplePoints = app.ui.settings.getSettingValue("KJNodes.splineEditor.showSamplePoints") ?? false; + this.drawHandles = app.ui.settings.getSettingValue("KJNodes.splineEditor.showControlLines") ?? false; + this.showChevrons = app.ui.settings.getSettingValue("KJNodes.splineEditor.showArrows") ?? true; + this.drawRuler = true; + + this.coordWidget = this.findWidget("coordinates"); + this.interpolationWidget = this.findWidget("interpolation"); + this.pointsWidget = this.findWidget("points_to_sample"); + this.pointsStoreWidget = this.findWidget("points_store"); + this.tensionWidget = this.findWidget("tension"); + this.samplingMethodWidget = this.findWidget("sampling_method"); + this.widthWidget = this.findWidget("mask_width"); + this.heightWidget = this.findWidget("mask_height"); + + this.interpolationWidget.callback = () => this.updatePath(true); + this.samplingMethodWidget.callback = () => { + if (this.samplingMethod === "controlpoints") this.drawSamplePoints = true; + else this.drawSamplePoints = false; + if (this.samplingMethod === "path" || this.samplingMethod === "speed") this.showChevrons = true; + this.updatePath(true); + }; + this.tensionWidget.callback = () => this.updatePath(true); + this.pointsWidget.callback = () => this.updatePath(true); + this.setupSizeCallbacks(); + this.pointsStoreWidget.callback = () => { + this.parseSplineData(); + this.render(); this.updatePath(true); + }; + + this.initDisplaySize(); + this.splines = []; + this.activeSplineIndex = 0; + this.hoverSplineIndex = -1; + this.hoverIndex = -1; + this.sampledCoords = null; + this.lastContextMenuPos = { x: this.coordWidth / 2, y: this.coordHeight / 2 }; + + if (!reset && this.pointsStoreWidget.value !== "") { + this.parseSplineData(); + } else { + this.initializeDefaultSplines(); + this.pointsStoreWidget.value = JSON.stringify(this.splines); + } + + this._onKeyUp = (e) => { + if (e.key === 'Control' && this.subdividePreview) { this.subdividePreview = null; this.render(); } + }; + document.addEventListener('keyup', this._onKeyUp); + + this.initEditor('splineEditor', 'splineEditorHeight', 460); + this.updatePath(); + this.refreshBackgroundImage(); + } + + // ─── Widget value getters ─── + get interpolation() { return this.interpolationWidget.value; } + get tension() { return this.tensionWidget.value; } + get points_to_sample() { return this.pointsWidget.value; } + get samplingMethod() { return this.samplingMethodWidget.value; } + + // ─── Base class hooks ─── + + destroy() { + super.destroy(); + if (this.sampler) this.sampler.release(); + } + + onDataChanged() { this.updatePath(true); } + + onImageResize() { this.drawRuler = false; } + + onCoordSpaceResized(oldWidth, oldHeight) { + const sx = this.coordWidth / oldWidth, sy = this.coordHeight / oldHeight; + const cw = this.coordWidth, ch = this.coordHeight; + const clamp = (v, max) => Math.max(0, Math.min(max, v)); + for (const spline of this.splines) { + for (const p of spline.points) { + p.x = clamp(p.x * sx, cw); p.y = clamp(p.y * sy, ch); + if (p.h1x !== undefined) { p.h1x = clamp(p.h1x * sx, cw); p.h1y = clamp(p.h1y * sy, ch); } + if (p.h2x !== undefined) { p.h2x = clamp(p.h2x * sx, cw); p.h2y = clamp(p.h2y * sy, ch); } + } + } + } + + // ─── Spline Helpers ─── + + parseSplineData() { + try { + const parsed = JSON.parse(this.pointsStoreWidget.value); + if (Array.isArray(parsed) && parsed.length > 0 && parsed[0].hasOwnProperty('points')) { + this.splines = parsed; + } else { + this.splines = [{ points: parsed, color: "#1f77b4", name: "Spline 1" }]; + } + } catch (e) { + console.error("Error parsing spline data:", e); + this.initializeDefaultSplines(); + } + if (this.activeSplineIndex >= this.splines.length) { + this.activeSplineIndex = Math.max(0, this.splines.length - 1); + } + } + + initializeDefaultSplines() { + this.splines = [{ + points: [ + { x: 0, y: this.coordHeight }, + { x: this.coordWidth / 5 * 2, y: 50 + Math.random() * (this.coordHeight - 100) }, + { x: this.coordWidth, y: 0 } + ], + color: this.getSplineColor(0), + name: "Spline 1" + }]; + } + + getSplineColor(index) { + const colors = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf"]; + return colors[index % colors.length]; + } + + // ─── Hit Testing ─── + + findPointAt(x, y) { + const activePoints = this.splines[this.activeSplineIndex]?.points; + if (!activePoints) return -1; + const r = 12 / Math.min(this.scaleX, this.scaleY); + for (let i = activePoints.length - 1; i >= 0; i--) { + const dx = activePoints[i].x - x, dy = activePoints[i].y - y; + if (dx * dx + dy * dy <= r * r) return i; + } + return -1; + } + + findHandleAt(x, y) { + if (this.interpolation !== 'bezier') return null; + const activePoints = this.splines[this.activeSplineIndex]?.points; + if (!activePoints) return null; + const r = 8 / Math.min(this.scaleX, this.scaleY); + for (let i = activePoints.length - 1; i >= 0; i--) { + const p = activePoints[i]; + if (i > 0 && p.h1x !== undefined) { const dx = p.h1x - x, dy = p.h1y - y; if (dx * dx + dy * dy <= r * r) return { index: i, type: 'h1' }; } + if (i < activePoints.length - 1 && p.h2x !== undefined) { const dx = p.h2x - x, dy = p.h2y - y; if (dx * dx + dy * dy <= r * r) return { index: i, type: 'h2' }; } + } + return null; + } + + findSplineAt(x, y) { + const threshold = 15 / Math.min(this.scaleX, this.scaleY); + let bestDist = Infinity, bestIdx = -1; + for (let si = 0; si < this.splines.length; si++) { + const spline = this.splines[si]; + if (spline.isSinglePoint || (spline.points && spline.points.length === 1)) { + const p = spline.points[0], d = Math.sqrt((p.x - x) ** 2 + (p.y - y) ** 2); + if (d < threshold && d < bestDist) { bestDist = d; bestIdx = si; } + continue; + } + const pathD = Interpolation.buildPathD(spline.points, this.interpolation, this.tension); + this.sampler.setPath(pathD); + const len = this.sampler.getTotalLength(); + const steps = Math.min(200, Math.max(50, Math.round(len / 3))); + for (let i = 0; i <= steps; i++) { + const pt = this.sampler.getPointAtLength(len * i / steps); + const d = Math.sqrt((pt.x - x) ** 2 + (pt.y - y) ** 2); + if (d < threshold && d < bestDist) { bestDist = d; bestIdx = si; } + } + } + return bestIdx; + } + + findClosestPoints(points, clickedPoint) { + if (points.length < 2) return null; + let distances = points.map((point, idx) => { + const dx = clickedPoint.x - point.x, dy = clickedPoint.y - point.y; + return { index: idx, distance: Math.sqrt(dx * dx + dy * dy) }; + }); + distances.sort((a, b) => a.distance - b.distance); + let i1 = distances[0].index, i2 = distances[1].index; + if (i1 > i2) [i1, i2] = [i2, i1]; + return { point1Index: i1, point2Index: i2 }; + } + + // ─── Mouse Handlers ─── + + onMouseDown(e) { + const mouse = this.getLocalMouse(e); + const clamped = this.clamp(mouse.x, mouse.y); + const activeSpline = this.splines[this.activeSplineIndex]; + const activePoints = activeSpline?.points; + + if (e.shiftKey && e.button === 0 && activePoints) { + activePoints.push({ x: clamped.x, y: clamped.y }); + this.render(); this.updatePath(); return; + } + if (e.ctrlKey && e.button === 0 && activePoints && activePoints.length >= 2) { + const closest = this.findClosestPoints(activePoints, clamped); + if (closest) { + const p1 = activePoints[closest.point1Index], p2 = activePoints[closest.point2Index]; + activePoints.splice(closest.point2Index, 0, { x: (p1.x + p2.x) / 2, y: (p1.y + p2.y) / 2 }); + this.render(); this.updatePath(); + } + return; + } + + if (e.button === 2) { + const ptIdx = this.findPointAt(mouse.x, mouse.y); + if (activePoints && ptIdx > 0 && ptIdx < activePoints.length - 1) { + activePoints.splice(ptIdx, 1); this.render(); this.updatePath(); return; + } + this.lastContextMenuPos = { x: clamped.x, y: clamped.y }; + this.showContextMenu(e); + return; + } + + if (e.button !== 0) return; + + // Bezier handle drag + const handle = this.findHandleAt(mouse.x, mouse.y); + if (handle && activePoints) { + this.dragIndex = handle.index; this.dragType = handle.type; + const p = activePoints[handle.index]; + const hx = handle.type === 'h1' ? p.h1x : p.h2x, hy = handle.type === 'h1' ? p.h1y : p.h2y; + this.dragOffset = { x: mouse.x - hx, y: mouse.y - hy }; + this.render(); this.startDocumentDrag(e); return; + } + + // Point drag + const ptIdx = this.findPointAt(mouse.x, mouse.y); + if (ptIdx >= 0) { + this.dragIndex = ptIdx; this.dragType = 'point'; this.hoverIndex = ptIdx; + const p = activePoints[ptIdx]; + this.dragOffset = { x: mouse.x - p.x, y: mouse.y - p.y }; + this.render(); this.startDocumentDrag(e); return; + } + + // Spline selection + const splineIdx = this.findSplineAt(mouse.x, mouse.y); + if (splineIdx >= 0 && splineIdx !== this.activeSplineIndex) { + this.activeSplineIndex = splineIdx; this.render(); this.updatePath(); return; + } + } + + onMouseMove(e) { + const mouse = this.getLocalMouse(e); + const clamped = this.clamp(mouse.x, mouse.y); + + // Bezier handle drag + if ((this.dragType === 'h1' || this.dragType === 'h2') && this.dragIndex >= 0) { + const activePoints = this.splines[this.activeSplineIndex]?.points; + if (activePoints && this.dragIndex < activePoints.length) { + const p = activePoints[this.dragIndex]; + const clampH = (hx, hy) => ({ x: Math.max(0, Math.min(this.coordWidth, hx)), y: Math.max(0, Math.min(this.coordHeight, hy)) }); + const ox = this.dragOffset?.x || 0, oy = this.dragOffset?.y || 0; + const prevH1x = p.h1x, prevH1y = p.h1y, prevH2x = p.h2x, prevH2y = p.h2y; + if (this.dragType === 'h1') { + const c = clampH(mouse.x - ox, mouse.y - oy); + p.h1x = c.x; p.h1y = c.y; + if (!e.altKey) { const dx = p.x - p.h1x, dy = p.y - p.h1y; const mc = clampH(p.x + dx, p.y + dy); p.h2x = mc.x; p.h2y = mc.y; } + } else { + const c = clampH(mouse.x - ox, mouse.y - oy); + p.h2x = c.x; p.h2y = c.y; + if (!e.altKey) { const dx = p.x - p.h2x, dy = p.y - p.h2y; const mc = clampH(p.x + dx, p.y + dy); p.h1x = mc.x; p.h1y = mc.y; } + } + if (p.h1x === prevH1x && p.h1y === prevH1y && p.h2x === prevH2x && p.h2y === prevH2y) return; + this.render(); + } + return; + } + + // Point drag + if (this.dragType === 'point' && this.dragIndex >= 0) { + const activePoints = this.splines[this.activeSplineIndex]?.points; + if (activePoints && this.dragIndex < activePoints.length) { + const p = activePoints[this.dragIndex]; + const ox = this.dragOffset?.x || 0, oy = this.dragOffset?.y || 0; + const target = this.clamp(mouse.x - ox, mouse.y - oy); + if (target.x === p.x && target.y === p.y) return; + const dx = target.x - p.x, dy = target.y - p.y; + if (this.interpolation === 'bezier' && p.h1x !== undefined) { + p.h1x += dx; p.h1y += dy; p.h2x += dx; p.h2y += dy; + } + p.x = target.x; p.y = target.y; + this.render(); + } + return; + } + + // Hover detection + if (!this.dragType) { + const ptIdx = this.findPointAt(mouse.x, mouse.y); + const activePoints = this.splines[this.activeSplineIndex]?.points; + + // Ctrl preview + if (e.ctrlKey && activePoints && activePoints.length >= 2) { + const closest = this.findClosestPoints(activePoints, clamped); + if (closest) { + const p1 = activePoints[closest.point1Index], p2 = activePoints[closest.point2Index]; + this.subdividePreview = { x: (p1.x + p2.x) / 2, y: (p1.y + p2.y) / 2 }; + this.canvas.style.cursor = 'copy'; this.render(); return; + } + } else if (this.subdividePreview) { this.subdividePreview = null; this.render(); } + + // Handle hover + const handleHit = this.findHandleAt(mouse.x, mouse.y); + if (handleHit) { this.canvas.style.cursor = 'crosshair'; if (ptIdx !== this.hoverIndex) { this.hoverIndex = ptIdx; this.render(); } return; } + + if (ptIdx !== this.hoverIndex) { + this.hoverIndex = ptIdx; + this.canvas.style.cursor = ptIdx >= 0 ? 'move' : 'default'; + this.render(); + } + + // Spline hover + if (ptIdx < 0 && this.splines.length > 1) { + const splineIdx = this.findSplineAt(mouse.x, mouse.y); + const newSplineHover = splineIdx >= 0 && splineIdx !== this.activeSplineIndex ? splineIdx : -1; + if (newSplineHover !== this.hoverSplineIndex) { + this.hoverSplineIndex = newSplineHover; + this.canvas.style.cursor = newSplineHover >= 0 ? 'pointer' : 'default'; + this.render(); + } + } else if (this.hoverSplineIndex >= 0) { this.hoverSplineIndex = -1; this.render(); } + } + } + + onMouseUp() { + if (this.dragType) { + this.endDrag(); + this.updatePath(); + } + } + + // ─── Path Building & Sampling ─── + + samplePoints(splineIndex, numSamples, samplingMethod) { + const spline = this.splines[splineIndex]; + if (!spline || !spline.points || spline.points.length < 1) return []; + if (spline.isSinglePoint || spline.points.length === 1) { + const point = spline.points[0]; + return Array(numSamples).fill().map(() => ({ x: point.x, y: point.y })); + } + if (numSamples < 2) { + const p = spline.points[0]; + return [{ x: p.x, y: p.y }]; + } + const pathD = Interpolation.buildPathD(spline.points, this.interpolation, this.tension); + if (!pathD) return []; + + this.sampler.setPath(pathD); + const pathLength = this.sampler.getTotalLength(); + const points = []; + + if (samplingMethod === "speed") { + if (pathLength < 0.001) return [{ x: spline.points[0].x, y: spline.points[0].y }]; + const controlPoints = spline.points; + // Pair each control point with its path position, then sort together + const cpWithPos = controlPoints.map(cp => { + let bestDist = Infinity, bestPos = 0; + for (let pos = 0; pos <= pathLength; pos += pathLength / 100) { + const pt = this.sampler.getPointAtLength(pos); + const dist = Math.sqrt((pt.x - cp.x) ** 2 + (pt.y - cp.y) ** 2); + if (dist < bestDist) { bestDist = dist; bestPos = pos; } + } + return { cp, pos: bestPos }; + }); + cpWithPos.sort((a, b) => a.pos - b.pos); + const pathPositions = cpWithPos.map(c => c.pos); + + const densities = []; + let totalWeight = 0; + for (let i = 0; i < cpWithPos.length - 1; i++) { + const segLength = pathPositions[i + 1] - pathPositions[i]; + const d = 1 / Math.max(segLength, 0.0001); + densities.push(d); + totalWeight += d; + } + const cumulativeWeights = []; + let cum = 0; + for (let i = 0; i < densities.length; i++) { + cum += densities[i] / totalWeight; + cumulativeWeights.push(cum); + } + const mapToPath = (t) => { + if (t === 0) return 0; if (t === 1) return pathLength; + let segIdx = cumulativeWeights.length - 1; + for (let i = 0; i < cumulativeWeights.length; i++) { if (t <= cumulativeWeights[i]) { segIdx = i; break; } } + const segStart = segIdx > 0 ? cumulativeWeights[segIdx - 1] : 0; + const segT = (t - segStart) / (cumulativeWeights[segIdx] - segStart); + return pathPositions[segIdx] + segT * (pathPositions[segIdx + 1] - pathPositions[segIdx]); + }; + for (let i = 0; i < numSamples; i++) { + const pt = this.sampler.getPointAtLength(mapToPath(i / (numSamples - 1))); + points.push({ x: pt.x, y: pt.y }); + } + return points; + } + + for (let i = 0; i < numSamples; i++) { + let point; + if (samplingMethod === "time") { + point = this.sampler.findPointAtX((this.coordWidth / (numSamples - 1)) * i, pathLength); + } else { + point = this.sampler.getPointAtLength((pathLength / (numSamples - 1)) * i); + } + points.push({ x: point.x, y: point.y }); + } + if (points.length > 0 && spline.points.length > 1) { + points[points.length - 1].y = spline.points[spline.points.length - 1].y; + } + return points; + } + + updatePath(allDirty = false) { + if (!this.splines || this.splines.length === 0) return; + + // Rebuild cache when spline count changes or full refresh requested + if (!this._sampledCache || this._sampledCache.length !== this.splines.length || allDirty) { + this._sampledCache = new Array(this.splines.length).fill(null); + } + + const method = this.samplingMethod; + const useControlPoints = method === "controlpoints"; + for (let i = 0; i < this.splines.length; i++) { + if (useControlPoints) { + this._sampledCache[i] = this.splines[i].points; + } else if (i === this.activeSplineIndex || !this._sampledCache[i]) { + this._sampledCache[i] = this.samplePoints(i, this.points_to_sample, method); + } + } + this.sampledCoords = this._sampledCache[this.activeSplineIndex] || []; + this.pointsStoreWidget.value = JSON.stringify(this.splines); + if (this.coordWidget) this.coordWidget.value = JSON.stringify(this._sampledCache); + this.render(); + } + + // ─── Rendering ─── + + _render() { + const ctx = this.ctx; + const h = this.coordHeight; + this.beginRender(); + + // Ruler lines + if (this.drawRuler && !this.bgImage) { + ctx.strokeStyle = 'rgba(128,128,128,0.3)'; ctx.lineWidth = 1; + for (let y = 64; y < h; y += 64) { + const cy = y * this.scaleY; + ctx.beginPath(); ctx.moveTo(0, cy); ctx.lineTo(this.width, cy); ctx.stroke(); + } + } + + const tc = (x, y) => this.toCanvas(x, y); + + // Draw all splines — use scale transform for curves, toCanvas for discrete elements + for (let si = 0; si < this.splines.length; si++) { + const spline = this.splines[si]; + const isActive = si === this.activeSplineIndex, isHover = si === this.hoverSplineIndex; + + if (spline.isSinglePoint || (spline.points && spline.points.length === 1)) { + const cp = tc(spline.points[0].x, spline.points[0].y), sz = isActive ? 8 : 6; + ctx.fillStyle = spline.color; ctx.fillRect(cp.x - sz, cp.y - sz, sz * 2, sz * 2); + ctx.strokeStyle = 'black'; ctx.lineWidth = 2; ctx.strokeRect(cp.x - sz, cp.y - sz, sz * 2, sz * 2); + continue; + } + if (spline.points.length < 2) continue; + + // Build path once, stroke twice (outline + color) to avoid recomputing interpolation + const pathD = Interpolation.buildPathD(spline.points, this.interpolation, this.tension); + const curvePath = new Path2D(pathD); + ctx.save(); ctx.scale(this.scaleX, this.scaleY); + ctx.lineWidth = (isActive ? 5 : isHover ? 4 : 3.5) / Math.min(this.scaleX, this.scaleY); ctx.strokeStyle = 'black'; + ctx.stroke(curvePath); + ctx.lineWidth = (isActive ? 3 : isHover ? 2 : 1.5) / Math.min(this.scaleX, this.scaleY); ctx.strokeStyle = spline.color; + ctx.stroke(curvePath); + ctx.restore(); + } + + const activeSpline = this.splines[this.activeSplineIndex]; + + // Handle lines + if (this.drawHandles) { + const activePoints = this.splines[this.activeSplineIndex]?.points; + if (activePoints && activePoints.length >= 2) { + ctx.strokeStyle = '#ff7f0e'; ctx.lineWidth = 1; ctx.beginPath(); + const c0 = tc(activePoints[0].x, activePoints[0].y); + ctx.moveTo(c0.x, c0.y); + for (let i = 1; i < activePoints.length; i++) { const ci = tc(activePoints[i].x, activePoints[i].y); ctx.lineTo(ci.x, ci.y); } + ctx.stroke(); + } + } + + // Control points for active spline + if (activeSpline && activeSpline.points && !activeSpline.isSinglePoint) { + const dotRadius = 12; + for (let i = 0; i < activeSpline.points.length; i++) { + const p = activeSpline.points[i]; + const cp = tc(p.x, p.y); + const isHovered = this.hoverIndex === i; + ctx.fillStyle = 'rgba(100, 100, 100, 0.3)'; + ctx.strokeStyle = isHovered ? '#ff7f0e' : '#1f77b4'; ctx.lineWidth = 2; + ctx.beginPath(); ctx.arc(cp.x, cp.y, dotRadius, 0, Math.PI * 2); ctx.fill(); ctx.stroke(); + + // Chevron + const isEndpoint = i === 0 || i === activeSpline.points.length - 1; + const showChevron = this.showChevrons && ( + (this.samplingMethod === 'path' || this.samplingMethod === 'speed') + || (this.samplingMethod === 'time' && isEndpoint)); + + if (showChevron && activeSpline.points.length > 1) { + let angle = 0; + if (i > 0 && i < activeSpline.points.length - 1) { const prev = activeSpline.points[i-1], next = activeSpline.points[i+1]; angle = Math.atan2((next.y - prev.y) * this.scaleY, (next.x - prev.x) * this.scaleX); } + else if (i === 0) { const next = activeSpline.points[1]; angle = Math.atan2((next.y - p.y) * this.scaleY, (next.x - p.x) * this.scaleX); } + else { const prev = activeSpline.points[i-1]; angle = Math.atan2((p.y - prev.y) * this.scaleY, (p.x - prev.x) * this.scaleX); } + ctx.save(); + ctx.translate(cp.x, cp.y); + ctx.rotate(angle); + ctx.strokeStyle = isHovered ? '#fff' : 'rgba(255,255,255,0.7)'; + ctx.lineWidth = 2.5; ctx.lineCap = 'round'; ctx.lineJoin = 'round'; + ctx.beginPath(); + ctx.moveTo(-dotRadius * 0.2, -dotRadius * 0.35); + ctx.lineTo(dotRadius * 0.3, 0); + ctx.lineTo(-dotRadius * 0.2, dotRadius * 0.35); + ctx.stroke(); + ctx.lineCap = 'butt'; ctx.lineJoin = 'miter'; + ctx.restore(); + } + + if (isHovered) { + ctx.font = '11px monospace'; + const ax = String(Math.round(p.x)), ay = String(Math.round(p.y)); + const nx = (p.x / this.coordWidth).toFixed(3), ny = (p.y / this.coordHeight).toFixed(3); + const ml = Math.max(ax.length, nx.length), mr = Math.max(ay.length, ny.length); + const pre = `${i}: `, L1 = `[${ax.padStart(ml)}, ${ay.padStart(mr)}]`, L2 = `[${nx.padStart(ml)}, ${ny.padStart(mr)}]`; + const preW = ctx.measureText(pre).width, totW = preW + Math.max(ctx.measureText(L1).width, ctx.measureText(L2).width); + const pad = dotRadius + 5, lineH = 13; + let lx = cp.x + pad, ly = cp.y - lineH; + if (lx + totW > this.width) lx = cp.x - pad - totW; + if (ly - lineH < 0) ly = cp.y + pad; + for (const [style, fill] of [['#000', 3], ['#fff', 0]]) { + ctx.strokeStyle = style; ctx.lineWidth = fill; ctx.fillStyle = style; + const draw = fill ? 'strokeText' : 'fillText'; + ctx[draw](pre + L1, lx, ly); + ctx[draw](L2, lx + preW, ly + lineH); + } + } + } + } + + // Bezier handles + if (this.interpolation === 'bezier' && activeSpline && activeSpline.points && !activeSpline.isSinglePoint) { + const pts = activeSpline.points; + Interpolation.ensureBezierHandles(pts); + const hr = 5; + for (let i = 0; i < pts.length; i++) { + const p = pts[i]; + const cp = tc(p.x, p.y); + if (i > 0 && p.h1x !== undefined) { + const h1 = tc(p.h1x, p.h1y); + ctx.strokeStyle = 'rgba(255,255,255,0.4)'; ctx.lineWidth = 1; + ctx.beginPath(); ctx.moveTo(cp.x, cp.y); ctx.lineTo(h1.x, h1.y); ctx.stroke(); + ctx.fillStyle = (this.dragType === 'h1' && this.dragIndex === i) ? '#fff' : 'rgba(255,255,255,0.8)'; + ctx.strokeStyle = '#1f77b4'; ctx.lineWidth = 1.5; + ctx.beginPath(); ctx.arc(h1.x, h1.y, hr, 0, Math.PI * 2); ctx.fill(); ctx.stroke(); + } + if (i < pts.length - 1 && p.h2x !== undefined) { + const h2 = tc(p.h2x, p.h2y); + ctx.strokeStyle = 'rgba(255,255,255,0.4)'; ctx.lineWidth = 1; + ctx.beginPath(); ctx.moveTo(cp.x, cp.y); ctx.lineTo(h2.x, h2.y); ctx.stroke(); + ctx.fillStyle = (this.dragType === 'h2' && this.dragIndex === i) ? '#fff' : 'rgba(255,255,255,0.8)'; + ctx.strokeStyle = '#ff7f0e'; ctx.lineWidth = 1.5; + ctx.beginPath(); ctx.arc(h2.x, h2.y, hr, 0, Math.PI * 2); ctx.fill(); ctx.stroke(); + } + } + } + + // Sample points + if (this.drawSamplePoints && this.sampledCoords) { + ctx.fillStyle = 'red'; ctx.strokeStyle = 'black'; ctx.lineWidth = 1; + for (const pt of this.sampledCoords) { const cp = tc(pt.x, pt.y); ctx.beginPath(); ctx.arc(cp.x, cp.y, 5, 0, Math.PI * 2); ctx.fill(); ctx.stroke(); } + } + + // Subdivide preview + if (this.subdividePreview) { + const sp = tc(this.subdividePreview.x, this.subdividePreview.y); + ctx.beginPath(); ctx.arc(sp.x, sp.y, 10, 0, Math.PI * 2); + ctx.strokeStyle = 'rgba(0,255,0,0.3)'; ctx.lineWidth = 4; ctx.stroke(); + ctx.beginPath(); ctx.arc(sp.x, sp.y, 8, 0, Math.PI * 2); + ctx.strokeStyle = '#0f0'; ctx.lineWidth = 2; ctx.stroke(); + } + + this.endRender(); + } + +} diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/fast_preview.js b/custom_nodes/ComfyUI-KJNodes/web/js/fast_preview.js new file mode 100644 index 0000000000000000000000000000000000000000..15c99d6d7dfa7570e365390cfff3cb81dfbc8f2f --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/fast_preview.js @@ -0,0 +1,40 @@ +import { chainCallback } from './utility.js'; +const { app } = window.comfyAPI.app; +const { api } = window.comfyAPI.api; + +let execId = null; +api.addEventListener("executing", e => { execId = e.detail ?? null; }); + +app.registerExtension({ + name: 'KJNodes.FastPreview', + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== 'FastPreview') return; + chainCallback(nodeType.prototype, "onNodeCreated", function () { + this.setSize([550, 550]); + const nodeRef = this; + + const show = blob => { + const img = new Image(); + img.onload = () => { nodeRef.imgs = [img]; nodeRef.setDirtyCanvas(true); }; + img.src = URL.createObjectURL(blob); + }; + + const metaHandler = e => { + const { blob, nodeId, displayNodeId } = e.detail; + if (String(displayNodeId || nodeId) === String(nodeRef.id)) show(blob); + }; + const plainHandler = e => { + if (api.serverSupportsFeature?.("supports_preview_metadata")) return; + if (String(execId) === String(nodeRef.id)) show(e.detail); + }; + + api.addEventListener("b_preview_with_metadata", metaHandler); + api.addEventListener("b_preview", plainHandler); + + chainCallback(nodeRef, "onRemoved", () => { + api.removeEventListener("b_preview_with_metadata", metaHandler); + api.removeEventListener("b_preview", plainHandler); + }); + }); + } +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/fast_preview_batch.js b/custom_nodes/ComfyUI-KJNodes/web/js/fast_preview_batch.js new file mode 100644 index 0000000000000000000000000000000000000000..f49e47cb3922b2ab04bd1488a846fea5a87cd229 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/fast_preview_batch.js @@ -0,0 +1,518 @@ +import { chainCallback, addMiddleClickPan, addWheelPassthrough } from './utility.js'; +const { app } = window.comfyAPI.app; +const { api } = window.comfyAPI.api; + +const STYLE_ID = 'kj-batch-preview-style'; +function ensureStyles() { + if (document.getElementById(STYLE_ID)) return; + const s = document.createElement('style'); + s.id = STYLE_ID; + s.textContent = ` +.kj-bp-root { width:100%; height:100%; position:relative; background:#1a1a1a; overflow:hidden; user-select:none; } +.kj-bp-grid { width:100%; height:100%; display:block; cursor:pointer; image-rendering:auto; } +.kj-bp-detail { position:absolute; inset:0; display:none; background:#0d0d0d; } +.kj-bp-detail.visible { display:flex; flex-direction:column; } +.kj-bp-detail-canvas { flex:1 1 auto; min-height:0; width:100%; display:block; } +.kj-bp-bar { display:flex; align-items:center; justify-content:space-between; padding:4px 8px; background:#222; color:#ddd; font:12px sans-serif; } +.kj-bp-btn { background:#333; color:#ddd; border:1px solid #444; padding:2px 10px; cursor:pointer; font:12px sans-serif; } +.kj-bp-btn:hover { background:#444; } +.kj-bp-btn:disabled { opacity:0.4; cursor:default; } +.kj-bp-scrubber { position:relative; flex:1 1 auto; height:22px; margin:0 8px; background:#2a2a2a; border:1px solid #444; cursor:ew-resize; user-select:none; touch-action:none; } +.kj-bp-scrubber-fill { position:absolute; top:0; bottom:0; left:0; width:0; background:#3a5a7a; pointer-events:none; } +.kj-bp-scrubber-text { position:absolute; inset:0; display:flex; align-items:center; justify-content:center; color:#ddd; font:11px sans-serif; pointer-events:none; text-shadow:0 1px 2px rgba(0,0,0,0.6); } +.kj-bp-status { position:absolute; inset:0; display:flex; align-items:center; justify-content:center; color:#888; font:12px sans-serif; pointer-events:none; } +`; + document.head.appendChild(s); +} + +function pickGridLayout(count, aspect, containerW, containerH) { + let best = { cols: 1, rows: count, cellW: 0, cellH: 0, area: 0 }; + for (let cols = 1; cols <= count; cols++) { + const rows = Math.ceil(count / cols); + const cellW = containerW / cols; + const cellH = containerH / rows; + const fitW = Math.min(cellW, cellH * aspect); + const fitH = fitW / aspect; + const area = fitW * fitH; + if (area > best.area) best = { cols, rows, cellW, cellH, area, fitW, fitH }; + } + return best; +} + +app.registerExtension({ + name: 'KJNodes.FastPreviewBatch', + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== 'FastPreviewBatch') return; + + chainCallback(nodeType.prototype, 'onNodeCreated', function () { + ensureStyles(); + this.setSize([520, 520]); + + const root = document.createElement('div'); + root.className = 'kj-bp-root'; + + const status = document.createElement('div'); + status.className = 'kj-bp-status'; + status.textContent = 'No preview yet'; + root.appendChild(status); + + const gridCanvas = document.createElement('canvas'); + gridCanvas.className = 'kj-bp-grid'; + gridCanvas.style.display = 'none'; + root.appendChild(gridCanvas); + + const detail = document.createElement('div'); + detail.className = 'kj-bp-detail'; + const detailCanvas = document.createElement('canvas'); + detailCanvas.className = 'kj-bp-detail-canvas'; + const bar = document.createElement('div'); + bar.className = 'kj-bp-bar'; + const prevBtn = document.createElement('button'); + prevBtn.className = 'kj-bp-btn'; + prevBtn.textContent = '❮'; + const nextBtn = document.createElement('button'); + nextBtn.className = 'kj-bp-btn'; + nextBtn.textContent = '❯'; + const scrubber = document.createElement('div'); + scrubber.className = 'kj-bp-scrubber'; + const scrubberFill = document.createElement('div'); + scrubberFill.className = 'kj-bp-scrubber-fill'; + const scrubberText = document.createElement('span'); + scrubberText.className = 'kj-bp-scrubber-text'; + scrubber.appendChild(scrubberFill); + scrubber.appendChild(scrubberText); + const playBtn = document.createElement('button'); + playBtn.className = 'kj-bp-btn'; + playBtn.textContent = '▶'; + const closeBtn = document.createElement('button'); + closeBtn.className = 'kj-bp-btn'; + closeBtn.textContent = '✕'; + bar.appendChild(playBtn); + bar.appendChild(prevBtn); + bar.appendChild(scrubber); + bar.appendChild(nextBtn); + bar.appendChild(closeBtn); + detail.appendChild(detailCanvas); + detail.appendChild(bar); + root.appendChild(detail); + + this.addDOMWidget('fast_preview_batch', 'div', root, { serialize: false }); + + const removeMiddleClickPan = addMiddleClickPan(root); + addWheelPassthrough(root); + + const state = { + video: null, + strip: null, + stripCols: 1, + stripCellW: 0, + stripCellH: 0, + frameCount: 0, + fps: 30, + thumbW: 0, + thumbH: 0, + layout: null, + gridReady: false, + currentIndex: 0, + hoverIndex: -1, + }; + + const setStatus = (msg) => { + if (msg) { + status.textContent = msg; + status.style.display = 'flex'; + } else { + status.style.display = 'none'; + } + }; + + const cleanupVideo = () => { + if (state.video) { + try { state.video.pause(); } catch (_) { } + state.video.removeAttribute('src'); + try { state.video.load(); } catch (_) { } + if (state.video.parentNode) state.video.parentNode.removeChild(state.video); + state.video = null; + } + }; + + const seekToFrame = (i) => new Promise((resolve) => { + if (!state.video) return resolve(); + const t = (i + 0.5) / state.fps; + const onSeeked = () => { + state.video.removeEventListener('seeked', onSeeked); + // give the decoder a tick to ensure frame is drawable + requestAnimationFrame(() => resolve()); + }; + state.video.addEventListener('seeked', onSeeked); + try { + state.video.currentTime = t; + } catch (_) { + state.video.removeEventListener('seeked', onSeeked); + resolve(); + } + }); + + let drawGeneration = 0; + const updateScrubberUI = () => { + const t = state.frameCount > 1 ? state.currentIndex / (state.frameCount - 1) : 0; + scrubberFill.style.width = (t * 100) + '%'; + scrubberText.textContent = `frame ${state.currentIndex + 1} / ${state.frameCount}`; + prevBtn.disabled = state.currentIndex <= 0; + nextBtn.disabled = state.currentIndex >= state.frameCount - 1; + }; + const paintFrameToDetail = () => { + const v = state.video; + if (!v) return; + const w = detailCanvas.clientWidth; + const h = detailCanvas.clientHeight; + if (w <= 0 || h <= 0) return; + const dpr = window.devicePixelRatio || 1; + const targetW = Math.max(1, Math.floor(w * dpr)); + const targetH = Math.max(1, Math.floor(h * dpr)); + if (detailCanvas.width !== targetW) detailCanvas.width = targetW; + if (detailCanvas.height !== targetH) detailCanvas.height = targetH; + const ctx = detailCanvas.getContext('2d'); + ctx.fillStyle = '#0d0d0d'; + ctx.fillRect(0, 0, detailCanvas.width, detailCanvas.height); + const aspect = state.thumbW / state.thumbH; + let dw = detailCanvas.width; + let dh = dw / aspect; + if (dh > detailCanvas.height) { + dh = detailCanvas.height; + dw = dh * aspect; + } + const dx = (detailCanvas.width - dw) / 2; + const dy = (detailCanvas.height - dh) / 2; + ctx.drawImage(v, dx, dy, dw, dh); + }; + const drawDetail = async () => { + if (!state.video || !state.frameCount) return; + updateScrubberUI(); + const gen = ++drawGeneration; + await seekToFrame(state.currentIndex); + if (gen !== drawGeneration) return; // superseded by a newer scrub + paintFrameToDetail(); + }; + + let playing = false; + let playHandle = 0; + const updatePlayBtn = () => { playBtn.textContent = playing ? '⏸' : '▶'; }; + const stopPlayback = () => { + if (!playing) return; + playing = false; + playHandle++; + if (state.video) { try { state.video.pause(); } catch (_) { } } + updatePlayBtn(); + }; + const startPlayback = async () => { + if (!state.video || !state.frameCount || playing) return; + if (typeof state.video.requestVideoFrameCallback !== 'function') return; + if (state.currentIndex >= state.frameCount - 1) state.currentIndex = 0; + drawGeneration++; // invalidate any pending drawDetail + await seekToFrame(state.currentIndex); + paintFrameToDetail(); + updateScrubberUI(); + playing = true; + updatePlayBtn(); + state.video.playbackRate = 1; + const myHandle = ++playHandle; + try { await state.video.play(); } catch (_) { stopPlayback(); return; } + if (myHandle !== playHandle) return; + const onFrame = (_now, metadata) => { + if (!playing || myHandle !== playHandle) return; + const idx = Math.min(state.frameCount - 1, + Math.floor(metadata.mediaTime * state.fps + 1e-6)); + if (idx !== state.currentIndex) { + state.currentIndex = idx; + updateScrubberUI(); + } + paintFrameToDetail(); + if (idx >= state.frameCount - 1) { stopPlayback(); return; } + state.video.requestVideoFrameCallback(onFrame); + }; + state.video.requestVideoFrameCallback(onFrame); + }; + + const renderGrid = () => { + if (!state.frameCount || !state.strip) return; + const containerW = root.clientWidth; + const containerH = root.clientHeight; + if (containerW <= 0 || containerH <= 0) return; + const dpr = window.devicePixelRatio || 1; + const aspect = state.thumbW / state.thumbH; + const layout = pickGridLayout(state.frameCount, aspect, containerW, containerH); + state.layout = layout; + gridCanvas.width = Math.floor(containerW * dpr); + gridCanvas.height = Math.floor(containerH * dpr); + gridCanvas.style.width = containerW + 'px'; + gridCanvas.style.height = containerH + 'px'; + const ctx = gridCanvas.getContext('2d'); + ctx.fillStyle = '#1a1a1a'; + ctx.fillRect(0, 0, gridCanvas.width, gridCanvas.height); + + const cellW = layout.cellW * dpr; + const cellH = layout.cellH * dpr; + const fitW = layout.fitW * dpr; + const fitH = layout.fitH * dpr; + const total = state.frameCount; + const sw = state.stripCellW; + const sh = state.stripCellH; + const sCols = state.stripCols; + for (let i = 0; i < total; i++) { + const sx = (i % sCols) * sw; + const sy = Math.floor(i / sCols) * sh; + const dCol = i % layout.cols; + const dRow = Math.floor(i / layout.cols); + const dx = dCol * cellW + (cellW - fitW) / 2; + const dy = dRow * cellH + (cellH - fitH) / 2; + ctx.drawImage(state.strip, sx, sy, sw, sh, dx, dy, fitW, fitH); + } + if (state.hoverIndex >= 0 && state.hoverIndex < total) { + const i = state.hoverIndex; + const sx = (i % sCols) * sw; + const sy = Math.floor(i / sCols) * sh; + const dCol = i % layout.cols; + const dRow = Math.floor(i / layout.cols); + const cx = dCol * cellW + cellW / 2; + const cy = dRow * cellH + cellH / 2; + const scale = 1.15; + const zw = fitW * scale; + const zh = fitH * scale; + let dx = cx - zw / 2; + let dy = cy - zh / 2; + dx = Math.max(0, Math.min(gridCanvas.width - zw, dx)); + dy = Math.max(0, Math.min(gridCanvas.height - zh, dy)); + ctx.save(); + ctx.shadowColor = 'rgba(0,0,0,0.65)'; + ctx.shadowBlur = 14 * dpr; + ctx.drawImage(state.strip, sx, sy, sw, sh, dx, dy, zw, zh); + ctx.restore(); + ctx.strokeStyle = '#5a8ec4'; + ctx.lineWidth = Math.max(2, 2 * dpr); + ctx.strokeRect(dx + ctx.lineWidth / 2, dy + ctx.lineWidth / 2, zw - ctx.lineWidth, zh - ctx.lineWidth); + } + state.gridReady = true; + setStatus(''); + gridCanvas.style.display = 'block'; + }; + + const enterDetail = (idx) => { + state.currentIndex = idx; + detail.classList.add('visible'); + requestAnimationFrame(drawDetail); + }; + const exitDetail = () => { + stopPlayback(); + detail.classList.remove('visible'); + }; + + playBtn.addEventListener('click', () => { + if (playing) stopPlayback(); + else startPlayback(); + }); + + const cellIndexFromEvent = (e) => { + if (!state.gridReady || !state.layout) return -1; + const rect = gridCanvas.getBoundingClientRect(); + if (rect.width <= 0 || rect.height <= 0) return -1; + const x = e.clientX - rect.left; + const y = e.clientY - rect.top; + const cellW = rect.width / state.layout.cols; + const cellH = rect.height / state.layout.rows; + const col = Math.floor(x / cellW); + const row = Math.floor(y / cellH); + if (col < 0 || col >= state.layout.cols || row < 0 || row >= state.layout.rows) return -1; + const idx = row * state.layout.cols + col; + return (idx >= 0 && idx < state.frameCount) ? idx : -1; + }; + + gridCanvas.addEventListener('click', (e) => { + const idx = cellIndexFromEvent(e); + if (idx >= 0) enterDetail(idx); + }); + + gridCanvas.addEventListener('pointermove', (e) => { + const idx = cellIndexFromEvent(e); + if (idx !== state.hoverIndex) { + state.hoverIndex = idx; + renderGrid(); + } + }); + gridCanvas.addEventListener('pointerleave', () => { + if (state.hoverIndex !== -1) { + state.hoverIndex = -1; + renderGrid(); + } + }); + + prevBtn.addEventListener('click', () => { + if (state.currentIndex > 0) { + stopPlayback(); + state.currentIndex--; + drawDetail(); + } + }); + nextBtn.addEventListener('click', () => { + if (state.currentIndex < state.frameCount - 1) { + stopPlayback(); + state.currentIndex++; + drawDetail(); + } + }); + closeBtn.addEventListener('click', exitDetail); + + let scrubbing = false; + const scrubFromEvent = (e) => { + if (!state.frameCount) return; + const rect = scrubber.getBoundingClientRect(); + const t = Math.max(0, Math.min(1, (e.clientX - rect.left) / rect.width)); + const idx = Math.round(t * (state.frameCount - 1)); + if (idx !== state.currentIndex) { + state.currentIndex = idx; + drawDetail(); + } + }; + scrubber.addEventListener('pointerdown', (e) => { + if (e.button !== 0) return; // let middle/right click bubble (canvas pan, etc.) + stopPlayback(); + scrubbing = true; + scrubber.setPointerCapture(e.pointerId); + scrubFromEvent(e); + e.preventDefault(); + }); + scrubber.addEventListener('pointermove', (e) => { + if (scrubbing) scrubFromEvent(e); + }); + const endScrub = (e) => { + if (!scrubbing) return; + scrubbing = false; + try { scrubber.releasePointerCapture(e.pointerId); } catch (_) { } + }; + scrubber.addEventListener('pointerup', endScrub); + scrubber.addEventListener('pointercancel', endScrub); + + const keyHandler = (e) => { + if (!detail.classList.contains('visible')) return; + if (e.key === 'Escape') { exitDetail(); e.stopPropagation(); } + else if (e.key === ' ') { + if (playing) stopPlayback(); else startPlayback(); + e.preventDefault(); e.stopPropagation(); + } + else if (e.key === 'ArrowLeft' && state.currentIndex > 0) { + stopPlayback(); state.currentIndex--; drawDetail(); e.stopPropagation(); + } else if (e.key === 'ArrowRight' && state.currentIndex < state.frameCount - 1) { + stopPlayback(); state.currentIndex++; drawDetail(); e.stopPropagation(); + } + }; + window.addEventListener('keydown', keyHandler); + + const ro = new ResizeObserver(() => { + if (detail.classList.contains('visible')) { + drawDetail(); + } else if (state.strip) { + renderGrid(); + } + }); + ro.observe(root); + + const buildViewUrl = (filename, subfolder, type) => { + const params = new URLSearchParams({ + filename: filename, + subfolder: subfolder || '', + type: type || 'temp', + t: String(Date.now()), + }); + return api.apiURL(`/view?${params.toString()}`); + }; + + const loadFromOutput = async (info) => { + stopPlayback(); + cleanupVideo(); + state.strip = null; + state.gridReady = false; + gridCanvas.style.display = 'none'; + setStatus('Loading preview…'); + + state.frameCount = info.frame_count; + state.fps = info.fps || 30; + state.thumbW = info.thumb_w; + state.thumbH = info.thumb_h; + state.stripCols = info.strip_cols || 1; + state.stripCellW = info.strip_cell_w || info.thumb_w; + state.stripCellH = info.strip_cell_h || info.thumb_h; + if (state.currentIndex >= state.frameCount || state.currentIndex < 0) { + state.currentIndex = 0; + } + + const v = document.createElement('video'); + v.muted = true; + v.playsInline = true; + v.preload = 'auto'; + v.style.position = 'absolute'; + v.style.left = '-99999px'; + v.style.width = '1px'; + v.style.height = '1px'; + document.body.appendChild(v); + state.video = v; + + // Video loads in parallel for the detail view; failure is non-fatal + // for the grid (strip handles that). Caught to avoid unhandled-rejection. + new Promise((resolve, reject) => { + v.addEventListener('loadeddata', resolve, { once: true }); + v.addEventListener('error', () => { + cleanupVideo(); + reject(new Error('video load failed')); + }, { once: true }); + v.src = buildViewUrl(info.filename, info.subfolder, info.type); + }).catch(() => { }); + + await new Promise((resolve, reject) => { + const img = new Image(); + img.onload = () => { state.strip = img; resolve(); }; + img.onerror = () => reject(new Error('strip load failed')); + img.src = buildViewUrl(info.strip_filename, info.subfolder, info.type); + }); + renderGrid(); + if (detail.classList.contains('visible')) { + // Refresh detail canvas with the new video; drawDetail awaits seek, + // which itself queues until the new video element finishes loading. + drawDetail(); + } + }; + + const tryLoad = (info, onMissing) => { + loadFromOutput(info).catch((err) => { + console.warn('FastPreviewBatch load failed', err); + if (onMissing) onMissing(); + else setStatus('Preview failed: ' + (err && err.message || err)); + }); + }; + + chainCallback(this, 'onExecuted', function (output) { + if (!output || !output.kj_batch_preview) return; + const info = output.kj_batch_preview[0]; + if (!info) return; + this.properties = this.properties || {}; + this.properties.kj_batch_preview = info; + tryLoad(info); + }); + + chainCallback(this, 'onConfigure', function () { + const info = this.properties && this.properties.kj_batch_preview; + if (!info) return; + tryLoad(info, () => { + setStatus('Preview unavailable (re-run to regenerate)'); + delete this.properties.kj_batch_preview; + }); + }); + + chainCallback(this, 'onRemoved', function () { + window.removeEventListener('keydown', keyHandler); + ro.disconnect(); + removeMiddleClickPan(); + cleanupVideo(); + }); + }); + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/fillconnect.js b/custom_nodes/ComfyUI-KJNodes/web/js/fillconnect.js new file mode 100644 index 0000000000000000000000000000000000000000..1903e415eb10237b0dbae50440c9787fbf8d3f5a --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/fillconnect.js @@ -0,0 +1,225 @@ +const { app } = window.comfyAPI.app; +import { getSlotPos } from "./utility.js"; + +/** + * Check type compatibility. Returns a match tier: + * 2 = exact single-type match, 1 = multi-type or wildcard match, -1 = incompatible + */ +function typeMatchTier(outType, inType) { + if (!LiteGraph.isValidConnection(outType, inType)) return -1; + if (typeof outType === "string" && typeof inType === "string" + && !outType.includes(",") && !inType.includes(",") + && outType.toUpperCase() === inType.toUpperCase()) return 2; + return 1; +} + + +/** + * Order selected nodes for connection. + * Respects existing links between them (topological sort), + * then inserts unconnected nodes by spatial position (left-to-right, top-to-bottom). + */ +function orderNodes(nodes, graph) { + const selectedIds = new Set(nodes.map(n => n.id)); + const nodeById = new Map(nodes.map(n => [n.id, n])); + + const outEdges = new Map(); + const inDegree = new Map(); + for (const n of nodes) { + outEdges.set(n.id, new Set()); + inDegree.set(n.id, 0); + } + + for (const node of nodes) { + for (const out of (node.outputs || [])) { + for (const linkId of (out.links || [])) { + const link = graph.getLink(linkId); + if (link && selectedIds.has(link.target_id) && link.target_id !== node.id) { + if (!outEdges.get(node.id).has(link.target_id)) { + outEdges.get(node.id).add(link.target_id); + inDegree.set(link.target_id, inDegree.get(link.target_id) + 1); + } + } + } + } + } + + const bySpatial = (a, b) => a.pos[0] - b.pos[0] || a.pos[1] - b.pos[1]; + const queue = nodes.filter(n => inDegree.get(n.id) === 0); + queue.sort(bySpatial); + + const ordered = []; + while (queue.length > 0) { + const node = queue.shift(); + ordered.push(node); + for (const targetId of outEdges.get(node.id)) { + const deg = inDegree.get(targetId) - 1; + inDegree.set(targetId, deg); + if (deg === 0 && nodeById.has(targetId)) { + queue.push(nodeById.get(targetId)); + queue.sort(bySpatial); + } + } + } + + // Cycle fallback: append remaining by position + if (ordered.length < nodes.length) { + const inOrdered = new Set(ordered.map(n => n.id)); + const remaining = nodes.filter(n => !inOrdered.has(n.id)); + remaining.sort(bySpatial); + ordered.push(...remaining); + } + + return ordered; +} + +/** + * Collect all candidate connections and assign globally by closest 2D distance. + * Each input and output is used at most once per invocation. + */ +function planConnections(ordered) { + // Collect all candidates globally across all target nodes + const candidates = []; + + for (let b = 1; b < ordered.length; b++) { + const nodeB = ordered[b]; + if (!nodeB.inputs) continue; + + for (let inIdx = 0; inIdx < nodeB.inputs.length; inIdx++) { + const inp = nodeB.inputs[inIdx]; + if (inp.link != null) continue; + + const inPos = getSlotPos(nodeB, true, inIdx); + const inputCandidates = []; + let hasExact = false; + + const inName = (inp.name || inp.label || "").toLowerCase(); + + for (let a = b - 1; a >= 0; a--) { + const nodeA = ordered[a]; + if (!nodeA.outputs) continue; + + for (let outIdx = 0; outIdx < nodeA.outputs.length; outIdx++) { + const out = nodeA.outputs[outIdx]; + const tier = typeMatchTier(out.type, inp.type); + if (tier < 0) continue; + if (tier === 2) hasExact = true; + + const outName = (out.name || out.label || "").toLowerCase(); + const nameMatch = inName !== "" && inName === outName ? 1 : 0; + + const outPos = getSlotPos(nodeA, false, outIdx); + const dx = outPos[0] - inPos[0]; + const dy = outPos[1] - inPos[1]; + inputCandidates.push({ + targetNode: nodeB, + inIdx, + sourceNode: nodeA, + outIdx, + tier, + nameMatch, + dist: dx * dx + dy * dy, + }); + } + } + + for (const c of inputCandidates) { + c.hasExact = hasExact; + } + candidates.push(...inputCandidates); + } + } + + // Filter out non-exact matches for inputs that have an exact candidate (keep name matches) + const filtered = candidates.filter(c => c.tier === 2 || c.nameMatch || !c.hasExact); + + // Sort globally: name match first, then higher type tier, then closest pair + filtered.sort((x, y) => (y.nameMatch - x.nameMatch) || (y.tier - x.tier) || (x.dist - y.dist)); + + // Greedy global assignment — closest pairs first + const planned = []; + const usedInputs = new Set(); + const usedOutputs = new Set(); + + for (const c of filtered) { + const inKey = `${c.targetNode.id}:${c.inIdx}`; + if (usedInputs.has(inKey)) continue; + const outKey = `${c.sourceNode.id}:${c.outIdx}`; + if (usedOutputs.has(outKey)) continue; + + planned.push({ + sourceNode: c.sourceNode, + outIdx: c.outIdx, + targetNode: c.targetNode, + inIdx: c.inIdx, + }); + usedInputs.add(inKey); + usedOutputs.add(outKey); + } + + return planned; +} + +function connectPlanned(planned, graph) { + for (const p of planned) { + p.sourceNode.connect(p.outIdx, p.targetNode, p.inIdx); + } + graph.change(); +} + +function fillConnectSelected() { + if (!app.ui.settings.getSettingValue("KJNodes.fillConnectEnabled")) return; + + const canvas = app.canvas; + const graph = canvas.graph; + const nodes = Object.values(canvas.selected_nodes || {}); + if (nodes.length < 2) return; + + const ordered = orderNodes(nodes, graph); + const planned = planConnections(ordered); + + if (planned.length > 0) { + connectPlanned(planned, graph); + } + + if (app.extensionManager?.toast) { + app.extensionManager.toast.add({ + severity: planned.length > 0 ? "info" : "warn", + summary: planned.length > 0 + ? `Connected ${planned.length} link${planned.length > 1 ? "s" : ""}` + : "No compatible connections found", + life: 2000, + }); + } +} + +app.registerExtension({ + name: "KJNodes.FillConnect", + + settings: [ + { + id: "KJNodes.fillConnectEnabled", + name: "Enable Fill-Connect hotkey", + category: ["KJNodes", "Fill-Connect"], + tooltip: "Fill-connect selected nodes by compatible types", + type: "boolean", + defaultValue: true, + }, + ], + + commands: [ + { + id: "KJNodes.fillConnectSelected", + label: "Fill-Connect Selected Nodes", + function: fillConnectSelected, + }, + ], + + keybindings: [ + { + commandId: "KJNodes.fillConnectSelected", + combo: { key: "f" }, + targetElementId: "graph-canvas", + }, + ], +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/fix_node.js b/custom_nodes/ComfyUI-KJNodes/web/js/fix_node.js new file mode 100644 index 0000000000000000000000000000000000000000..aec5cbd9b621a9a2030a1d040e826b1cd0f9a071 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/fix_node.js @@ -0,0 +1,93 @@ +const { app } = window.comfyAPI.app; + +function snapshotOld(oldNode) { + const data = oldNode.serialize(); + const widgetValuesByName = {}; + for (const w of oldNode.widgets || []) widgetValuesByName[w.name] = w.value; + + const inputSnapshot = (oldNode.inputs || []).map(i => ({ + name: i.name, + widgetName: i.widget?.name ?? null, + linkId: i.link ?? null, + })); + const outputLinksByName = {}; + for (const o of oldNode.outputs || []) { + if (o.links?.length) outputLinksByName[o.name] = [...o.links]; + } + return { data, widgetValuesByName, inputSnapshot, outputLinksByName }; +} + +function applyCosmetics(newNode, old) { + if (old.pos) newNode.pos = [...old.pos]; + if (old.size) { + const min = newNode.computeSize?.() || [0, 0]; + newNode.size = [Math.max(old.size[0], min[0]), Math.max(old.size[1], min[1])]; + } + if (old.title) newNode.title = old.title; + if (old.color) newNode.color = old.color; + if (old.bgcolor) newNode.bgcolor = old.bgcolor; + if (old.flags) newNode.flags = { ...(newNode.flags || {}), ...old.flags }; + if (typeof old.mode === "number") newNode.mode = old.mode; + if (old.properties) newNode.properties = { ...(newNode.properties || {}), ...old.properties }; +} + +function fixNode(oldNode, comfyClass, { resetValues = false } = {}) { + const graph = app.canvas.graph; + const snap = snapshotOld(oldNode); + + const newNode = LiteGraph.createNode(comfyClass); + if (!newNode) { + console.error(`[KJNodes.FixNode] Unknown node type: ${comfyClass}`); + return null; + } + + try { + graph.add(newNode, false); + applyCosmetics(newNode, snap.data); + + // Restore widget values by name (live widgets, not widgets_values array). + if (!resetValues) { + for (const w of newNode.widgets || []) { + if (Object.prototype.hasOwnProperty.call(snap.widgetValuesByName, w.name)) { + try { w.value = snap.widgetValuesByName[w.name]; } catch {} + } + } + } + + // Reconnect inputs by name. + for (const inp of snap.inputSnapshot) { + if (inp.linkId == null) continue; + const link = app.graph.links[inp.linkId]; + if (!link) continue; + const src = app.graph.getNodeById(link.origin_id); + const newSlot = newNode.findInputSlot?.(inp.name); + if (!src || newSlot == null || newSlot < 0) continue; + try { src.connect(link.origin_slot, newNode, newSlot); } catch {} + } + + // Reconnect outputs by name. + (newNode.outputs || []).forEach((out, i) => { + const linkIds = snap.outputLinksByName[out.name]; + if (!linkIds) return; + for (const id of linkIds) { + const link = app.graph.links[id]; + if (!link) continue; + const tgt = app.graph.getNodeById(link.target_id); + if (!tgt) continue; + try { newNode.connect(i, tgt, link.target_slot); } catch {} + } + }); + + graph.remove(oldNode); + app.graph.afterChange(); + requestAnimationFrame(() => app.canvas.setDirty(true, true)); + return newNode; + } catch (err) { + console.error("[KJNodes.FixNode] Aborting, rolling back:", err); + try { graph.remove(newNode); } catch {} + return null; + } +} + +window.kjNodes = window.kjNodes || {}; +window.kjNodes.recreateNode = fixNode; diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/hdr_preview.js b/custom_nodes/ComfyUI-KJNodes/web/js/hdr_preview.js new file mode 100644 index 0000000000000000000000000000000000000000..c7c454f58e4b2f2d82f6114a8cb7c04ff4b7a8b6 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/hdr_preview.js @@ -0,0 +1,574 @@ +import { chainCallback, addMiddleClickPan, addWheelPassthrough } from './utility.js'; +const { app } = window.comfyAPI.app; + +// Shared across all HDR Preview nodes so synced nodes can drive each other. +const hdrSyncGroup = new Set(); + +const VERTEX_SHADER = `#version 300 es +out vec2 v_texCoord; +void main() { + vec2 verts[3] = vec2[](vec2(-1.0, -1.0), vec2(3.0, -1.0), vec2(-1.0, 3.0)); + v_texCoord = verts[gl_VertexID] * 0.5 + 0.5; + v_texCoord.y = 1.0 - v_texCoord.y; + gl_Position = vec4(verts[gl_VertexID], 0.0, 1.0); +}`; + +// LogC3 constants ported from ComfyUI-LTXVideo/hdr.py +const FRAGMENT_SHADER = `#version 300 es +precision highp float; +in vec2 v_texCoord; +out vec4 fragColor; +uniform sampler2D u_image; +uniform float u_exposure; +uniform float u_saturation; +uniform float u_linearScale; +uniform int u_space; // 0 = logc3, 1 = linear + +const float LC_A = 5.555556; +const float LC_B = 0.052272; +const float LC_C = 0.247190; +const float LC_D = 0.385537; +const float LC_E = 5.367655; +const float LC_F = 0.092809; +const float LC_CUT = 0.010591; +const float LC_CUT_LOG = LC_E * LC_CUT + LC_F; // ~0.14966 + +vec3 logc3_decompress(vec3 logc) { + logc = clamp(logc, 0.0, 1.0); + vec3 lin_from_log = (pow(vec3(10.0), (logc - LC_D) / LC_C) - LC_B) / LC_A; + vec3 lin_from_lin = (logc - LC_F) / LC_E; + vec3 is_log = step(vec3(LC_CUT_LOG), logc); + return mix(lin_from_lin, lin_from_log, is_log); +} + +vec3 linear_to_srgb(vec3 x) { + vec3 cutoff = vec3(0.0031308); + vec3 low = 12.92 * x; + vec3 high = 1.055 * pow(max(x, cutoff), vec3(1.0 / 2.4)) - 0.055; + return clamp(mix(low, high, step(cutoff, x)), 0.0, 1.0); +} + +vec3 srgb_to_linear(vec3 x) { + vec3 cutoff = vec3(0.04045); + vec3 low = x / 12.92; + vec3 high = pow(max(x, vec3(0.0)) + 0.055, vec3(2.4)) / pow(vec3(1.055), vec3(2.4)); + return mix(low, high, step(cutoff, x)); +} + +void main() { + vec3 col = texture(u_image, v_texCoord).rgb; + vec3 hdr; + if (u_space == 0) { + hdr = logc3_decompress(col); + } else if (u_space == 1) { + hdr = col * u_linearScale; + } else { + hdr = srgb_to_linear(col); + } + hdr = max(hdr, vec3(0.0)); + vec3 exposed = hdr * exp2(u_exposure); + float luma = dot(exposed, vec3(0.2126, 0.7152, 0.0722)); + vec3 saturated = max(mix(vec3(luma), exposed, u_saturation), vec3(0.0)); + vec3 tm = (u_space == 2) ? clamp(saturated, 0.0, 1.0) : saturated / (1.0 + saturated); + fragColor = vec4(linear_to_srgb(tm), 1.0); +}`; + +function compileShader(gl, type, src) { + const sh = gl.createShader(type); + gl.shaderSource(sh, src); + gl.compileShader(sh); + if (!gl.getShaderParameter(sh, gl.COMPILE_STATUS)) { + const log = gl.getShaderInfoLog(sh) || "shader compile failed"; + gl.deleteShader(sh); + throw new Error(log); + } + return sh; +} + +function createProgram(gl, vsSrc, fsSrc) { + const vs = compileShader(gl, gl.VERTEX_SHADER, vsSrc); + const fs = compileShader(gl, gl.FRAGMENT_SHADER, fsSrc); + const prog = gl.createProgram(); + gl.attachShader(prog, vs); + gl.attachShader(prog, fs); + gl.linkProgram(prog); + if (!gl.getProgramParameter(prog, gl.LINK_STATUS)) { + const log = gl.getProgramInfoLog(prog) || "program link failed"; + gl.deleteProgram(prog); + throw new Error(log); + } + gl.deleteShader(vs); + gl.deleteShader(fs); + return prog; +} + +function loadImageBitmapFromView(filename, type) { + return fetch(`/view?filename=${encodeURIComponent(filename)}&type=${type}`) + .then(r => r.ok ? r.blob() : null) + .then(blob => blob ? createImageBitmap(blob, { colorSpaceConversion: "none", premultiplyAlpha: "none" }) : null) + .catch(() => null); +} + +app.registerExtension({ + name: "KJNodes.HDRPreview", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== "HDRPreviewKJ") return; + + chainCallback(nodeType.prototype, "onNodeCreated", function () { + const node = this; + + const state = { + frames: [], + frameCount: 0, + width: 512, height: 512, + fps: 24, + inputSpace: "logc3", + linearScale: 1.0, + exposure: 0.0, + saturation: 1.0, + currentFrame: 0, + playing: false, + playTimer: null, + gl: null, program: null, texture: null, vao: null, + uniforms: {}, + uploadedFrame: -1, + needsRender: false, + execGen: 0, + }; + node._hdrState = state; + + let controlsHeight = 28; // measured from DOM after mount; this is the fallback + node._widgetHeight = controlsHeight; + + const container = document.createElement("div"); + container.style.cssText = + "position:relative;width:100%;background:#111;display:flex;flex-direction:column;align-items:center;overflow:hidden;"; + + const viewport = document.createElement("div"); + // flex-shrink:0 — legacy container can be shorter than this height; without it the viewport collapses to 0. + viewport.style.cssText = + "position:relative;width:100%;height:0;background:#000;overflow:hidden;flex-shrink:0;"; + + const canvas = document.createElement("canvas"); + canvas.style.cssText = + "display:block;width:100%;height:100%;"; + viewport.appendChild(canvas); + + const frameRow = document.createElement("div"); + frameRow.style.cssText = + "display:flex;align-items:center;gap:6px;padding:3px 6px;color:#ccc;font-size:11px;background:#1a1a1a;width:100%;box-sizing:border-box;flex:0 0 auto;"; + + const playBtn = document.createElement("button"); + playBtn.type = "button"; + playBtn.textContent = "▶"; + playBtn.style.cssText = + "background:#333;color:#ccc;border:1px solid #555;cursor:pointer;padding:1px 8px;font-size:11px;min-width:24px;"; + + const syncBtn = document.createElement("button"); + syncBtn.type = "button"; + syncBtn.textContent = "⛓"; + syncBtn.title = "Sync playback with other HDR Preview nodes that have sync enabled."; + syncBtn.style.cssText = + "background:#333;color:#888;border:1px solid #555;cursor:pointer;padding:1px 6px;font-size:11px;min-width:22px;"; + + const frameSlider = document.createElement("input"); + frameSlider.type = "range"; + frameSlider.min = "0"; + frameSlider.max = "0"; + frameSlider.value = "0"; + frameSlider.step = "1"; + frameSlider.style.cssText = "flex:1;min-width:40px;accent-color:#5af;"; + + const frameLabel = document.createElement("span"); + frameLabel.textContent = "0/0"; + frameLabel.style.cssText = "min-width:50px;text-align:right;font-variant-numeric:tabular-nums;"; + + frameRow.appendChild(playBtn); + frameRow.appendChild(syncBtn); + frameRow.appendChild(frameSlider); + frameRow.appendChild(frameLabel); + + container.appendChild(viewport); + container.appendChild(frameRow); + + const stopProp = (e) => e.stopPropagation(); + for (const el of [frameSlider, playBtn, syncBtn]) { + el.addEventListener("pointerdown", stopProp); + el.addEventListener("mousedown", stopProp); + } + + addMiddleClickPan(container); + addWheelPassthrough(canvas); + + function initGL() { + try { + const gl = canvas.getContext("webgl2", { antialias: false, premultipliedAlpha: false, alpha: false }); + if (!gl) { + console.error("[HDRPreviewKJ] WebGL2 not available"); + return false; + } + const program = createProgram(gl, VERTEX_SHADER, FRAGMENT_SHADER); + gl.useProgram(program); + + const uniforms = { + u_image: gl.getUniformLocation(program, "u_image"), + u_exposure: gl.getUniformLocation(program, "u_exposure"), + u_saturation: gl.getUniformLocation(program, "u_saturation"), + u_space: gl.getUniformLocation(program, "u_space"), + u_linearScale: gl.getUniformLocation(program, "u_linearScale"), + }; + gl.uniform1i(uniforms.u_image, 0); + + const texture = gl.createTexture(); + gl.activeTexture(gl.TEXTURE0); + gl.bindTexture(gl.TEXTURE_2D, texture); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE); + + const vao = gl.createVertexArray(); + gl.bindVertexArray(vao); + + state.gl = gl; + state.program = program; + state.texture = texture; + state.vao = vao; + state.uniforms = uniforms; + return true; + } catch (err) { + console.error("[HDRPreviewKJ] WebGL init failed:", err); + return false; + } + } + + function requestRender() { + if (state.needsRender) return; + state.needsRender = true; + requestAnimationFrame(() => { + state.needsRender = false; + render(); + }); + } + + function render() { + if (!state.gl || !state.frames.length) return; + const frame = state.frames[state.currentFrame]; + if (!frame) return; + + const gl = state.gl; + // Backing buffer sized to source resolution for quality; CSS handles display scaling + if (canvas.width !== state.width || canvas.height !== state.height) { + canvas.width = state.width; + canvas.height = state.height; + } + gl.viewport(0, 0, canvas.width, canvas.height); + gl.useProgram(state.program); + gl.bindVertexArray(state.vao); + gl.activeTexture(gl.TEXTURE0); + gl.bindTexture(gl.TEXTURE_2D, state.texture); + + if (state.uploadedFrame !== state.currentFrame) { + gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, false); + gl.pixelStorei(gl.UNPACK_ALIGNMENT, 1); + gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGB8, gl.RGB, gl.UNSIGNED_BYTE, frame); + state.uploadedFrame = state.currentFrame; + } + + gl.uniform1f(state.uniforms.u_exposure, state.exposure); + gl.uniform1f(state.uniforms.u_saturation, state.saturation); + const spaceIdx = state.inputSpace === "logc3" ? 0 : state.inputSpace === "srgb" ? 2 : 1; + gl.uniform1i(state.uniforms.u_space, spaceIdx); + gl.uniform1f(state.uniforms.u_linearScale, state.linearScale); + + gl.clearColor(0, 0, 0, 1); + gl.clear(gl.COLOR_BUFFER_BIT); + gl.drawArrays(gl.TRIANGLES, 0, 3); + } + + function stopPlayback() { + if (state.playTimer !== null) { + clearInterval(state.playTimer); + state.playTimer = null; + } + state.playing = false; + playBtn.textContent = "▶"; + } + + function updateFrameLabel() { + const total = Math.max(state.frameCount, 1); + frameLabel.textContent = `${state.currentFrame + 1}/${total}`; + } + + const exposureWidget = node.widgets?.find(w => w.name === "exposure"); + const saturationWidget = node.widgets?.find(w => w.name === "saturation"); + + function hookLiveWidget(widget, stateKey) { + if (!widget) return; + state[stateKey] = isFinite(widget.value) ? widget.value : state[stateKey]; + const original = widget.callback; + widget.callback = function (value) { + const num = isFinite(value) ? value : state[stateKey]; + state[stateKey] = num; + requestRender(); + return original?.apply(this, arguments); + }; + } + hookLiveWidget(exposureWidget, "exposure"); + hookLiveWidget(saturationWidget, "saturation"); + + // Sync handle: other nodes in the group call these to follow along. + let syncEnabled = false; + const syncHandle = { + setFrameFraction(fraction) { + if (state.frameCount <= 0) return; + const maxIdx = Math.max(state.frameCount - 1, 0); + const frame = Math.min(maxIdx, Math.max(0, Math.round(fraction * maxIdx))); + if (frame === state.currentFrame) return; + state.currentFrame = frame; + frameSlider.value = String(frame); + updateFrameLabel(); + requestRender(); + }, + showPlaying(play) { + state.playing = play; + playBtn.textContent = play ? "■" : "▶"; + if (!play && state.playTimer !== null) { + clearInterval(state.playTimer); + state.playTimer = null; + } + }, + }; + + function broadcastFraction() { + if (!syncEnabled || state.frameCount <= 1) return; + const fraction = state.currentFrame / (state.frameCount - 1); + for (const other of hdrSyncGroup) { + if (other !== syncHandle) other.setFrameFraction(fraction); + } + } + + function broadcastPlaying(play) { + if (!syncEnabled) return; + for (const other of hdrSyncGroup) { + if (other !== syncHandle) other.showPlaying(play); + } + } + + syncBtn.addEventListener("click", () => { + syncEnabled = !syncEnabled; + syncBtn.style.color = syncEnabled ? "#5fa" : "#888"; + syncBtn.style.borderColor = syncEnabled ? "#5fa" : "#555"; + if (syncEnabled) { + hdrSyncGroup.add(syncHandle); + } else { + hdrSyncGroup.delete(syncHandle); + } + }); + + frameSlider.addEventListener("input", () => { + state.currentFrame = parseInt(frameSlider.value, 10) || 0; + updateFrameLabel(); + requestRender(); + broadcastFraction(); + }); + frameSlider.addEventListener("pointerdown", () => { + stopPlayback(); + broadcastPlaying(false); + }); + function persistCurrentFrame() { + if (!node.properties?.hdrLastPreview) return; + node.properties.hdrLastPreview.current_frame = state.currentFrame; + node.graph?.change?.(); + try { app.extensionManager?.workflow?.activeWorkflow?.changeTracker?.checkState?.(); } catch {} + } + frameSlider.addEventListener("change", persistCurrentFrame); + + playBtn.addEventListener("click", () => { + if (state.playing) { + stopPlayback(); + broadcastPlaying(false); + persistCurrentFrame(); + return; + } + if (state.frameCount <= 1) return; + // Ensure no other synced node has an active timer — we become the sole driver. + if (syncEnabled) { + for (const other of hdrSyncGroup) { + if (other !== syncHandle) other.showPlaying(false); + } + } + const intervalMs = Math.max(16, 1000 / Math.max(state.fps, 1)); + state.playing = true; + playBtn.textContent = "■"; + state.playTimer = setInterval(() => { + state.currentFrame = (state.currentFrame + 1) % state.frameCount; + frameSlider.value = String(state.currentFrame); + updateFrameLabel(); + requestRender(); + broadcastFraction(); + }, intervalMs); + broadcastPlaying(true); + }); + + const domWidget = node.addDOMWidget("hdr_preview", "hdr_preview", container, { + serialize: false, + hideOnZoom: false, + margin: 0, + getMinHeight: () => node._widgetHeight, + getMaxHeight: () => node._widgetHeight, + getHeight: () => node._widgetHeight, + }); + node.resizable = true; + + // LiteGraph's computeSize uses a fixed per-widget height that ignores DOM widgets — sum manually. + const NATIVE_ROW_H = (LiteGraph?.NODE_WIDGET_HEIGHT ?? 20) + 4; + const TITLE_H = LiteGraph?.NODE_TITLE_HEIGHT ?? 30; + function computeNodeHeight() { + let nativeCount = 0; + for (const w of node.widgets || []) { + if (w === domWidget) continue; + if (w.hidden || w.type === "converted-widget") continue; + nativeCount++; + } + return TITLE_H + nativeCount * NATIVE_ROW_H + node._widgetHeight + 8; + } + + // Don't gate on state.frames — a fire during async bitmap load would wipe _widgetHeight and collapse the node. + const controlsObserver = new ResizeObserver(() => { + const measured = frameRow.offsetHeight; + if (measured <= 0 || measured === controlsHeight) return; + controlsHeight = measured; + const viewportPx = parseInt(viewport.style.height, 10) || 0; + node._widgetHeight = viewportPx + controlsHeight; + node.setSize([node.size[0], computeNodeHeight()]); + node.graph?.setDirtyCanvas(true, true); + }); + controlsObserver.observe(frameRow); + + let resizing = false; + function resizeToFit() { + if (resizing) return; + resizing = true; + try { + const srcW = state.width || 512, srcH = state.height || 512; + const availW = Math.max(100, node.size[0] - 30); + const ratio = srcH / srcW; + const displayH = Math.round(availW * ratio); + const totalH = displayH + controlsHeight; + viewport.style.width = availW + "px"; + viewport.style.height = displayH + "px"; + // min-height keeps children visible when ComfyUI sets the container height from a stale computeSize. + container.style.minHeight = totalH + "px"; + node._widgetHeight = totalH; + node.setSize([node.size[0], computeNodeHeight()]); + node.graph?.setDirtyCanvas(true, true); + } finally { + resizing = false; + } + } + + chainCallback(node, "onResize", function () { + if (state.frames.length) resizeToFit(); + }); + + async function applyPreviewData(data) { + const gen = ++state.execGen; + + stopPlayback(); + + for (const f of state.frames) { + try { f.close?.(); } catch {} + } + state.frames = []; + state.uploadedFrame = -1; + + state.frameCount = data.frame_count || 0; + state.width = data.width || 512; + state.height = data.height || 512; + state.fps = data.fps || 24; + state.inputSpace = data.input_space || "logc3"; + state.linearScale = data.linear_scale || 1.0; + const restoredFrame = Number.isInteger(data.current_frame) ? data.current_frame : 0; + state.currentFrame = Math.min(Math.max(0, restoredFrame), Math.max(0, state.frameCount - 1)); + + frameSlider.max = String(Math.max(0, state.frameCount - 1)); + frameSlider.value = String(state.currentFrame); + updateFrameLabel(); + + if (!state.gl) initGL(); + + resizeToFit(); + + const bitmaps = await Promise.all( + (data.frames || []).map(f => loadImageBitmapFromView(f.filename, f.type)) + ); + + if (gen !== state.execGen) { + for (const b of bitmaps) { + try { b?.close?.(); } catch {} + } + return; + } + + state.frames = bitmaps.filter(Boolean); + // Re-assert size in case the ResizeObserver fired mid-await with empty frames. + resizeToFit(); + requestRender(); + } + + chainCallback(node, "onExecuted", async function (message) { + const data = message?.hdr_preview_data?.[0]; + if (!data) return; + + node.properties = node.properties || {}; + node.properties.hdrLastPreview = { + frames: data.frames, + width: data.width, + height: data.height, + fps: data.fps, + input_space: data.input_space, + linear_scale: data.linear_scale, + frame_count: data.frame_count, + }; + // Autosave snapshots before execution, so post-execute property updates need an explicit dirty nudge. + node.graph?.change?.(); + try { + const ct = app.extensionManager?.workflow?.activeWorkflow?.changeTracker; + ct?.checkState?.(); + } catch {} + try { app.workflowManager?.activeWorkflow?.changeTracker?.checkState?.(); } catch {} + + await applyPreviewData(data); + }); + + chainCallback(node, "onConfigure", function () { + if (exposureWidget && isFinite(exposureWidget.value)) state.exposure = exposureWidget.value; + if (saturationWidget && isFinite(saturationWidget.value)) state.saturation = saturationWidget.value; + const saved = node.properties?.hdrLastPreview; + if (saved?.frames?.length) applyPreviewData(saved); + }); + + chainCallback(node, "onRemoved", function () { + stopPlayback(); + controlsObserver.disconnect(); + hdrSyncGroup.delete(syncHandle); + for (const f of state.frames) { + try { f.close?.(); } catch {} + } + state.frames = []; + if (state.gl) { + try { + state.gl.deleteProgram(state.program); + state.gl.deleteTexture(state.texture); + state.gl.deleteVertexArray(state.vao); + const lose = state.gl.getExtension("WEBGL_lose_context"); + if (lose) lose.loseContext(); + } catch {} + state.gl = null; + } + }); + }); + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/help_popup.js b/custom_nodes/ComfyUI-KJNodes/web/js/help_popup.js new file mode 100644 index 0000000000000000000000000000000000000000..d089fbb11c76ceb7df4066638717b3df546ca3d3 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/help_popup.js @@ -0,0 +1,468 @@ +const { app } = window.comfyAPI.app; + +// code based on mtb nodes by Mel Massadian https://github.com/melMass/comfy_mtb/ +export const loadScript = ( + FILE_URL, + async = true, + type = 'text/javascript', +) => { + return new Promise((resolve, reject) => { + try { + // Check if the script already exists + const existingScript = document.querySelector(`script[src="${FILE_URL}"]`) + if (existingScript) { + resolve({ status: true, message: 'Script already loaded' }) + return + } + + const scriptEle = document.createElement('script') + scriptEle.type = type + scriptEle.async = async + scriptEle.src = FILE_URL + + scriptEle.addEventListener('load', (ev) => { + resolve({ status: true }) + }) + + scriptEle.addEventListener('error', (ev) => { + reject({ + status: false, + message: `Failed to load the script ${FILE_URL}`, + }) + }) + + document.body.appendChild(scriptEle) + } catch (error) { + reject(error) + } + }) +} + +loadScript('kjweb_async/marked.min.js').catch((e) => { + console.error(e) +}) +loadScript('kjweb_async/purify.min.js').catch((e) => { + console.error(e) +}) + +const categories = ["KJNodes", "SUPIR", "VoiceCraft", "Marigold", "IC-Light", "WanVideoWrapper"]; +const nodeDescriptions = new Map(); + +function isHelpPopupEnabled() { + return app.ui.settings.getSettingValue("KJNodes.helpPopup") !== false; +} + +app.registerExtension({ + name: "KJNodes.HelpPopup", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (!isHelpPopupEnabled()) return; + try { + categories.forEach(category => { + if (nodeData?.category?.startsWith(category)) { + if (nodeData.description) { + nodeDescriptions.set(nodeData.name, nodeData.description); + } + addDocumentation(nodeData, nodeType); + } + else return + }); + } catch (error) { + console.error("Error in registering KJNodes.HelpPopup", error); + } + }, + nodeCreated(node) { + if (!isHelpPopupEnabled()) return; + const description = nodeDescriptions.get(node.type) || nodeDescriptions.get(node.comfyClass); + if (!description) return; + node._kjHelpDescription = description; + }, + setup() { + if (!isHelpPopupEnabled()) return; + setupHelpObserver(); + }, +}); + +const create_documentation_stylesheet = () => { + const tag = 'kj-documentation-stylesheet' + + let styleTag = document.getElementById(tag) + + if (!styleTag) { + styleTag = document.createElement('style') + styleTag.type = 'text/css' + styleTag.id = tag + styleTag.innerHTML = ` + .kj-documentation-popup { + background: var(--comfy-menu-bg); + position: absolute; + color: var(--fg-color); + font: 12px monospace; + line-height: 1.5em; + padding: 10px; + border-radius: 10px; + border-style: solid; + border-width: medium; + border-color: var(--border-color); + z-index: 5; + overflow: hidden; + } + .content-wrapper { + overflow: auto; + max-height: 100%; + /* Scrollbar styling for Chrome */ + &::-webkit-scrollbar { + width: 6px; + } + &::-webkit-scrollbar-track { + background: var(--bg-color); + } + &::-webkit-scrollbar-thumb { + background-color: var(--fg-color); + border-radius: 6px; + border: 3px solid var(--bg-color); + } + + /* Scrollbar styling for Firefox */ + scrollbar-width: thin; + scrollbar-color: var(--fg-color) var(--bg-color); + a { + color: yellow; + } + a:visited { + color: orange; + } + a:hover { + color: red; + } + } + ` + document.head.appendChild(styleTag) + } + } + +/** + * Creates the documentation popup DOM and wires up resize/close interactions. + * Returns { docElement, contentWrapper }. + * @param {string} description - Markdown description text + * @param {AbortSignal} signal - Signal to clean up event listeners + * @param {() => void} onClose - Called when the close button is clicked + * @param {{ scaleResize: boolean }} opts - If scaleResize is true, resize deltas are divided by canvas scale + */ +function createDocPopup(description, signal, onClose, opts = {}) { + create_documentation_stylesheet() + + const docElement = document.createElement('div') + const contentWrapper = document.createElement('div') + docElement.appendChild(contentWrapper) + + contentWrapper.classList.add('content-wrapper') + docElement.classList.add('kj-documentation-popup') + // Try ComfyUI's built-in markdown renderer first (available after frontend PR #10700) + if (app.extensionManager?.renderMarkdownToHtml) { + contentWrapper.innerHTML = app.extensionManager.renderMarkdownToHtml(description) + } else if (typeof marked !== 'undefined' && typeof DOMPurify !== 'undefined') { + contentWrapper.innerHTML = DOMPurify.sanitize(marked.parse(description)) + } else { + // Fallback: convert markdown links to
    tags, auto-link bare URLs, preserve line breaks + const escaped = description + .replace(/&/g, '&').replace(//g, '>') + .replace(/\[([^\]]+)\]\((https?:\/\/[^)]+)\)/g, '$1') + .replace(/(^|[^"'])(https?:\/\/[^\s<]+)/g, '$1$2') + .replace(/\n/g, '
    ') + contentWrapper.innerHTML = escaped + } + + // resize handle + const resizeHandle = document.createElement('div') + resizeHandle.style.width = '0' + resizeHandle.style.height = '0' + resizeHandle.style.position = 'absolute' + resizeHandle.style.bottom = '0' + resizeHandle.style.right = '0' + resizeHandle.style.cursor = 'se-resize' + const borderColor = getComputedStyle(document.documentElement).getPropertyValue('--border-color').trim() + resizeHandle.style.borderTop = '10px solid transparent' + resizeHandle.style.borderLeft = '10px solid transparent' + resizeHandle.style.borderBottom = `10px solid ${borderColor}` + resizeHandle.style.borderRight = `10px solid ${borderColor}` + docElement.appendChild(resizeHandle) + + let isResizing = false + let startX, startY, startWidth, startHeight + resizeHandle.addEventListener('mousedown', (e) => { + e.preventDefault() + e.stopPropagation() + isResizing = true + startX = e.clientX + startY = e.clientY + startWidth = parseInt(document.defaultView.getComputedStyle(docElement).width, 10) + startHeight = parseInt(document.defaultView.getComputedStyle(docElement).height, 10) + }, { signal }) + + document.addEventListener('mousemove', (e) => { + if (!isResizing) return + const scaleFactor = opts.scaleResize ? app.canvas.ds.scale : 1 + const newWidth = startWidth + (e.clientX - startX) / scaleFactor + const newHeight = startHeight + (e.clientY - startY) / scaleFactor + docElement.style.width = `${newWidth}px` + docElement.style.height = `${newHeight}px` + }, { signal }) + + document.addEventListener('mouseup', () => { + isResizing = false + }, { signal }) + + // close button + const closeButton = document.createElement('div') + closeButton.textContent = '❌' + closeButton.style.position = 'absolute' + closeButton.style.top = '0' + closeButton.style.right = '0' + closeButton.style.cursor = 'pointer' + closeButton.style.padding = '5px' + closeButton.style.color = 'red' + closeButton.style.fontSize = '12px' + docElement.appendChild(closeButton) + + closeButton.addEventListener('mousedown', (e) => { + e.stopPropagation() + onClose() + }, { signal }) + + document.body.appendChild(docElement) + return { docElement, contentWrapper } +} + +// ─── Legacy canvas mode (onDrawForeground) ─── + +/** Add documentation widget to the selected node (legacy canvas rendering) */ +export const addDocumentation = ( + nodeData, + nodeType, + opts = { icon_size: 14, icon_margin: 4 },) => { + + opts = opts || {} + const iconSize = opts.icon_size ? opts.icon_size : 14 + const iconMargin = opts.icon_margin ? opts.icon_margin : 4 + let docElement = null + let contentWrapper = null + + if (!nodeData.description) { + return + } + + const drawFg = nodeType.prototype.onDrawForeground + nodeType.prototype.onDrawForeground = function (ctx) { + const r = drawFg ? drawFg.apply(this, arguments) : undefined + if (this.flags.collapsed) return r + + const x = this.size[0] - iconSize - iconMargin + + // create the popup + if (this.show_doc && docElement === null) { + const popup = createDocPopup( + nodeData.description, + this.docCtrl.signal, + () => { + this.show_doc = !this.show_doc + docElement.parentNode.removeChild(docElement) + docElement = null + contentWrapper = null + }, + { scaleResize: true } + ) + docElement = popup.docElement + contentWrapper = popup.contentWrapper + } + // close the popup + else if (!this.show_doc && docElement !== null) { + docElement.parentNode.removeChild(docElement) + docElement = null + } + // update position of the popup + if (this.show_doc && docElement !== null) { + const rect = ctx.canvas.getBoundingClientRect() + const scaleX = rect.width / ctx.canvas.width + const scaleY = rect.height / ctx.canvas.height + + const transform = new DOMMatrix() + .scaleSelf(scaleX, scaleY) + .multiplySelf(ctx.getTransform()) + .translateSelf(this.size[0] * scaleX * Math.max(1.0,window.devicePixelRatio) , 0) + .translateSelf(10, -32) + + const scale = new DOMMatrix() + .scaleSelf(transform.a, transform.d); + const bcr = app.canvas.canvas.getBoundingClientRect() + + const styleObject = { + transformOrigin: '0 0', + transform: scale, + left: `${transform.a + bcr.x + transform.e}px`, + top: `${transform.d + bcr.y + transform.f}px`, + }; + Object.assign(docElement.style, styleObject); + } + + ctx.save() + ctx.translate(x - 2, iconSize - 34) + ctx.scale(iconSize / 32, iconSize / 32) + ctx.strokeStyle = 'rgba(255,255,255,0.3)' + ctx.lineCap = 'round' + ctx.lineJoin = 'round' + ctx.lineWidth = 2.4 + ctx.font = 'bold 36px monospace' + ctx.fillStyle = 'orange'; + ctx.fillText('?', 0, 24) + ctx.restore() + return r + } + + // handle clicking of the icon + const mouseDown = nodeType.prototype.onMouseDown + nodeType.prototype.onMouseDown = function (e, localPos, canvas) { + const r = mouseDown ? mouseDown.apply(this, arguments) : undefined + const iconX = this.size[0] - iconSize - iconMargin + const iconY = iconSize - 34 + if ( + localPos[0] > iconX && + localPos[0] < iconX + iconSize && + localPos[1] > iconY && + localPos[1] < iconY + iconSize + ) { + if (this.show_doc === undefined) { + this.show_doc = true + } else { + this.show_doc = !this.show_doc + } + if (this.show_doc) { + this.docCtrl = new AbortController() + } else { + this.docCtrl.abort() + } + return true; + } + return r; + } + + const onRem = nodeType.prototype.onRemoved + nodeType.prototype.onRemoved = function () { + const r = onRem ? onRem.apply(this, []) : undefined + if (docElement) { + docElement.remove() + docElement = null + contentWrapper = null + } + return r + } +} + +// ─── Vue nodes mode (DOM injection via MutationObserver) ─── + +/** Per-node popup state, keyed by node ID */ +const popupState = new Map() + +function getNodeById(nodeId) { + return app.graph?.getNodeById(nodeId) +} + +function closeNodePopup(nodeId) { + const state = popupState.get(nodeId) + if (!state) return + if (state.docElement) { + state.docElement.remove() + } + if (state.abortCtrl) { + state.abortCtrl.abort() + } + if (state.animFrame) { + cancelAnimationFrame(state.animFrame) + } + popupState.delete(nodeId) +} + +function openNodePopup(nodeId, description) { + closeNodePopup(nodeId) + const state = popupState.get(nodeId) || {} + popupState.set(nodeId, state) + + state.abortCtrl = new AbortController() + const popup = createDocPopup( + description, + state.abortCtrl.signal, + () => closeNodePopup(nodeId), + { scaleResize: false } + ) + state.docElement = popup.docElement + + function updatePosition() { + if (!state.docElement || !state.docElement.parentNode) return + const nodeEl = document.querySelector(`[data-node-id="${nodeId}"]`) + if (nodeEl) { + const rect = nodeEl.getBoundingClientRect() + state.docElement.style.left = `${rect.right + 10}px` + state.docElement.style.top = `${rect.top}px` + } + state.animFrame = requestAnimationFrame(updatePosition) + } + state.animFrame = requestAnimationFrame(updatePosition) +} + +/** Try to inject a "?" button into a Vue node header */ +function tryInjectHelpButton(header) { + if (header.querySelector('.kj-help-btn')) return + + const nodeEl = header.closest('[data-node-id]') + if (!nodeEl) return + const nodeId = nodeEl.dataset.nodeId + const node = getNodeById(nodeId) + if (!node) return + + const description = node._kjHelpDescription + if (!description) return + + const flexContainer = header.querySelector(':scope > div') + if (!flexContainer) return + + const helpBtn = document.createElement('span') + helpBtn.className = 'kj-help-btn' + helpBtn.textContent = '?' + helpBtn.style.cssText = ` + color: orange; + font-weight: bold; + font-size: 14px; + cursor: pointer; + flex-shrink: 0; + padding: 0 4px; + line-height: 1; + user-select: none; + ` + helpBtn.title = 'Show help' + helpBtn.addEventListener('click', (e) => { + e.stopPropagation() + const state = popupState.get(nodeId) + if (state?.docElement) { + closeNodePopup(nodeId) + } else { + openNodePopup(nodeId, description) + } + }) + flexContainer.appendChild(helpBtn) +} + +/** Observe the DOM for Vue node headers appearing and inject help buttons */ +function setupHelpObserver() { + // Inject into any headers already in the DOM + document.querySelectorAll('.lg-node-header').forEach(tryInjectHelpButton) + + let pending = false + const observer = new MutationObserver(() => { + if (pending) return + pending = true + requestAnimationFrame(() => { + pending = false + document.querySelectorAll('.lg-node-header:not(:has(.kj-help-btn))').forEach(tryInjectHelpButton) + }) + }) + observer.observe(document.body, { childList: true, subtree: true }) +} diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/ideogram4_prompt_builder.js b/custom_nodes/ComfyUI-KJNodes/web/js/ideogram4_prompt_builder.js new file mode 100644 index 0000000000000000000000000000000000000000..626cffed68cbeaf99f738fb4ddf8b605efa356d7 --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/ideogram4_prompt_builder.js @@ -0,0 +1,2544 @@ +import { chainCallback, addMiddleClickPan, addWheelPassthrough, cursorForBboxMode, watchImageInputs, captureVideoFrame } from './utility.js'; +const { app } = window.comfyAPI.app; + +const HANDLE = 8; // hit radius (canvas px) for corners/edges +const MAX_ELEM_COLORS = 5; // Ideogram 4 per-element palette cap +const MAX_STYLE_COLORS = 16; // Ideogram 4 style palette cap +let copiedBoxes = null; // internal clipboard for copy/paste of regions (array; shared across nodes) + +// Track the most recent generated image so it can be grabbed as a background. +let lastResultImage = null; +try { + app.api?.addEventListener?.("executed", (e) => { + const imgs = e?.detail?.output?.images; + if (Array.isArray(imgs) && imgs.length) { + lastResultImage = imgs[imgs.length - 1]; + // Live nodes swap their preview for the full-res final result automatically. + for (const n of livePreviewNodes) n._ideoGrabFinal?.(); + } + }); +} catch (e) {} +function resultViewUrl(img) { + const p = new URLSearchParams({ filename: img.filename || "", subfolder: img.subfolder || "", type: img.type || "output" }); + return "/view?" + p.toString(); +} + +// Nodes opted into "live background": feed them the sampling preview frames as they arrive. +const livePreviewNodes = new Set(); +let _lastPreviewBlob = null; +function feedPreviewBlob(blob) { + // Newer ComfyUI dispatches "b_preview_with_metadata"; older ones only "b_preview" — and both can + // fire for the same frame (same Blob), so dedupe by identity to avoid decoding twice. + if (!blob || blob === _lastPreviewBlob || !livePreviewNodes.size) return; + _lastPreviewBlob = blob; + createImageBitmap(blob).then((bmp) => { + let used = false; + for (const n of livePreviewNodes) { n._ideoSetLiveBg?.(bmp); used = true; } + if (!used && bmp.close) bmp.close(); + }).catch(() => {}); +} +try { + app.api?.addEventListener?.("b_preview", (e) => feedPreviewBlob(e?.detail)); + app.api?.addEventListener?.("b_preview_with_metadata", (e) => feedPreviewBlob(e?.detail?.blob)); +} catch (e) {} + +// Named caption-JSON templates, each stored as its own file server-side via ComfyUI's userdata +// API (under the ComfyUI user dir) so they persist across browsers/machines and survive clears. +const TPL_DIR = "kjnodes/ideogram4/templates"; +const tplSafe = (s) => (s || "").replace(/[\/\\:*?"<>|]+/g, "_").trim(); // filesystem-safe name +const tplFile = (name) => `${TPL_DIR}/${name}.json`; +async function listTemplateNames() { + try { + const items = await app.api.listUserDataFullInfo(TPL_DIR); + return items.map((it) => it.path.split(/[\\/]/).pop() || "") + .filter((f) => /\.json$/i.test(f)) + .map((f) => f.replace(/\.json$/i, "")) + .filter(Boolean).sort((a, b) => a.localeCompare(b)); + } catch (e) { return []; } +} +async function loadTemplate(name) { + try { const r = await app.api.getUserData(tplFile(name)); if (r.status === 200) return await r.text(); } catch (e) {} + return null; +} +async function saveTemplate(name, caption) { + try { + await app.api.storeUserData(tplFile(name), caption, { overwrite: true, stringify: false, throwOnError: true }); + return true; + } catch (e) { window.alert("Couldn't save the template to the server."); return false; } +} +async function deleteTemplate(name) { + try { await app.api.deleteUserData(tplFile(name)); } catch (e) {} +} + +// Parse a #rrggbb hex into {r,g,b}, or null if malformed. +function hexRgb(hex) { + const h = (hex || "").replace("#", ""); + if (h.length < 6) return null; + return { r: parseInt(h.slice(0, 2), 16), g: parseInt(h.slice(2, 4), 16), b: parseInt(h.slice(4, 6), 16) }; +} +// Perceived luminance (0-255) of an {r,g,b}. +function luminance({ r, g, b }) { return 0.299 * r + 0.587 * g + 0.114 * b; } + +// The hex color lightened toward white if too dark, so it stays readable on the dark canvas. +function readableText(hex) { + const c = hexRgb(hex); + if (!c) return "#d4d4d4"; + let { r, g, b } = c; + const lum = luminance(c), MIN = 130; + if (lum < MIN) { + const t = (MIN - lum) / (255 - lum); + r = Math.round(r + (255 - r) * t); g = Math.round(g + (255 - g) * t); b = Math.round(b + (255 - b) * t); + } + return `rgb(${r},${g},${b})`; +} + +// Black or white, whichever contrasts better with the given hex background. +function textOn(hex) { + const c = hexRgb(hex); + if (!c) return "#000"; + return luminance(c) > 140 ? "#000" : "#fff"; +} + +// Parse a clipboard string into a normalized #rrggbb, or null if it isn't a color. +function parseColorString(s) { + if (!s) return null; + s = s.trim(); + let m = s.match(/^#?([0-9a-fA-F]{6})$/); + if (m) return "#" + m[1].toLowerCase(); + m = s.match(/^#?([0-9a-fA-F]{3})$/); + if (m) { const h = m[1]; return ("#" + h[0] + h[0] + h[1] + h[1] + h[2] + h[2]).toLowerCase(); } + m = s.match(/^rgba?\(\s*(\d+)\s*[, ]\s*(\d+)\s*[, ]\s*(\d+)/i); + if (m) { + const h2 = (n) => Math.max(0, Math.min(255, parseInt(n, 10))).toString(16).padStart(2, "0"); + return "#" + h2(m[1]) + h2(m[2]) + h2(m[3]); + } + return null; +} +// The palette swatch under the cursor + a setter, so Ctrl+V pastes a color onto it +// and Ctrl+C copies its hex. +let hoveredSwatch = null; +document.addEventListener("keydown", (e) => { + if (!hoveredSwatch) return; + const ctrl = e.ctrlKey || e.metaKey; + if (!ctrl) return; + const key = e.key.toLowerCase(); + if (key !== "v" && key !== "c") return; + const ae = document.activeElement; // don't steal copy/paste from a focused text field + if (ae && (ae.tagName === "TEXTAREA" || ae.tagName === "INPUT" || ae.isContentEditable)) return; + const target = hoveredSwatch; + if (!target.sw || !target.sw.isConnected) return; + e.preventDefault(); e.stopPropagation(); + if (key === "c") { + navigator.clipboard?.writeText?.(target.sw.dataset.hex || "").catch(() => {}); + } else { + navigator.clipboard?.readText?.().then((txt) => { + const c = parseColorString(txt); + if (c && target.sw && target.sw.isConnected) target.setColor(c); + }).catch(() => {}); + } +}, true); + +// The node whose canvas is under the cursor — lets H toggle box visibility on hover (no click needed). +let hoveredCanvasNode = null; +document.addEventListener("keydown", (e) => { + if (!hoveredCanvasNode) return; + if (e.ctrlKey || e.metaKey || e.altKey || (e.key !== "h" && e.key !== "H")) return; + const ae = document.activeElement; + if (ae && (ae.tagName === "TEXTAREA" || ae.tagName === "INPUT" || ae.isContentEditable)) return; + e.preventDefault(); e.stopPropagation(); + hoveredCanvasNode._toggleHideBoxes?.(); +}, true); +// Alt held over a canvas: preview the alt-click cycle target even without moving the mouse. +document.addEventListener("keydown", (e) => { if (hoveredCanvasNode && e.key === "Alt") hoveredCanvasNode._altRefresh?.(true); }, true); +document.addEventListener("keyup", (e) => { if (hoveredCanvasNode && e.key === "Alt") hoveredCanvasNode._altRefresh?.(false); }, true); + +// Pinned dock follows the node: Nodes 2.0 hosts it inside the node element (inherits its transform); legacy is body-fixed. +const pinnedDocks = new Set(); +const liveDocks = new Set(); // every node with a dock — swept so an orphaned dock can't linger +let _dockRectNonce = 0; +window.addEventListener("resize", () => { _dockRectNonce++; wakeDocks(); }); // canvas may shift without pan/zoom +// Safety net: a dock whose node has left the graph (deleted, workflow cleared) is torn down even if onRemoved didn't fire. +function sweepOrphanDocks() { + for (const n of liveDocks) { + if (n.graph?.getNodeById?.(n.id) === n) continue; // still registered in its graph + try { n._dockRO?.disconnect(); } catch (e) {} + try { n._visObserver?.disconnect(); } catch (e) {} + n._dockEl?.remove(); + pinnedDocks.delete(n); liveDocks.delete(n); + } +} +// gr.x/gr.y = offset (graph units) from the node's content top-left (node.pos). +function applyDockTransform(n) { + const c = app.canvas, fl = n._dockEl, gr = n.properties.dockGraph; + if (!c || !fl || !gr) return; + let nodeEl = null; + if (window.LiteGraph?.vueNodesMode && n.id != null) { // cache; re-query only if it's gone (virtualization) + nodeEl = n._dockNodeEl; + if (!nodeEl || !nodeEl.isConnected) nodeEl = n._dockNodeEl = document.querySelector(`[data-node-id="${n.id}"]`); + } + if (nodeEl) { // host inside the node element → inherits its transform + const title = window.LiteGraph?.NODE_TITLE_HEIGHT ?? 30; // element top is the title bar; node.pos sits below it + const sig = "v|" + gr.x + "|" + gr.y; + if (fl.parentElement !== nodeEl) { + nodeEl.appendChild(fl); + fl.style.position = "absolute"; fl.style.transform = ""; fl.style.transformOrigin = ""; fl.style.zIndex = ""; + n._dockSig = ""; + } + if (n._dockSig !== sig) { fl.style.left = gr.x + "px"; fl.style.top = (title + gr.y) + "px"; n._dockSig = sig; } + return; + } + // Legacy / fallback (no Vue element, e.g. just after a mode switch): body-fixed from node.pos + transform. + if (fl.parentElement !== document.body) { document.body.appendChild(fl); fl.style.left = ""; fl.style.top = ""; n._dockSig = ""; } // re-attach if orphaned + if (!n.pos) return; + const ds = c.ds, scale = ds.scale, rect = c.canvas.getBoundingClientRect(); + const baseLeft = rect.left + (n.pos[0] + gr.x + ds.offset[0]) * scale; + const baseTop = rect.top + (n.pos[1] + gr.y + ds.offset[1]) * scale; + const tf = `translate(${baseLeft}px,${baseTop}px) scale(${scale})`; + if (n._dockSig !== tf) { fl.style.position = "fixed"; fl.style.transformOrigin = "top left"; fl.style.transform = tf; n._dockSig = tf; } + const order = n.graph?.nodes?.indexOf(n); // sit at the node DOM-widget layer (not above everything) + if (order != null && order >= 0) fl.style.zIndex = String(order); +} +// Dock chrome follows the node's color theme (title → header, body → background); falls back to dark when uncolored. +function applyDockTheme(n) { + const fl = n._dockEl; if (!fl) return; + const sig = (n.color || "") + "|" + (n.bgcolor || ""); + if (n._dockTheme === sig) return; + n._dockTheme = sig; + const set = (k, v) => (v ? fl.style.setProperty(k, v) : fl.style.removeProperty(k)); + set("--kj-dock-bg", n.bgcolor); + set("--kj-dock-head", n.color); + set("--kj-dock-border", n.color || n.bgcolor); +} +// Event-driven loop: woken by canvas redraws (dirty-gated) + resize, self-stops when settled. +const DOCK_IDLE_STOP = 6; // stop after this many frames with no further wake +let _dockRAF = 0, _dockLoopSig = "", _dockLastMode = null, _dockIdle = 0, _dockWakesInstalled = false; +function tickDocks() { + const c = app.canvas; + if (c && pinnedDocks.size) { + for (const n of pinnedDocks) applyDockTheme(n); // cheap: cached, only writes on color change + const vue = !!window.LiteGraph?.vueNodesMode; + if (vue !== _dockLastMode) { // legacy↔2.0 flip rebuilds the node DOM — force re-parent + re-place + _dockLastMode = vue; _dockLoopSig = ""; + for (const n of pinnedDocks) { n._dockSig = ""; n._dockNodeEl = null; } + } + if (vue) { + for (const n of pinnedDocks) { + applyDockTransform(n); + const w = n._dockWrap; // Vue re-mounts the widget into the node body — pull it back + if (w && n._dockBody && !n._fullscreen && w.parentElement !== n._dockBody) { n._dockBody.appendChild(w); n._dockSig = ""; } + } + } else { // legacy: skip rect reads when nothing moved + const ds = c.ds; + let sig = _dockRectNonce + "|" + ds.scale + "|" + ds.offset[0] + "|" + ds.offset[1]; + for (const n of pinnedDocks) sig += n.pos ? "|" + n.pos[0] + "," + n.pos[1] : "|"; + if (sig !== _dockLoopSig) { _dockLoopSig = sig; for (const n of pinnedDocks) applyDockTransform(n); } + } + } + if (pinnedDocks.size && ++_dockIdle < DOCK_IDLE_STOP) _dockRAF = requestAnimationFrame(tickDocks); + else { _dockRAF = 0; _dockLoopSig = ""; } // settled — stop until the next wake +} +function wakeDocks() { // a canvas change happened — run/continue the loop until it settles + _dockIdle = 0; + if (!_dockRAF && pinnedDocks.size) _dockRAF = requestAnimationFrame(tickDocks); +} +function installDockWakes() { // onDrawForeground only fires when dirty_canvas is set + if (_dockWakesInstalled) return; + const c = app.canvas; if (!c) return; + _dockWakesInstalled = true; + chainCallback(c, "onDrawForeground", () => { sweepOrphanDocks(); wakeDocks(); }); +} +function startDockLoop() { installDockWakes(); wakeDocks(); } + +// Pointer-capture drag: filters moves by pointerId, auto-removes its listeners on release, calls onEnd. +function dragPointer(e, target, onMove, onEnd) { + try { target.setPointerCapture(e.pointerId); } catch (er) {} + const move = (me) => { if (me.pointerId === e.pointerId) onMove(me); }; + const end = (ue) => { + if (ue.pointerId !== e.pointerId) return; + target.removeEventListener("pointermove", move); + target.removeEventListener("pointerup", end); + target.removeEventListener("pointercancel", end); + if (onEnd) onEnd(ue); + }; + target.addEventListener("pointermove", move); + target.addEventListener("pointerup", end); + target.addEventListener("pointercancel", end); +} + +// Outside-click dismissal for a popup: arm() after showing, disarm() after hiding. +function outsideDismiss(menu, onDismiss, anchor) { + let handler = null; + const disarm = () => { + if (!handler) return; + document.removeEventListener("pointerdown", handler, true); + document.removeEventListener("mousedown", handler, true); + handler = null; + }; + const arm = () => { + disarm(); + handler = (e) => { if (!menu.contains(e.target) && e.target !== anchor) onDismiss(); }; + setTimeout(() => { // defer so the opening click itself doesn't dismiss + if (!handler) return; + document.addEventListener("pointerdown", handler, true); + document.addEventListener("mousedown", handler, true); + }, 0); + }; + return { arm, disarm }; +} + +function injectStyle() { + if (document.getElementById("kjideo-style")) return; + const s = document.createElement("style"); + s.id = "kjideo-style"; + s.textContent = ` + .kjideo-wrap { display:flex; flex-direction:column; overflow:hidden; position:relative; pointer-events:auto; gap:4px; } + .kjideo-cv { flex:1 1 auto; min-height:60px; display:flex; align-items:center; justify-content:center; overflow:hidden; } + .kjideo-canvas { cursor:crosshair; display:block; flex:0 0 auto; background:#1a1a1a; border-radius:4px; outline:none; touch-action:none; } + .kjideo-bar { display:flex; align-items:center; gap:6px; font:11px sans-serif; color:#aaa; user-select:none; padding:0 2px; flex:0 0 auto; } + .kjideo-panel { display:flex; flex-direction:column; gap:5px; padding:6px; background:#262626; border-radius:4px; font:11px sans-serif; color:#bbb; flex:0 0 auto; overflow-y:auto; min-height:0; } + .kjideo-split { flex:0 0 auto; height:8px; cursor:ns-resize; position:relative; } + .kjideo-split::before { content:""; position:absolute; left:50%; top:50%; transform:translate(-50%,-50%); width:34px; height:3px; background:#555; border-radius:2px; } + .kjideo-split:hover::before { background:#46b4e6; } + .kjideo-row { display:flex; align-items:center; gap:6px; flex-wrap:wrap; } + .kjideo-btn { background:#333; border:1px solid #555; border-radius:4px; color:#bbb; font:11px sans-serif; cursor:pointer; padding:2px 8px; line-height:16px; white-space:nowrap; flex-shrink:0; } + .kjideo-btn:hover { border-color:#46b4e6; color:#fff; } + .kjideo-btn.active { border-color:#46b4e6; color:#46b4e6; background:#2a3a42; } + .kjideo-area { width:100%; box-sizing:border-box; background:#1d1d1d; border:1px solid #444; border-radius:4px; color:#ddd; font:13px monospace; padding:4px 6px; resize:none; min-height:36px; flex:1 1 auto; } + .kjideo-sw { width:20px; height:20px; border:1px solid #666; border-radius:3px; cursor:pointer; flex-shrink:0; position:relative; touch-action:none; transition:transform .18s ease, box-shadow .12s ease, opacity .12s ease; } + .kjideo-sw:hover { transform:scale(1.2); box-shadow:0 0 0 2px #46b4e6; z-index:3; } + .kjideo-sw.dragging { opacity:.4; box-shadow:0 0 0 2px #46b4e6; } + body.kjideo-dragging, body.kjideo-dragging * { cursor:move !important; } + .kjideo-sw input { position:absolute; opacity:0; width:0; height:0; pointer-events:none; } + .kjideo-inline { position:absolute; box-sizing:border-box; background:rgba(18,18,18,0.92); border:2px solid #46b4e6; border-radius:3px; color:#fff; font:13px monospace; padding:3px 4px; resize:none; outline:none; z-index:10; } + .kjideo-bbox { width:128px; box-sizing:border-box; background:#1d1d1d; border:1px solid #444; border-radius:4px; color:#bbb; font:11px monospace; padding:2px 5px; } + .kjideo-bbox:focus { border-color:#46b4e6; outline:none; color:#fff; } + .kjideo-menu { position:fixed; z-index:10000; background:#262626; border:1px solid #555; border-radius:6px; padding:4px; box-shadow:0 6px 20px rgba(0,0,0,0.55); font:12px sans-serif; color:#ddd; max-height:60vh; overflow-y:auto; min-width:210px; max-width:340px; } + .kjideo-mhdr { font:11px sans-serif; color:#888; padding:2px 6px 4px; user-select:none; } + .kjideo-lrow { display:flex; align-items:center; gap:6px; padding:3px 5px; border-radius:4px; cursor:move; user-select:none; touch-action:none; transition:transform .18s ease, box-shadow .12s ease, opacity .12s ease, background .12s; } + .kjideo-lrow:hover { background:#333; } + .kjideo-lrow.active { background:#2a3a42; box-shadow:inset 0 0 0 1px #46b4e6; } + .kjideo-lrow.dragging { opacity:.4; box-shadow:0 0 0 2px #46b4e6; background:#333; } + .kjideo-lsw { width:16px; height:16px; border-radius:3px; border:1px solid #666; flex:0 0 auto; } + .kjideo-lnum { font:bold 11px monospace; color:#888; flex:0 0 auto; width:18px; } + .kjideo-ltext { flex:1 1 auto; min-width:0; white-space:nowrap; overflow:hidden; text-overflow:ellipsis; } + .kjideo-ltext.empty { color:#777; font-style:italic; } + .kjideo-lbtn { background:none; border:none; color:#999; cursor:pointer; font:13px sans-serif; line-height:1; padding:2px 5px; border-radius:3px; flex:0 0 auto; } + .kjideo-lbtn:hover { color:#fff; background:#444; } + .kjideo-lbtn.del:hover { color:#fff; background:#a33; } + .kjideo-lbtn.on { background:#3a3320; } + .kjideo-lock { filter:grayscale(1); opacity:0.4; } /* unlocked: faded grey */ + .kjideo-lock.on, .kjideo-lock:hover { filter:none; opacity:1; } /* locked / hover: full colour */ + .kjideo-lbtn:disabled { opacity:0.25; cursor:default; background:none; } + .kjideo-fs { position:fixed; inset:0; z-index:9000; background:rgba(0,0,0,0.72); display:flex; align-items:center; justify-content:center; } + .kjideo-fs-inner { position:relative; width:88vw; height:90vh; background:#1a1a1a; border:1px solid #444; border-radius:8px; box-shadow:0 12px 48px rgba(0,0,0,0.6); padding:12px; box-sizing:border-box; } + .kjideo-fs-inner .kjideo-wrap { height:100%; } + .kjideo-fs-close { position:absolute; top:14px; right:18px; z-index:5; padding:4px 12px; font-size:14px; } + .kjideo-dock { position:fixed; z-index:8500; pointer-events:auto; display:flex; flex-direction:column; background:var(--kj-dock-bg,#1a1a1a); border:1px solid var(--kj-dock-border,#555); border-radius:8px; box-shadow:0 8px 30px rgba(0,0,0,0.55); min-width:300px; min-height:240px; overflow:hidden; } + .kjideo-rsz { position:absolute; z-index:20; touch-action:none; } + .kjideo-rsz.n { top:0; left:11px; right:11px; height:6px; cursor:ns-resize; } + .kjideo-rsz.s { bottom:0; left:11px; right:11px; height:6px; cursor:ns-resize; } + .kjideo-rsz.e { right:0; top:11px; bottom:11px; width:6px; cursor:ew-resize; } + .kjideo-rsz.w { left:0; top:11px; bottom:11px; width:6px; cursor:ew-resize; } + .kjideo-rsz.ne { top:0; right:0; width:12px; height:12px; cursor:nesw-resize; } + .kjideo-rsz.nw { top:0; left:0; width:12px; height:12px; cursor:nwse-resize; } + .kjideo-rsz.se { bottom:0; right:0; width:12px; height:12px; cursor:nwse-resize; } + .kjideo-rsz.sw { bottom:0; left:0; width:12px; height:12px; cursor:nesw-resize; } + .kjideo-dock.minimized { min-height:0 !important; height:auto !important; } + .kjideo-dock.minimized .kjideo-dock-body { display:none; } + .kjideo-dock.minimized .kjideo-rsz { display:none; } + .kjideo-dock-head { display:flex; align-items:center; gap:6px; padding:4px 8px; background:var(--kj-dock-head,#262626); cursor:move; font:12px sans-serif; color:#ccc; user-select:none; border-bottom:1px solid rgba(0,0,0,0.25); flex:0 0 auto; } + .kjideo-dock-head .kjideo-btn { padding:1px 7px; } + .kjideo-dock-body { flex:1 1 auto; min-height:0; padding:8px; box-sizing:border-box; overflow:hidden; } + .kjideo-dock-body .kjideo-wrap { height:100%; } + .kjideo-bgmenu { padding:7px; display:flex; flex-direction:column; gap:7px; min-width:170px; } + .kjideo-bgrow { display:flex; align-items:center; gap:8px; } + .kjideo-bglbl { color:#888; font:11px sans-serif; flex:0 0 auto; min-width:62px; } + .kjideo-trow { padding:2px 4px; border-radius:4px; } + .kjideo-trow:hover { background:#333; } + `; + document.head.appendChild(s); +} + +app.registerExtension({ + name: "KJNodes.Ideogram4PromptBuilder", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== "Ideogram4PromptBuilderKJ") return; + injectStyle(); + + chainCallback(nodeType.prototype, "onNodeCreated", function () { + const node = this; + const findW = (n) => node.widgets?.find((w) => w.name === n); + const elementsWidget = findW("elements_data"); + const stylePaletteWidget = findW("style_palette_data"); + const bgBrightnessWidget = findW("bg_brightness"); + const formatWidget = findW("output_format"); // "pretty" | "compact" (set via the Text menu) + if (bgBrightnessWidget && typeof bgBrightnessWidget.value !== "number") bgBrightnessWidget.value = 25; + const wWidget = findW("width"), hWidget = findW("height"); + // Hide the data widgets while keeping them serializable. + function hideDataWidgets() { + for (const w of [elementsWidget, stylePaletteWidget, bgBrightnessWidget, formatWidget]) { + if (!w) continue; + w.hidden = true; + w.computeSize = () => [0, -4]; + } + for (const name of ["elements_data", "style_palette_data", "bg_brightness"]) { + const i = node.inputs?.findIndex((inp) => inp.name === name); + if (i != null && i !== -1) node.removeInput(i); + } + } + hideDataWidgets(); + + node._boxes = []; // {x,y,w,h normalized 0-1, type, text, desc, palette[]} + node._stylePalette = []; // global style color palette (hex[]) + node._activeIdx = -1; + node._drawing = false; + node._placing = false; // duplicate-placement mode: active box follows the cursor until clicked + node._dragMode = null; + node._dragStartN = null; // mouse-down point, normalized + node._boxAtStart = null; // active box snapshot at drag start + node._selection = new Set(); // selected box indices (multi-select); always contains _activeIdx + node._groupStart = null; // {idx: {x,y,w,h}} snapshot of all selected boxes at drag start + node._pendingCollapse = -1; // box to collapse selection to if a click (not drag) on a multi-selection + node._marquee = null; // {x0,y0,x,y} rubber-band selection rect (shift-drag), normalized + node._hoverTitle = null; // index of the title chip under the cursor + node._hoverBox = null; // index of the box under the cursor + node._focused = false; // editor (DOM) focused — gates the active-box highlight + node._selected = false; // node selected in the graph + node._bgImg = null; // optional reference image shown as the canvas background + node._bgManual = false; // bg set via "use last result" (not the image input) + node._lastImported = ""; // last import_json applied to the editor (avoid re-apply) + node._areaObservers = []; // (reserved) live ResizeObservers to disconnect on rebuild + + // ── DOM ── + const wrap = document.createElement("div"); + wrap.className = "kjideo-wrap"; + const bar = document.createElement("div"); + bar.className = "kjideo-bar"; + const hint = document.createElement("span"); + hint.style.flex = "1"; + const copyBtn = document.createElement("button"); + copyBtn.className = "kjideo-btn"; + copyBtn.textContent = "Copy"; + copyBtn.title = "Copy the current caption JSON to the clipboard"; + const importBtn = document.createElement("button"); + importBtn.className = "kjideo-btn"; + importBtn.textContent = "Paste"; + importBtn.title = "Parse a caption JSON (clipboard, else paste prompt) and populate the node"; + const clearBtn = document.createElement("button"); + clearBtn.className = "kjideo-btn"; + clearBtn.textContent = "Clear all"; + const tokenSpan = document.createElement("span"); + tokenSpan.style.cssText = "color:#888; white-space:nowrap;"; + tokenSpan.title = "Rough token estimate (~chars/4). Grey <256, green healthy, orange nearing, red ≥2048 (model cap — will error)"; + // Compact-output toggle next to the token count (compact is the format Ideogram 4 expects — default on). + const compactLbl = document.createElement("label"); + compactLbl.style.cssText = "display:flex;align-items:center;gap:3px;cursor:pointer;flex:0 0 auto;color:#aaa;white-space:nowrap;"; + compactLbl.title = "Compact JSON output (the format Ideogram 4 was trained on). Uncheck for pretty/indented. Copy reflects it."; + const compactCb = document.createElement("input"); compactCb.type = "checkbox"; + compactCb.checked = (formatWidget?.value) !== "pretty"; // default: compact + stopProp(compactCb); + compactCb.addEventListener("change", () => { if (formatWidget) formatWidget.value = compactCb.checked ? "compact" : "pretty"; updateTokens(); }); + compactLbl.appendChild(compactCb); compactLbl.appendChild(document.createTextNode("compact")); + compactLbl._cb = compactCb; + const grabBtn = document.createElement("button"); + grabBtn.className = "kjideo-btn"; + grabBtn.addEventListener("mousedown", (e) => e.stopPropagation()); + grabBtn.addEventListener("click", () => { (node._bgManual && node._bgImg) ? node._clearBg() : node._grabResultBg(); }); + function updateGrabBtn() { + const clear = node._bgManual && node._bgImg; + grabBtn.textContent = clear ? "Clear BG" : "Grab BG"; + grabBtn.title = clear ? "Remove the grabbed background" + : "Use the last generated image as the background"; + } + const liveLabel = document.createElement("label"); + liveLabel.style.cssText = "display:flex;align-items:center;gap:3px;flex:0 0 auto;cursor:pointer;"; + liveLabel.title = "Use the live sampling preview as the background while generating"; + const liveChk = document.createElement("input"); + liveChk.type = "checkbox"; + liveChk.checked = !!node.properties.liveBg; + liveChk.addEventListener("mousedown", (e) => e.stopPropagation()); + liveChk.addEventListener("change", () => { + node.properties.liveBg = liveChk.checked; + if (liveChk.checked) livePreviewNodes.add(node); else livePreviewNodes.delete(node); + }); + liveLabel.appendChild(liveChk); liveLabel.appendChild(document.createTextNode("Live")); + if (liveChk.checked) livePreviewNodes.add(node); + const bgSlider = document.createElement("input"); + bgSlider.type = "range"; bgSlider.min = "0"; bgSlider.max = "100"; bgSlider.step = "1"; + bgSlider.value = bgBrightnessWidget ? bgBrightnessWidget.value : 25; + bgSlider.title = "Background brightness (image or blank canvas)"; + bgSlider.style.cssText = "width:64px;flex:0 0 auto;"; + stopProp(bgSlider); + bgSlider.addEventListener("input", () => { if (bgBrightnessWidget) bgBrightnessWidget.value = parseInt(bgSlider.value, 10); drawCanvas(); }); + const guideSel = document.createElement("select"); + guideSel.className = "kjideo-btn"; guideSel.style.cssText = "flex:0 0 auto;"; + guideSel.title = "Composition guide overlay (editor view only)"; + for (const [val, label] of [["none", "no guide"], ["thirds", "thirds"], ["grid", "grid"], ["golden", "golden ratio"], ["spiral", "golden spiral"]]) { + const o = document.createElement("option"); o.value = val; o.textContent = label; guideSel.appendChild(o); + } + guideSel.value = node.properties.guide || "none"; + stopProp(guideSel); + guideSel.addEventListener("change", () => { node.properties.guide = guideSel.value; drawCanvas(); }); + // Group the background/guide controls into one popup to keep the toolbar tidy. + const bgBtn = document.createElement("button"); + bgBtn.className = "kjideo-btn"; bgBtn.textContent = "Background ▾"; + bgBtn.title = "Background & guides: live preview, grab/clear, brightness, composition guide"; + stopProp(bgBtn); + const GRID_INV = 130; // slider shows cell SIZE: position = GRID_INV - divisions (right = larger cells) + const gridSlider = document.createElement("input"); + gridSlider.type = "range"; gridSlider.min = "2"; gridSlider.max = "128"; gridSlider.step = "1"; + gridSlider.value = GRID_INV - (node.properties.gridSize || 10); + gridSlider.style.cssText = "width:90px;flex:0 0 auto;"; + gridSlider.title = "Grid cell size (drag right for larger cells); also the snap step"; + stopProp(gridSlider); + gridSlider.addEventListener("input", () => { node.properties.gridSize = GRID_INV - parseInt(gridSlider.value, 10); drawCanvas(); }); + const snapLabel = document.createElement("label"); + snapLabel.style.cssText = "display:flex;align-items:center;gap:4px;cursor:pointer;"; + const snapChk = document.createElement("input"); + snapChk.type = "checkbox"; snapChk.checked = !!node.properties.snap; + snapChk.addEventListener("change", () => { node.properties.snap = snapChk.checked; }); + snapLabel.appendChild(snapChk); snapLabel.appendChild(document.createTextNode("snap to grid")); + const guideColor = document.createElement("input"); + guideColor.type = "color"; guideColor.value = node.properties.guideColor || "#ffffff"; + guideColor.style.cssText = "width:32px;height:20px;flex:0 0 auto;padding:0;border:1px solid #555;background:none;"; + stopProp(guideColor); + guideColor.addEventListener("input", () => { node.properties.guideColor = guideColor.value; drawCanvas(); }); + const opacitySlider = document.createElement("input"); + opacitySlider.type = "range"; opacitySlider.min = "0"; opacitySlider.max = "100"; opacitySlider.step = "1"; + opacitySlider.value = node.properties.guideOpacity == null ? 100 : node.properties.guideOpacity; + opacitySlider.style.cssText = "width:90px;flex:0 0 auto;"; + opacitySlider.title = "Guide/grid line opacity"; + stopProp(opacitySlider); + opacitySlider.addEventListener("input", () => { node.properties.guideOpacity = parseInt(opacitySlider.value, 10); drawCanvas(); }); + const bgMenu = document.createElement("div"); + bgMenu.className = "kjideo-menu kjideo-bgmenu"; + bgMenu.style.display = "none"; + const bgRow = (labelText, el) => { + const r = document.createElement("div"); r.className = "kjideo-bgrow"; + if (labelText) { const l = document.createElement("span"); l.className = "kjideo-bglbl"; l.textContent = labelText; r.appendChild(l); } + r.appendChild(el); bgMenu.appendChild(r); + }; + bgRow("", liveLabel); bgRow("", grabBtn); bgRow("Brightness", bgSlider); + bgRow("Guide", guideSel); bgRow("Grid size", gridSlider); bgRow("", snapLabel); + bgRow("Line color", guideColor); bgRow("Line opacity", opacitySlider); + document.body.appendChild(bgMenu); + node._bgMenu = bgMenu; + const bgDismiss = outsideDismiss(bgMenu, () => closeBgMenu(), bgBtn); + function closeBgMenu() { bgMenu.style.display = "none"; bgDismiss.disarm(); } + bgBtn.addEventListener("click", () => { + if (bgMenu.style.display !== "none") { closeBgMenu(); return; } + bgMenu.style.display = ""; + const r = bgBtn.getBoundingClientRect(); + bgMenu.style.left = Math.max(4, Math.min(r.left, window.innerWidth - bgMenu.offsetWidth - 4)) + "px"; + bgMenu.style.top = Math.min(r.bottom + 4, window.innerHeight - bgMenu.offsetHeight - 4) + "px"; + bgDismiss.arm(); + }); + // ── Text style popup: show/hide, outline, font size, auto-placement ── + const txtBtn = document.createElement("button"); + txtBtn.className = "kjideo-btn"; txtBtn.textContent = "Text ▾"; + txtBtn.title = "Region text: show/hide, outline, font size, overlap-avoiding placement"; + stopProp(txtBtn); + const txtToggle = (prop, label, title) => { + const l = document.createElement("label"); + l.style.cssText = "display:flex;align-items:center;gap:4px;cursor:pointer;"; l.title = title || ""; + const cb = document.createElement("input"); cb.type = "checkbox"; cb.checked = node.properties[prop] !== false; + cb.addEventListener("change", () => { node.properties[prop] = cb.checked; drawCanvas(); }); + l.appendChild(cb); l.appendChild(document.createTextNode(label)); + l._cb = cb; return l; + }; + const showLbl = txtToggle("showBoxText", "show text", "Draw each region's text inside its box"); + const strokeLbl = txtToggle("textStroke", "outline", "Dark halo behind the text for legibility"); + const autoLbl = txtToggle("textAutoPlace", "auto-place", "Stagger overlapping labels to reduce overlap"); + const sizeSlider = document.createElement("input"); + sizeSlider.type = "range"; sizeSlider.min = "8"; sizeSlider.max = "22"; sizeSlider.step = "1"; + sizeSlider.value = node.properties.textSize || 12; + sizeSlider.style.cssText = "width:90px;flex:0 0 auto;"; + sizeSlider.title = "Region text font size"; + stopProp(sizeSlider); + sizeSlider.addEventListener("input", () => { node.properties.textSize = parseInt(sizeSlider.value, 10); drawCanvas(); }); + const boxOpacSlider = document.createElement("input"); + boxOpacSlider.type = "range"; boxOpacSlider.min = "0"; boxOpacSlider.max = "100"; boxOpacSlider.step = "1"; + boxOpacSlider.value = node.properties.boxOpacity == null ? 14 : node.properties.boxOpacity; + boxOpacSlider.style.cssText = "width:90px;flex:0 0 auto;"; + boxOpacSlider.title = "Box fill opacity"; + stopProp(boxOpacSlider); + boxOpacSlider.addEventListener("input", () => { node.properties.boxOpacity = parseInt(boxOpacSlider.value, 10); drawCanvas(); }); + const txtMenu = document.createElement("div"); + txtMenu.className = "kjideo-menu kjideo-bgmenu"; + txtMenu.style.display = "none"; + const txtRow = (labelText, el) => { + const r = document.createElement("div"); r.className = "kjideo-bgrow"; + if (labelText) { const l = document.createElement("span"); l.className = "kjideo-bglbl"; l.textContent = labelText; r.appendChild(l); } + r.appendChild(el); txtMenu.appendChild(r); + }; + txtRow("", showLbl); txtRow("Font size", sizeSlider); txtRow("Box opacity", boxOpacSlider); txtRow("", strokeLbl); txtRow("", autoLbl); + document.body.appendChild(txtMenu); + node._txtMenu = txtMenu; + const txtDismiss = outsideDismiss(txtMenu, () => closeTxtMenu(), txtBtn); + function closeTxtMenu() { txtMenu.style.display = "none"; txtDismiss.disarm(); } + txtBtn.addEventListener("click", () => { + if (txtMenu.style.display !== "none") { closeTxtMenu(); return; } + txtMenu.style.display = ""; + const r = txtBtn.getBoundingClientRect(); + txtMenu.style.left = Math.max(4, Math.min(r.left, window.innerWidth - txtMenu.offsetWidth - 4)) + "px"; + txtMenu.style.top = Math.min(r.bottom + 4, window.innerHeight - txtMenu.offsetHeight - 4) + "px"; + txtDismiss.arm(); + }); + // ── Templates popup: save/load named caption JSONs (server-side userdata) ── + const tplBtn = document.createElement("button"); + tplBtn.className = "kjideo-btn"; tplBtn.textContent = "Templates ▾"; + tplBtn.title = "Save / load the caption JSON as templates (stored on the server ComfyUI\\user\\default\\kjnodes\\ideogram4)"; + stopProp(tplBtn); + const tplMenu = document.createElement("div"); + tplMenu.className = "kjideo-menu kjideo-bgmenu"; + tplMenu.style.display = "none"; + document.body.appendChild(tplMenu); + node._tplMenu = tplMenu; + const tplDismiss = outsideDismiss(tplMenu, () => closeTplMenu(), tplBtn); + function closeTplMenu() { tplMenu.style.display = "none"; tplDismiss.disarm(); } + async function buildTplMenu() { + tplMenu.innerHTML = ""; + const saveRow = document.createElement("div"); saveRow.className = "kjideo-bgrow"; + const saveBtn = document.createElement("button"); saveBtn.className = "kjideo-btn"; saveBtn.textContent = "+ Save as…"; + saveBtn.addEventListener("click", async () => { + const name = tplSafe(window.prompt("Save template as:", "") || ""); + if (!name) return; + const existing = await listTemplateNames(); + if (existing.includes(name) && !window.confirm(`Overwrite template "${name}"?`)) return; + if (await saveTemplate(name, buildCaption())) buildTplMenu(); + }); + saveRow.appendChild(saveBtn); tplMenu.appendChild(saveRow); + const names = await listTemplateNames(); + if (!names.length) { + const empty = document.createElement("div"); empty.className = "kjideo-mhdr"; empty.textContent = "No templates saved."; + tplMenu.appendChild(empty); + } + for (const name of names) { + const row = document.createElement("div"); row.className = "kjideo-bgrow kjideo-trow"; + const txt = document.createElement("span"); + txt.className = "kjideo-ltext"; txt.style.cssText = "flex:1 1 auto;cursor:pointer;"; txt.textContent = name; txt.title = "Load " + name + " (replaces everything)"; + const ins = document.createElement("button"); ins.className = "kjideo-lbtn"; ins.textContent = "⊞"; ins.title = "Insert this template's boxes only into the current canvas"; + const upd = document.createElement("button"); upd.className = "kjideo-lbtn"; upd.textContent = "⤓"; upd.title = "Save current (overwrite)"; + const del = document.createElement("button"); del.className = "kjideo-lbtn del"; del.textContent = "✕"; del.title = "Delete template"; + row.append(txt, ins, upd, del); tplMenu.appendChild(row); + txt.addEventListener("click", async () => { + const cap = tryParseCaption(await loadTemplate(name)); + if (!cap) { window.alert("That template isn't a valid caption JSON."); return; } + loadCaption(cap); closeTplMenu(); + }); + ins.addEventListener("click", async (e) => { + e.stopPropagation(); + const cap = tryParseCaption(await loadTemplate(name)); + if (!cap) { window.alert("That template isn't a valid caption JSON."); return; } + insertCaptionBoxes(cap); closeTplMenu(); + }); + upd.addEventListener("click", async (e) => { + e.stopPropagation(); + if (await saveTemplate(name, buildCaption())) buildTplMenu(); + }); + del.addEventListener("click", async (e) => { + e.stopPropagation(); + if (!window.confirm(`Delete template "${name}"?`)) return; + await deleteTemplate(name); buildTplMenu(); + }); + } + } + tplBtn.addEventListener("click", async () => { + if (tplMenu.style.display !== "none") { closeTplMenu(); return; } + tplMenu.style.display = ""; + await buildTplMenu(); + const r = tplBtn.getBoundingClientRect(); + tplMenu.style.left = Math.max(4, Math.min(r.left, window.innerWidth - tplMenu.offsetWidth - 4)) + "px"; + tplMenu.style.top = Math.min(r.bottom + 4, window.innerHeight - tplMenu.offsetHeight - 4) + "px"; + tplDismiss.arm(); + }); + bar.appendChild(hint); bar.appendChild(tokenSpan); bar.appendChild(compactLbl); bar.appendChild(bgBtn); bar.appendChild(txtBtn); bar.appendChild(copyBtn); bar.appendChild(importBtn); bar.appendChild(tplBtn); bar.appendChild(clearBtn); + updateGrabBtn(); + + // Persistent global style-palette row + const styleBar = document.createElement("div"); + styleBar.className = "kjideo-bar"; + const styleLbl = document.createElement("span"); + styleLbl.textContent = "Style colors:"; + styleBar.appendChild(styleLbl); + + const canvasEl = document.createElement("canvas"); + canvasEl.className = "kjideo-canvas"; + canvasEl.tabIndex = 0; // focusable, so it can receive key events + canvasEl.title = "Drag to draw · Ctrl-drag force-draw over a box · click to select · shift-drag marquee-select · " + + "shift-click toggle · drag a group to move all · alt-click overlap · dbl-click edit · right-click region list · " + + "Del remove (all selected) · Ctrl/Cmd+C/V/D copy/paste/duplicate · H hide boxes (view)"; + const ctx = canvasEl.getContext("2d"); + addWheelPassthrough(wrap); + addMiddleClickPan(canvasEl); + + const panel = document.createElement("div"); + panel.className = "kjideo-panel"; + node._panelH = node._panelH || 150; // height of the prompt panel (set by the splitter) + panel.style.height = node._panelH + "px"; + + // Canvas letterbox-fits into a flex container (cvBox), so node resize never aspect-locks the height. + const cvBox = document.createElement("div"); + cvBox.className = "kjideo-cv"; + cvBox.appendChild(canvasEl); + + // Draggable separator between the canvas and the prompt panel — drag up for more prompt, down for more canvas. + const splitter = document.createElement("div"); + splitter.className = "kjideo-split"; + splitter.title = "Drag to resize the description area"; + splitter.addEventListener("pointerdown", (e) => { + if (e.button !== 0) return; + e.preventDefault(); e.stopPropagation(); + const scale = (node.properties.dockPinned && app.canvas) ? (app.canvas.ds.scale || 1) : 1; + const sy = e.clientY, h0 = panel.offsetHeight; + dragPointer(e, splitter, (me) => { + const dy = (me.clientY - sy) / scale; + const h = Math.max(60, Math.min(h0 - dy, wrap.clientHeight - 110)); // up → bigger panel; keep room for the canvas + node._panelH = Math.round(h); panel.style.height = node._panelH + "px"; // persisted via the onSerialize blob + fitCanvas(); + }, flushChange); + }); + + wrap.appendChild(bar); wrap.appendChild(styleBar); wrap.appendChild(cvBox); wrap.appendChild(splitter); wrap.appendChild(panel); + + const EDITOR_MIN = 220; // fixed floor; the node fills the rest and resizes freely + node.ideoEditor = node.addDOMWidget("ideo_editor", "Ideogram4Editor", wrap, { + serialize: false, hideOnZoom: false, + getMinHeight: () => (node._detached ? 0 : EDITOR_MIN), // collapse when popped out + }); + node.resizable = true; + + // DOM widgets are HTML layered over the canvas; ComfyUI only repositions/clips them + // during a foreground draw. When the node returns from off-screen the element can + // briefly float over the canvas until the next draw — force one when it re-enters view. + try { + node._visObserver = new IntersectionObserver((entries) => { + if (entries.some((en) => en.isIntersecting)) { + app.canvas?.setDirtyCanvas?.(true, true); + drawCanvas(); + } + }); + node._visObserver.observe(wrap); + } catch (e) {} + + // Letterbox-fit the canvas (keeping width/height aspect) into its flex container cvBox. + function fitCanvas() { + if (!node._detached && wrap.offsetParent === null) return; // hidden tab — measurements are 0 + const availW = cvBox.clientWidth, availH = cvBox.clientHeight; + if (availW < 4 || availH < 4) return; + const aspect = (wWidget?.value || 1) / (hWidget?.value || 1); + let cw = availW, ch = cw / aspect; + if (ch > availH) { ch = availH; cw = ch * aspect; } + canvasEl.style.width = Math.round(cw) + "px"; + canvasEl.style.height = Math.round(ch) + "px"; + drawCanvas(); + } + const fitFsCanvas = fitCanvas; // dock/fullscreen use the same fitter + function syncCanvasToDims() { fitCanvas(); } // aspect changed (width/height widgets) + function fitNode() { + if (node._detached) { fitCanvas(); return; } + if (wrap.offsetParent === null) return; // hidden tab / off-screen — skip + const minH = node.computeSize()[1]; // never sit below the floor + if (node.size[1] < minH) node.setSize([node.size[0], minH]); + fitCanvas(); + } + function detachInto(container) { // move the editor out of the node + node._wrapHome = wrap.parentNode; + container.appendChild(wrap); + if (node.ideoEditor) node.ideoEditor.hidden = true; + node._detached = true; + node._fsInner = container; + node._dockWrap = wrap; // so the dock can reclaim it if Vue re-mounts the widget + } + function reattach() { // fullscreen exit → back into the floating dock + canvasEl.style.width = ""; canvasEl.style.height = ""; + const home = node._dockBody || node._wrapHome; + if (home) home.appendChild(wrap); + node._fsInner = node._dockBody || null; // stays detached, living in the dock + if (node.graph) node.graph.setDirtyCanvas(true, true); + requestAnimationFrame(fitCanvas); + } + + function onFsEsc(e) { if (e.key === "Escape") { e.preventDefault(); e.stopPropagation(); exitFs(); } } + function enterFs() { + if (node._fullscreen) return; + node._fullscreen = true; + const ov = document.createElement("div"); ov.className = "kjideo-fs"; + const inner = document.createElement("div"); inner.className = "kjideo-fs-inner"; + ov.appendChild(inner); + ov.addEventListener("mousedown", (e) => { if (e.target === ov) exitFs(); }); // backdrop closes + document.body.appendChild(ov); + node._fsOverlay = ov; + detachInto(inner); + const closeBtn = document.createElement("button"); // visible exit (⛶ in the dock header is hidden here) + closeBtn.className = "kjideo-btn kjideo-fs-close"; closeBtn.textContent = "✕"; closeBtn.title = "Close (Esc)"; + stopProp(closeBtn); closeBtn.addEventListener("click", exitFs); + ov.appendChild(closeBtn); // on the backdrop corner, clear of the editor + document.addEventListener("keydown", onFsEsc, true); + window.addEventListener("resize", fitFsCanvas); + requestAnimationFrame(fitFsCanvas); + } + function exitFs() { + if (!node._fullscreen) return; + node._fullscreen = false; + document.removeEventListener("keydown", onFsEsc, true); + window.removeEventListener("resize", fitFsCanvas); + node._fsOverlay?.remove(); node._fsOverlay = null; + reattach(); + } + + // ── floating / dockable editor panel ── + // Persist the panel geometry — screen-space rect when unpinned, graph-space size when pinned. + function saveDockGeom() { + const fl = node._dockEl; if (!fl) return; + if (node.properties.dockMin) return; // collapsed: don't persist the header-only size + if (node.properties.dockPinned) { + node.properties.dockGraph = node.properties.dockGraph || { x: 0, y: 0 }; + node.properties.dockGraph.w = fl.offsetWidth; node.properties.dockGraph.h = fl.offsetHeight; + } else { + const r = fl.getBoundingClientRect(); + node.properties.dockRect = { x: Math.round(r.left), y: Math.round(r.top), w: Math.round(r.width), h: Math.round(r.height) }; + } + } + // Nudge ComfyUI's change tracker (it snapshots on mouseup, which our preventDefault'd drags suppress). + function flushChange() { try { window.dispatchEvent(new MouseEvent("mouseup")); } catch (e) {} } + function setPinned(on, pinBtn) { + const fl = node._dockEl, c = app.canvas; if (!fl || !c) return; + const rect = c.canvas.getBoundingClientRect(), ds = c.ds; + if (on) { // return to where it was pinned (the saved dockGraph) + let g = node.properties.dockGraph; + if (!g) { // never pinned before — derive from the current float position + const r = fl.getBoundingClientRect(); + const gx = (r.left - rect.left) / ds.scale - ds.offset[0], gy = (r.top - rect.top) / ds.scale - ds.offset[1]; + g = node.properties.dockGraph = { + x: gx - (node.pos?.[0] ?? 0), y: gy - (node.pos?.[1] ?? 0), + w: Math.round(fl.offsetWidth / ds.scale), h: Math.round(fl.offsetHeight / ds.scale), + }; + } + node.properties.dockPinned = true; + fl.style.left = ""; fl.style.top = ""; + fl.style.width = g.w + "px"; fl.style.height = g.h + "px"; + node._dockSig = ""; node._dockNodeEl = null; // force a fresh place + re-parent + pinnedDocks.add(node); startDockLoop(); + applyDockTransform(node); // place this frame (avoid a flash at the old float spot) + } else { // freeze current visual rect as a body-fixed screen panel + const r = fl.getBoundingClientRect(); + node.properties.dockPinned = false; + pinnedDocks.delete(node); + node._dockSig = ""; node._dockNodeEl = null; + document.body.appendChild(fl); // back out of the node element + fl.style.position = "fixed"; fl.style.transform = ""; fl.style.transformOrigin = ""; fl.style.zIndex = ""; + fl.style.left = Math.round(r.left) + "px"; fl.style.top = Math.round(r.top) + "px"; + fl.style.width = Math.round(r.width) + "px"; fl.style.height = Math.round(r.height) + "px"; + node.properties.dockRect = { x: Math.round(r.left), y: Math.round(r.top), w: Math.round(r.width), h: Math.round(r.height) }; + } + if (pinBtn) { pinBtn.classList.toggle("active", on); pinBtn.title = on ? "Unpin from canvas (float in screen)" : "Pin to canvas (move/zoom with the graph)"; } + c.setDirtyCanvas(true, true); requestAnimationFrame(fitFsCanvas); + flushChange(); // click's mouseup already fired before this handler — snapshot the new state + } + // Resize from any edge/corner; N/W edges shift the anchor. Units: graph px pinned (÷scale), else screen px. + function startDockResize(e, dir, fl) { + if (e.button !== 0) return; + e.preventDefault(); e.stopPropagation(); + const pinned = !!node.properties.dockPinned; + const scale = (pinned && app.canvas) ? (app.canvas.ds.scale || 1) : 1; + const sx = e.clientX, sy = e.clientY, w0 = fl.offsetWidth, h0 = fl.offsetHeight; + const gx0 = pinned ? node.properties.dockGraph.x : (parseFloat(fl.style.left) || 0); + const gy0 = pinned ? node.properties.dockGraph.y : (parseFloat(fl.style.top) || 0); + const MINW = 300, MINH = 240; + dragPointer(e, e.currentTarget, (me) => { + const dx = (me.clientX - sx) / scale, dy = (me.clientY - sy) / scale; + let w = w0, h = h0, gx = gx0, gy = gy0; + if (dir.indexOf("e") >= 0) w = w0 + dx; + if (dir.indexOf("s") >= 0) h = h0 + dy; + if (dir.indexOf("w") >= 0) { w = w0 - dx; gx = gx0 + dx; } + if (dir.indexOf("n") >= 0) { h = h0 - dy; gy = gy0 + dy; } + if (w < MINW) { if (dir.indexOf("w") >= 0) gx -= (MINW - w); w = MINW; } + if (h < MINH) { if (dir.indexOf("n") >= 0) gy -= (MINH - h); h = MINH; } + fl.style.width = Math.round(w) + "px"; fl.style.height = Math.round(h) + "px"; + if (pinned) { node.properties.dockGraph.x = gx; node.properties.dockGraph.y = gy; node._dockSig = ""; applyDockTransform(node); } + else { fl.style.left = Math.round(gx) + "px"; fl.style.top = Math.round(gy) + "px"; } + fitFsCanvas(); + }, () => { saveDockGeom(); flushChange(); }); + } + function addDockResizeHandles(fl) { + for (const dir of ["n", "s", "e", "w", "ne", "nw", "se", "sw"]) { + const h = document.createElement("div"); + h.className = "kjideo-rsz " + dir; + h.addEventListener("pointerdown", (e) => startDockResize(e, dir, fl)); + fl.appendChild(h); + } + } + function undock() { + if (node._docked) return; + if (node._fullscreen) exitFs(); + node._docked = true; node.properties.docked = true; + const fl = document.createElement("div"); fl.className = "kjideo-dock"; + fl.dataset.captureWheel = "true"; // 2.0: let focused inputs in the dock scroll instead of zooming the graph + stopProp(fl); // hosted inside the node element — don't let dock interactions drag/zoom the node + const head = document.createElement("div"); head.className = "kjideo-dock-head"; + const title = document.createElement("span"); title.textContent = "Ideogram 4 editor"; title.style.flex = "1"; + const minBtn = document.createElement("button"); minBtn.className = "kjideo-btn"; + stopProp(minBtn); + const applyMin = (on) => { // collapse the editor to just the header bar + node.properties.dockMin = !!on; + fl.classList.toggle("minimized", !!on); + minBtn.textContent = on ? "▢" : "—"; + minBtn.title = on ? "Restore editor" : "Minimize editor"; + if (!on) requestAnimationFrame(fitFsCanvas); // re-letterbox the canvas after it's visible again + }; + minBtn.addEventListener("click", () => { applyMin(!node.properties.dockMin); flushChange(); }); + const fsBtn = document.createElement("button"); fsBtn.className = "kjideo-btn"; fsBtn.textContent = "⛶"; + fsBtn.title = "Open in a larger window (Esc to close)"; + stopProp(fsBtn); fsBtn.addEventListener("click", () => node._fullscreen ? exitFs() : enterFs()); + const pinBtn = document.createElement("button"); pinBtn.className = "kjideo-btn"; pinBtn.textContent = "📌"; + stopProp(pinBtn); pinBtn.addEventListener("click", () => setPinned(!node.properties.dockPinned, pinBtn)); + head.append(title, fsBtn, minBtn, pinBtn); + const body = document.createElement("div"); body.className = "kjideo-dock-body"; + fl.append(head, body); + addDockResizeHandles(fl); + document.body.appendChild(fl); + node._dockEl = fl; node._dockBody = body; node._dockSig = ""; node._dockNodeEl = null; liveDocks.add(node); + applyDockTheme(node); // match the node's color theme right away + detachInto(body); + // initial geometry / mode + const pinned = !!node.properties.dockPinned; + if (pinned) { + const g = node.properties.dockGraph || { x: 0, y: 0, w: 540, h: 470 }; + fl.style.position = "absolute"; fl.style.width = g.w + "px"; fl.style.height = g.h + "px"; + pinnedDocks.add(node); startDockLoop(); + applyDockTransform(node); // parent + place this frame (avoid a flash at 0,0) + pinBtn.classList.add("active"); + } else { + const rc = node.properties.dockRect || { x: 90, y: 90, w: 540, h: 470 }; + fl.style.left = Math.max(0, Math.min(rc.x, window.innerWidth - 80)) + "px"; + fl.style.top = Math.max(0, Math.min(rc.y, window.innerHeight - 40)) + "px"; + fl.style.width = rc.w + "px"; fl.style.height = rc.h + "px"; + } + pinBtn.title = pinned ? "Unpin from canvas (float in screen)" : "Pin to canvas (move/zoom with the graph)"; + applyMin(!!node.properties.dockMin); // restore minimized state + // drag the panel by its header (graph-space when pinned, screen-space otherwise) + head.addEventListener("pointerdown", (e) => { + if (e.target === pinBtn || e.target === minBtn || e.target === fsBtn || e.button !== 0) return; + e.preventDefault(); + const sx0 = e.clientX, sy0 = e.clientY; + const pinned = node.properties.dockPinned; + const scale = pinned ? (app.canvas.ds.scale || 1) : 1; + const gx0 = pinned ? node.properties.dockGraph.x : 0, gy0 = pinned ? node.properties.dockGraph.y : 0; + const ox = fl.offsetLeft, oy = fl.offsetTop; + dragPointer(e, head, (me) => { + if (pinned) { // move the panel by writing its transform directly (no full redraw) + node.properties.dockGraph.x = gx0 + (me.clientX - sx0) / scale; + node.properties.dockGraph.y = gy0 + (me.clientY - sy0) / scale; + applyDockTransform(node); + } else { + fl.style.left = Math.max(0, Math.min(ox + me.clientX - sx0, window.innerWidth - 60)) + "px"; + fl.style.top = Math.max(0, Math.min(oy + me.clientY - sy0, window.innerHeight - 30)) + "px"; + } + }, () => { saveDockGeom(); flushChange(); }); + }); + try { node._dockRO = new ResizeObserver(() => { fitFsCanvas(); saveDockGeom(); }); node._dockRO.observe(fl); } catch (e) {} + // editor slot is collapsed (getMinHeight→0) — shrink the node to its widgets + requestAnimationFrame(() => { node.setSize([node.size[0], node.computeSize()[1]]); fitFsCanvas(); }); + } + // Pin the editor under the node. `fresh` (new node only) matches the node width; reloads honor the blob. + function ensureDocked(fresh) { + if (node._docked) return; + if (node.properties.dockPinned == null) node.properties.dockPinned = true; + const honor = !!node.properties.dockGraph && node._dockGeomRestored; // only trust blob-restored geometry + if (!honor) { + // old workflow: reuse its saved (aspect-locked) node height so the canvas isn't tiny + const sh = (!fresh && node._savedSize && node._savedSize[1] > 240) ? Math.round(node._savedSize[1]) : 470; + node.properties.dockGraph = { x: 0, y: 0, w: 560, h: sh }; + } + undock(); + if (!honor) requestAnimationFrame(() => { // no reliable geometry — place a default under the node + const g = node.properties.dockGraph; + g.x = 0; g.y = node.size[1] + 12; + if (fresh) { g.w = Math.round(node.size[0]); if (node._dockEl) node._dockEl.style.width = g.w + "px"; } + node._dockSig = ""; applyDockTransform(node); saveDockGeom(); + }); + } + + // ── geometry helpers ── (logical CSS px = the displayed canvas size) + function logW() { return canvasEl.offsetWidth || 1; } + function logH() { return canvasEl.offsetHeight || 1; } + function toPx(b) { + const W = logW(), H = logH(); + return { x1: b.x * W, y1: b.y * H, x2: (b.x + b.w) * W, y2: (b.y + b.h) * H }; + } + function mouseN(e) { + const r = canvasEl.getBoundingClientRect(); + return { x: (e.clientX - r.left) / r.width, y: (e.clientY - r.top) / r.height }; + } + function clamp01(v) { return Math.max(0, Math.min(1, v)); } + // greedy word-wrap to maxW px (ctx.font must be set by caller) + function wrapLines(text, maxW) { + const lines = []; + for (const para of text.split("\n")) { + let line = ""; + for (const word of para.split(/\s+/)) { + if (!word) continue; + const test = line ? line + " " + word : word; + if (line && ctx.measureText(test).width > maxW) { lines.push(line); line = word; } + else line = test; + } + lines.push(line); + } + return lines; + } + function normalizeBox(b) { + // collapse negative size to positive top-left + w/h, clamp into canvas + let x = b.x, y = b.y, w = b.w, h = b.h; + if (w < 0) { x += w; w = -w; } + if (h < 0) { y += h; h = -h; } + x = clamp01(x); y = clamp01(y); + w = Math.min(w, 1 - x); h = Math.min(h, 1 - y); + return { ...b, x, y, w: Math.max(0, w), h: Math.max(0, h) }; + } + + // All boxes under the point, top-first to match draw order: the active box is + // drawn last (on top), then the rest by index low→high (index 0 = front). + function boxesAt(mN) { + const baseRx = HANDLE / logW(), baseRy = HANDLE / logH(); + const res = []; + for (let i = 0; i < node._boxes.length; i++) { + const b = node._boxes[i]; + if (b.locked) continue; // locked regions aren't grabbable on the canvas + const rx = Math.min(baseRx, b.w / 3), ry = Math.min(baseRy, b.h / 3); // shrink handles on small boxes so a central move zone remains + const mode = rectHitTestN(mN.x, mN.y, b.x, b.y, b.x + b.w, b.y + b.h, rx, ry); + if (mode) res.push({ index: i, mode }); + } + const ai = res.findIndex((c) => c.index === node._activeIdx); + if (ai > 0) res.unshift(res.splice(ai, 1)[0]); + return res; + } + // Hover / right-click: prefer a resize handle on the active box, else topmost. + function hitTest(mN) { + const cands = boxesAt(mN); + if (!cands.length) return null; + return cands.find((c) => c.index === node._activeIdx && c.mode !== "move") || cands[0]; + } + // Tag-chip rects (canvas px), placed to avoid overlapping each other: each + // box's tag tries top-left, top-right, bottom-right, bottom-left in turn. + function tagRects() { + ctx.font = "bold 11px monospace"; + const W = logW(), H = logH(), h = 14; + const placed = [], rects = []; + const hits = (a, b) => a.x < b.x + b.w && a.x + a.w > b.x && a.y < b.y + b.h && a.y + a.h > b.y; + for (let i = 0; i < node._boxes.length; i++) { + const b = node._boxes[i]; + const x1 = b.x * W, y1 = b.y * H, x2 = (b.x + b.w) * W, y2 = (b.y + b.h) * H; + const tag = String(i + 1).padStart(2, "0"); + const w = ctx.measureText(tag).width + 8; + // Box too small to hold the badge → place it just outside (above/below/beside) so it doesn't + // cover the box and the box stays grabbable; otherwise tuck it into a free corner inside. + const small = (x2 - x1) < w || (y2 - y1) < h; + const cands = small + ? [[x1, y1 - h - 1], [x1, y2 + 1], [x2 + 1, y1], [x1 - w - 1, y1]] + : [[x1, y1], [x2 - w, y1], [x2 - w, y2 - h], [x1, y2 - h]]; + let pick = cands[0]; + for (const [cx, cy] of cands) { + if (cx < 0 || cy < 0 || cx + w > W || cy + h > H) continue; // keep on-canvas + if (!placed.some((p) => hits({ x: cx, y: cy, w, h }, p))) { pick = [cx, cy]; break; } + } + const r = { x: pick[0], y: pick[1], w, h, tag }; + placed.push(r); rects[i] = r; + } + return rects; + } + function titleAt(mN) { + const px = mN.x * logW(), py = mN.y * logH(); + const rects = tagRects(); + for (let i = node._boxes.length - 1; i >= 0; i--) { + if (node._boxes[i].locked) continue; // locked badges aren't interactive + const r = rects[i]; + if (r && px >= r.x && px <= r.x + r.w && py >= r.y && py <= r.y + r.h) return i; + } + return null; + } + // In-box text placement: when auto-place is on, put each box's text in whichever vertical slot overlaps + // least with the text already placed for other boxes (greedy), so overlapping labels don't stack. + function textBlocks() { + if (node.properties.showBoxText === false) return []; + const { fs, lh } = txtFont(); + ctx.font = fs + "px monospace"; + const W = logW(), H = logH(), pad = 4, tagH = 14, auto = node.properties.textAutoPlace !== false; + const placed = [], blocks = []; + const overlap = (a, b) => Math.max(0, Math.min(a.x + a.w, b.x + b.w) - Math.max(a.x, b.x)) * + Math.max(0, Math.min(a.y + a.h, b.y + b.h) - Math.max(a.y, b.y)); + for (let i = 0; i < node._boxes.length; i++) { + const b = node._boxes[i]; + let body = b.desc || ""; + if (b.type === "text" && b.text) body = `"${b.text}"` + (body ? " — " + body : ""); + if (!body) { blocks[i] = null; continue; } + const x1 = b.x * W, y1 = b.y * H, x2 = (b.x + b.w) * W, y2 = (b.y + b.h) * H; + const lines = wrapLines(body, (x2 - x1) - pad * 2); // full width + let lw = 0; for (const ln of lines) lw = Math.max(lw, ctx.measureText(ln).width); + const bw = Math.min(x2 - x1, lw + pad * 2), bh = Math.min(y2 - y1, lines.length * lh + 4); + // candidates = several vertical slots (more when the box is taller); text spans the width, so only Y varies + const top = y1 + ((y2 - y1) > bh + tagH ? tagH : 0), bottom = y2 - bh; + const cands = []; + if (!auto || bottom <= top) cands.push([x1, top]); // off, or fills the box → fixed top + else { + const n = Math.min(6, 2 + Math.floor((bottom - top) / (lh * 2))); + for (let k = 0; k < n; k++) cands.push([x1, top + ((bottom - top) * k) / (n - 1)]); + } + let best = cands[0], bestScore = Infinity; + for (const [cx, cy] of cands) { + const rect = { x: cx, y: cy, w: bw, h: bh }; + let s = 0; for (const p of placed) s += overlap(rect, p); + if (s < bestScore) { bestScore = s; best = [cx, cy]; if (s === 0) break; } + } + const r = { x: best[0], y: best[1], w: bw, h: bh, lines, lh, fs, pad }; + placed.push(r); blocks[i] = r; + } + return blocks; + } + // Click selection: active box's resize handle wins (corner resize); then a + // title-chip click selects that box (drawn to front); Alt-click cycles the + // overlap stack; else the topmost box. + function pickForSelection(mN, cycle) { + const cands = boxesAt(mN); + if (!cands.length) return null; + const ah = cands.find((c) => c.index === node._activeIdx && c.mode !== "move"); + if (ah && !cycle) return ah; + const ti = titleAt(mN); + if (ti !== null && !cycle) return { index: ti, mode: "move" }; + if (cycle && cands.length > 1) { + // cycle over a stable order (boxesAt moves the active box to the front, which would otherwise ping-pong 2) + const ordered = [...cands].sort((a, b) => a.index - b.index); + const pos = ordered.findIndex((c) => c.index === node._activeIdx); + return ordered[(pos + 1) % ordered.length]; + } + return cands.find((c) => c.index === node._activeIdx && c.mode !== "move") || cands[0]; + } + // normalized variant of rectHitTest with separate x/y radii + function rectHitTestN(mx, my, x1, y1, x2, y2, rx, ry) { + const h = (cx, cy) => Math.abs(mx - cx) < rx && Math.abs(my - cy) < ry; + if (h(x1, y1)) return "resize-tl"; + if (h(x2, y1)) return "resize-tr"; + if (h(x1, y2)) return "resize-bl"; + if (h(x2, y2)) return "resize-br"; + if (mx >= x1 && mx <= x2 && Math.abs(my - y1) < ry) return "resize-t"; + if (mx >= x1 && mx <= x2 && Math.abs(my - y2) < ry) return "resize-b"; + if (my >= y1 && my <= y2 && Math.abs(mx - x1) < rx) return "resize-l"; + if (my >= y1 && my <= y2 && Math.abs(mx - x2) < rx) return "resize-r"; + if (mx >= x1 && mx <= x2 && my >= y1 && my <= y2) return "move"; + return null; + } + + function applyDrag(mode, start, dN) { + const { x, y, w, h } = start; + const dx = dN.x, dy = dN.y; + if (mode === "move") { + return { ...start, x: clamp01(Math.min(x + dx, 1 - w)), y: clamp01(Math.min(y + dy, 1 - h)) }; + } + if (mode === "draw") { + return normalizeBox({ ...start, w: w + dx, h: h + dy }); + } + // resize: move only the dragged edges (clamped to the canvas); the others stay anchored. + const suf = mode.slice(7); // "tl"|"tr"|"bl"|"br"|"t"|"b"|"l"|"r" + let l = x, t = y, r = x + w, b = y + h; + if (suf.includes("l")) l = clamp01(l + dx); + if (suf.includes("r")) r = clamp01(r + dx); + if (suf.includes("t")) t = clamp01(t + dy); + if (suf.includes("b")) b = clamp01(b + dy); + if (r < l) [l, r] = [r, l]; // crossing the opposite edge flips, no anchor drift + if (b < t) [t, b] = [b, t]; + return { ...start, x: l, y: t, w: r - l, h: b - t }; + } + + // ── grid / guides ── + function gridN() { return Math.max(2, Math.min(128, node.properties.gridSize || 10)); } + function txtFont() { const fs = Math.max(6, Math.min(40, node.properties.textSize || 12)); return { fs, lh: Math.round(fs * 1.2) }; } + // Whole number of cells per axis so they fill the canvas exactly (no split cells at the edges); + // near-square (the target size is rounded to fit). Shared by the grid guide and snap-to-grid. + function gridStep() { + const W = logW(), H = logH(), target = Math.min(W, H) / gridN(); + const nx = Math.max(1, Math.round(W / target)), ny = Math.max(1, Math.round(H / target)); + return { nx, ny, cw: W / nx, ch: H / ny, sx: 1 / nx, sy: 1 / ny }; + } + function guideStroke(a) { + const c = hexRgb(node.properties.guideColor || "#ffffff") || { r: 255, g: 255, b: 255 }; + const op = (node.properties.guideOpacity == null ? 100 : node.properties.guideOpacity) / 100; + return `rgba(${c.r},${c.g},${c.b},${(a * op).toFixed(3)})`; + } + // Snap a box's free edges to the (square) grid when snap-to-grid is on. + function snapBox(b, mode) { + if (!node.properties.snap) return b; + const { sx, sy } = gridStep(); + const sn = (v, s) => Math.round(v / s) * s; + const { x, y, w, h } = b; + if (mode === "move") { // snap position, preserve size + return { ...b, x: clamp01(Math.min(sn(x, sx), 1 - w)), y: clamp01(Math.min(sn(y, sy), 1 - h)) }; + } + const x2 = sn(x + w, sx), y2 = sn(y + h, sy); // draw/resize: snap the edges + return normalizeBox({ ...b, x: sn(x, sx), y: sn(y, sy), w: x2 - sn(x, sx), h: y2 - sn(y, sy) }); + } + + // ── drawing ── + // Golden spiral: largest φ-rectangle fitting the canvas, subdivided into squares with quarter arcs. + function goldenSpiral(W, H) { + const phi = 1.6180339887; + let w, h; + if (W >= H) { if (W / H >= phi) { h = H; w = H * phi; } else { w = W; h = W / phi; } } + else { if (H / W >= phi) { w = W; h = W * phi; } else { h = H; w = H / phi; } } + let x = (W - w) / 2, y = (H - h) / 2; + ctx.save(); + ctx.strokeStyle = guideStroke(0.25); ctx.lineWidth = 1; + ctx.strokeRect(x + 0.5, y + 0.5, w, h); + ctx.strokeStyle = guideStroke(0.6); ctx.lineWidth = 1.5; + let phase = w >= h ? 0 : 1; + for (let i = 0; i < 12 && w > 1 && h > 1; i++, phase = (phase + 1) % 4) { + const s = Math.min(w, h); + let cx, cy, a0, a1; + if (phase === 0) { cx = x + s; cy = y + s; a0 = Math.PI; a1 = Math.PI * 1.5; x += s; w -= s; } + else if (phase === 1) { cx = x; cy = y + s; a0 = Math.PI * 1.5; a1 = Math.PI * 2; y += s; h -= s; } + else if (phase === 2) { cx = x + w - s; cy = y; a0 = 0; a1 = Math.PI * 0.5; w -= s; } + else { cx = x + w; cy = y + h - s; a0 = Math.PI * 0.5; a1 = Math.PI; h -= s; } + ctx.beginPath(); ctx.arc(cx, cy, s, a0, a1); ctx.stroke(); + } + ctx.restore(); + } + // Composition guide overlay (rule of thirds / grid / golden ratio / spiral), drawn on the bg. + function drawGuide(W, H) { + const kind = node.properties.guide; + if (!kind || kind === "none") return; + if (kind === "spiral") { goldenSpiral(W, H); return; } + ctx.save(); + ctx.lineWidth = 1; + if (kind === "grid") { // whole cells across both axes (no split cells at edges) + ctx.strokeStyle = guideStroke(0.24); + const { nx, ny, cw, ch } = gridStep(); + for (let i = 1; i < nx; i++) { const X = Math.round(i * cw) + 0.5; ctx.beginPath(); ctx.moveTo(X, 0); ctx.lineTo(X, H); ctx.stroke(); } + for (let j = 1; j < ny; j++) { const Y = Math.round(j * ch) + 0.5; ctx.beginPath(); ctx.moveTo(0, Y); ctx.lineTo(W, Y); ctx.stroke(); } + ctx.restore(); return; + } + const v = [], h = []; // thirds / golden: per-axis fractions (correct at any aspect) + if (kind === "thirds") { v.push(1 / 3, 2 / 3); h.push(1 / 3, 2 / 3); } + else if (kind === "golden") { const g = 1 / 1.6180339887; v.push(1 - g, g); h.push(1 - g, g); } + ctx.strokeStyle = guideStroke(kind === "golden" ? 0.45 : 0.28); + for (const fx of v) { const px = Math.round(fx * W) + 0.5; ctx.beginPath(); ctx.moveTo(px, 0); ctx.lineTo(px, H); ctx.stroke(); } + for (const fy of h) { const py = Math.round(fy * H) + 0.5; ctx.beginPath(); ctx.moveTo(0, py); ctx.lineTo(W, py); ctx.stroke(); } + ctx.restore(); + } + let _rafPending = false; + function drawCanvas() { + if (_rafPending) return; + _rafPending = true; + requestAnimationFrame(() => { + _rafPending = false; + _draw(); + }); + } + function _draw() { + // Size the backing store to display × DPR and draw in logical px (crisp text/lines). + const W = logW(), H = logH(), d = window.devicePixelRatio || 1; + const bw = Math.round(W * d), bh = Math.round(H * d); + if (canvasEl.width !== bw || canvasEl.height !== bh) { canvasEl.width = bw; canvasEl.height = bh; } + ctx.setTransform(d, 0, 0, d, 0, 0); + ctx.clearRect(0, 0, W, H); + let bri = bgBrightnessWidget ? bgBrightnessWidget.value : 25; + if (typeof bri !== "number" || isNaN(bri)) bri = 25; // guard against unset widget value + if (node._bgImg) { // reference image, dimmed by brightness + ctx.drawImage(node._bgImg, 0, 0, W, H); + const dim = 1 - bri / 100; + if (dim > 0) { ctx.fillStyle = `rgba(0,0,0,${dim})`; ctx.fillRect(0, 0, W, H); } + } else { // blank canvas grey from brightness + const g = Math.round(bri / 100 * 128); + ctx.fillStyle = `rgb(${g},${g},${g})`; ctx.fillRect(0, 0, W, H); + } + drawGuide(W, H); // composition guide overlay + if (node._hideBoxes) return; // H: temporary background-only view + // selection highlight only when the editor is focused or the node is selected + const showSel = node._focused || node._selected; + const aIdx = showSel ? node._activeIdx : -1; + const selSet = new Set(showSel ? node._selection : []); + if (aIdx >= 0) selSet.add(aIdx); + // index 0 = front (drawn last); selected drawn above non-selected, active last of all + const nonSel = node._boxes.map((_, i) => i).filter((i) => !selSet.has(i)).reverse(); + const selOthers = node._boxes.map((_, i) => i).filter((i) => selSet.has(i) && i !== aIdx).reverse(); + const order = [...nonSel, ...selOthers]; + if (aIdx >= 0 && aIdx < node._boxes.length) order.push(aIdx); + const tagR = tagRects(); // collision-avoided tag positions + const textB = textBlocks(); // collision-avoided in-box text positions + for (const i of order) { + const b = node._boxes[i], active = i === aIdx, selected = selSet.has(i) && !b.locked; // locked never shows as selected + const pal = (b.palette || []).filter(Boolean); + const col = pal.length ? pal[0] : "#8c8c8c"; // box color = first palette color, else neutral grey + const { x1, y1, x2, y2 } = toPx(b); + const w = x2 - x1, h = y2 - y1; + const hovered = (i === node._hoverBox && !b.locked) || selected; // locked boxes don't hover; selected stay highlighted + if (selected) { // opaque backing so contents read clearly over boxes behind + ctx.fillStyle = "rgba(26,26,26,0.88)"; + ctx.fillRect(x1, y1, w, h); + } + const baseA = (node.properties.boxOpacity == null ? 14 : node.properties.boxOpacity) / 100; + const fillA = Math.min(1, hovered ? baseA + 0.1 : baseA); // box-color tint at the chosen opacity + ctx.fillStyle = col + Math.round(fillA * 255).toString(16).padStart(2, "0"); + ctx.fillRect(x1, y1, w, h); + if (b.locked) ctx.setLineDash([3, 3]); // locked: frozen on the canvas + else if (b.nobbox) ctx.setLineDash([6, 4]); // unplaced (no bbox in source) + const lw = selected ? 2 : (hovered ? 1.5 : 1); + ctx.strokeStyle = col; ctx.lineWidth = lw; + ctx.strokeRect(x1 + lw / 2, y1 + lw / 2, w - lw, h - lw); // inside the box so strip/badge align at y1 + ctx.setLineDash([]); + if (pal.length) { // palette shown as a strip along the top edge + const sw = w / pal.length, n = pal.length, sh = 7; + for (let p = 0; p < n; p++) { + const sx = x1 + Math.round(p * sw); + ctx.fillStyle = pal[p]; + ctx.fillRect(sx, y1, x1 + Math.round((p + 1) * sw) - sx, sh); + } + } + // in-box content (clipped to the box): prompt text (at its de-conflicted corner) + lock badge + ctx.save(); + ctx.beginPath(); ctx.rect(x1, y1, w, h); ctx.clip(); + const tb = textB[i]; + if (tb) { + const stroke = node.properties.textStroke !== false; + ctx.font = tb.fs + "px monospace"; + if (stroke) { ctx.lineWidth = 3; ctx.lineJoin = "round"; ctx.strokeStyle = "rgba(0,0,0,0.85)"; } // dark halo + ctx.fillStyle = readableText(col); // box color, lightened if too dark + let ty = tb.y + tb.fs; + for (const line of tb.lines) { + if (stroke) ctx.strokeText(line, tb.x + tb.pad, ty); + ctx.fillText(line, tb.x + tb.pad, ty); + ty += tb.lh; + } + } + ctx.restore(); + // tag chip on top, unclipped (sits outside boxes too small to hold it — its collision-avoided position) + const tr = tagR[i]; + ctx.font = "bold 11px monospace"; + ctx.fillStyle = col; // tag chip = box color + ctx.fillRect(tr.x, tr.y, tr.w, 14); + if (i === node._hoverTitle && !b.locked) { // hover highlight (never on locked) + ctx.fillStyle = "rgba(255,255,255,0.25)"; ctx.fillRect(tr.x, tr.y, tr.w, 14); + ctx.strokeStyle = "#fff"; ctx.lineWidth = 1; ctx.strokeRect(tr.x + 0.5, tr.y + 0.5, tr.w - 1, 13); + } + ctx.fillStyle = textOn(col); + ctx.fillText(tr.tag, tr.x + 4, tr.y + 11); + if (b.locked) { ctx.font = "11px sans-serif"; ctx.fillStyle = "#ddd"; ctx.fillText("🔒", tr.x + tr.w + 2, tr.y + 11); } // lock symbol next to the badge + if (selected) { // orange selection ring on top (above strip/tag): solid = primary, dashed = others + const olw = active ? 2 : 1; + ctx.strokeStyle = "#ff8c00"; ctx.lineWidth = olw; + if (!active) ctx.setLineDash([5, 3]); + ctx.strokeRect(x1 + olw / 2, y1 + olw / 2, w - olw, h - olw); + ctx.setLineDash([]); + } + } + const apIdx = node._altPreview ?? -1; // Alt-held: ring the box an alt-click would select next + if (apIdx >= 0 && apIdx < node._boxes.length) { + const { x1, y1, x2, y2 } = toPx(node._boxes[apIdx]); + ctx.strokeStyle = "#46b4e6"; ctx.lineWidth = 2; ctx.setLineDash([5, 3]); + ctx.strokeRect(x1 + 1, y1 + 1, (x2 - x1) - 2, (y2 - y1) - 2); + ctx.setLineDash([]); + } + if (node._marquee && node._marqueeActive) { // rubber-band selection rectangle + const r = marqueeRect(); + const mx = r.x * W, my = r.y * H, mw = r.w * W, mh = r.h * H; + ctx.fillStyle = "rgba(70,180,230,0.12)"; + ctx.fillRect(mx, my, mw, mh); + ctx.strokeStyle = "#46b4e6"; ctx.lineWidth = 1; ctx.setLineDash([4, 3]); + ctx.strokeRect(mx + 0.5, my + 0.5, mw - 1, mh - 1); + ctx.setLineDash([]); + } + } + + // ── serialization ── + // Wired AND the source node isn't muted(2)/bypassed(4) — those keep the link but emit nothing. + const importConnected = () => { + const link = node.graph?.links?.[node.inputs?.find((i) => i.name === "import_json")?.link]; + return !!link && ![2, 4].includes(node.graph.getNodeById(link.origin_id)?.mode); + }; + function serialize() { // saved/restored value: clean boxes + if (elementsWidget) elementsWidget.value = node._boxes.length ? JSON.stringify(node._boxes) : ""; + if (stylePaletteWidget) stylePaletteWidget.value = node._stylePalette.length ? JSON.stringify(node._stylePalette) : ""; + } + // Queue-time value (not the saved value): when a wired import should drive the output — "always" + // mode, or empty in "when empty" mode — return a unique marker so ComfyUI can't cache-skip the + // node. It then re-executes and pushes the import back via ui, refreshing the editor. The server + // treats a non-list elements_data as empty (and ignores it entirely in "always" mode), so the + // nonce never affects the output. + if (elementsWidget) { + elementsWidget.serializeValue = () => { + const always = findW("import_mode")?.value === "always"; + if (importConnected() && (always || !node._boxes.length)) { + return JSON.stringify({ _refresh: (node._serialSeq = (node._serialSeq || 0) + 1) }); + } + return node._boxes.length ? JSON.stringify(node._boxes) : ""; + }; + } + + function commit() { serialize(); renderPanel(); drawCanvas(); updateTokens(); flushChange(); } // flush so canvas edits persist without a defocus + // Live text edit: persist + repaint + token count, without rebuilding the panel. + function touch() { serialize(); drawCanvas(); updateTokens(); } + + function removeBox(i) { + node._boxes.splice(i, 1); + node._selection = new Set(); + if (node._boxes.length === 0) node._activeIdx = -1; + else if (i <= node._activeIdx) node._activeIdx = Math.max(0, node._activeIdx - 1); + if (node._activeIdx >= 0) node._selection.add(node._activeIdx); + } + // Delete every selected box (or the active one) and reset the selection. + function removeSelected() { + const idxs = [...node._selection].sort((a, b) => b - a).filter((i) => !node._boxes[i]?.locked); // keep locked + if (!idxs.length && node._activeIdx >= 0 && !node._boxes[node._activeIdx]?.locked) idxs.push(node._activeIdx); + if (!idxs.length) return; + for (const i of idxs) node._boxes.splice(i, 1); + node._selection = new Set(); + node._activeIdx = node._boxes.length ? Math.min(idxs[idxs.length - 1], node._boxes.length - 1) : -1; + if (node._activeIdx >= 0) node._selection.add(node._activeIdx); + } + // Replace the selection with a single box (used by clicks / programmatic selects). + function selectOnly(idx) { + node._activeIdx = idx; + node._selection = idx >= 0 ? new Set([idx]) : new Set(); + } + + // ── pointer interaction ── + canvasEl.addEventListener("pointerdown", (e) => { + if (node._placing) { // drop the duplicate being placed + if (e.button === 0) { placeFollower(mouseN(e)); finishPlacing(); } + else cancelPlacing(); + e.preventDefault(); e.stopPropagation(); + return; + } + if (e.button !== 0) return; + if (node._hideBoxes) return; // view-only while boxes are hidden (H) + canvasEl.focus(); // so Delete/Backspace targets this editor + // capture the pointer so move/up keep coming even when the cursor leaves the node + try { canvasEl.setPointerCapture(e.pointerId); } catch (e2) {} + node._hoverTitle = null; node._hoverBox = null; // clear hover highlight while interacting + const mN = mouseN(e); + // Ctrl/Cmd forces drawing a new box even when starting over an existing one. + const hit = (e.ctrlKey || e.metaKey) ? null : pickForSelection(mN, e.altKey); + // Touch double-tap opens the inline editor (dblclick may not fire on touch). + if (e.pointerType && e.pointerType !== "mouse") { + const last = node._lastTap, now = e.timeStamp; + if (hit && last && now - last.t < 350 && Math.abs(mN.x - last.x) < 0.03 && Math.abs(mN.y - last.y) < 0.03) { + node._lastTap = null; openInlineEditor(hit.index); + e.preventDefault(); e.stopPropagation(); return; + } + node._lastTap = { t: now, x: mN.x, y: mN.y }; + } + node._pendingCollapse = -1; + node._groupStart = null; + if (e.shiftKey) { // shift-drag = marquee select; shift-click = toggle + startMarquee(mN, hit ? hit.index : -1); + e.preventDefault(); e.stopPropagation(); + return; + } + if (hit) { + if (!node._selection.has(hit.index)) node._selection = new Set([hit.index]); // outside selection → pick it + else if (node._selection.size > 1) node._pendingCollapse = hit.index; // click (no drag) collapses + node._activeIdx = hit.index; + node._dragMode = hit.mode; + node._boxAtStart = { ...node._boxes[hit.index] }; + if (node._selection.size > 1) { // snapshot the whole group for group move/resize + node._groupStart = {}; + for (const i of node._selection) node._groupStart[i] = { ...node._boxes[i] }; + } + } else { + node._dragMode = "draw"; + const nb = { x: mN.x, y: mN.y, w: 0, h: 0, type: "obj", text: "", desc: "", palette: [] }; + node._boxes.push(nb); + node._activeIdx = node._boxes.length - 1; + node._selection = new Set([node._activeIdx]); + node._boxAtStart = { ...nb }; + } + node._drawing = true; + node._dragStartN = mN; + canvasEl.addEventListener("pointermove", onMove); + canvasEl.addEventListener("pointerup", onUp); + canvasEl.addEventListener("pointercancel", onUp); // touch can cancel instead of up + e.preventDefault(); e.stopPropagation(); + drawCanvas(); // panel rebuild/resize deferred to onUp so the canvas doesn't shift mid-drag + }); + + // Hover targets: with Alt, preview what an alt-click would select next (the box under the current one). + function hoverTargets(mN, alt, force) { + if (alt && !force) { const pick = pickForSelection(mN, true); return { ti: null, hb: null, ap: pick ? pick.index : -1, hit: null }; } + const ti = force ? null : titleAt(mN); + const hit = force ? null : hitTest(mN); + return { ti, hb: ti != null ? ti : (hit ? hit.index : null), ap: -1, hit }; + } + function applyHover(t) { + if (t.ti !== node._hoverTitle || t.hb !== node._hoverBox || t.ap !== (node._altPreview ?? -1)) { + node._hoverTitle = t.ti; node._hoverBox = t.hb; node._altPreview = t.ap; drawCanvas(); + } + } + canvasEl.addEventListener("pointermove", (e) => { + node._lastMouseN = mouseN(e); // track cursor for paste-under-cursor + if (node._placing) { placeFollower(node._lastMouseN); return; } + if (node._drawing || node._marquee || node._hideBoxes) return; + const mN = node._lastMouseN; + const force = e.ctrlKey || e.metaKey, alt = e.altKey; // Ctrl/Cmd = force-draw, Alt = cycle preview + const t = hoverTargets(mN, alt, force); + applyHover(t); + canvasEl.style.cursor = (alt && !force) ? (t.ap >= 0 ? "pointer" : "crosshair") + : (t.ti != null ? "pointer" : (t.hit ? (cursorForBboxMode(t.hit.mode) || "crosshair") : "crosshair")); + }); + // Alt pressed/released without moving the mouse — refresh the preview from the last cursor position. + node._altRefresh = (altDown) => { + if (node._drawing || node._marquee || node._placing || node._hideBoxes || !node._lastMouseN) return; + applyHover(hoverTargets(node._lastMouseN, altDown, false)); + }; + canvasEl.addEventListener("pointerleave", () => { + if (hoveredCanvasNode === node) hoveredCanvasNode = null; + if (node._hoverTitle !== null || node._hoverBox !== null || (node._altPreview ?? -1) >= 0) { + node._hoverTitle = null; node._hoverBox = null; node._altPreview = -1; drawCanvas(); + } + }); + canvasEl.addEventListener("pointerenter", () => { hoveredCanvasNode = node; }); + // H (while hovering the canvas): temporary background-only view (not serialized). + node._toggleHideBoxes = () => { node._hideBoxes = !node._hideBoxes; drawCanvas(); }; + + // ── inline description editing (double-click a region) ── + let inlineTa = null; + function closeInlineEditor() { + if (inlineTa) { inlineTa.remove(); inlineTa = null; } + } + function openInlineEditor(idx) { + closeInlineEditor(); + const b = node._boxes[idx]; + if (!b) return; + selectOnly(idx); + const dw = canvasEl.offsetWidth, dh = canvasEl.offsetHeight; // CSS display size + const ox = canvasEl.offsetLeft, oy = canvasEl.offsetTop; + const w = Math.min(dw, Math.max(70, b.w * dw)); + const h = Math.min(dh, Math.max(42, b.h * dh)); + // clamp so the editor stays inside the canvas (wrapper is overflow:hidden) + const left = Math.max(ox, Math.min(ox + b.x * dw, ox + dw - w)); + const top = Math.max(oy, Math.min(oy + b.y * dh, oy + dh - h)); + const ta = document.createElement("textarea"); + ta.className = "kjideo-inline"; + ta.value = b.desc || ""; + ta.style.left = left + "px"; + ta.style.top = top + "px"; + ta.style.width = w + "px"; + ta.style.height = h + "px"; + ta.style.borderColor = (b.palette || []).find(Boolean) || "#46b4e6"; // first palette color, else accent + stopProp(ta); + wrap.appendChild(ta); + inlineTa = ta; + ta.focus(); ta.select(); + const orig = b.desc || ""; + let cancelled = false; + ta.addEventListener("input", () => { b.desc = ta.value; drawCanvas(); updateTokens(); }); + ta.addEventListener("keydown", (e) => { + e.stopPropagation(); + if (e.key === "Escape") { cancelled = true; b.desc = orig; ta.blur(); } + else if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) ta.blur(); + }); + ta.addEventListener("blur", () => { + if (!cancelled) b.desc = ta.value; + closeInlineEditor(); + commit(); + }); + } + canvasEl.addEventListener("dblclick", (e) => { + if (node._hideBoxes) return; + e.preventDefault(); e.stopPropagation(); + const cands = boxesAt(mouseN(e)); // edit the active box if it's under the cursor, else topmost + const target = cands.find((c) => c.index === node._activeIdx) || cands[0]; + if (target) openInlineEditor(target.index); + }); + + // Snapshot the current selection (or the active box) for the clipboard. + function copySelection() { + const idxs = node._selection.size ? [...node._selection].sort((a, b) => a - b) + : (node._activeIdx >= 0 ? [node._activeIdx] : []); + return idxs.map((i) => JSON.parse(JSON.stringify(node._boxes[i]))); + } + // Paste the clipboard regions as a group, centered under the cursor, keeping their layout. + function pasteBoxes() { + if (!copiedBoxes || !copiedBoxes.length) return; + const clones = copiedBoxes.map((b) => { const c = JSON.parse(JSON.stringify(b)); delete c.locked; return c; }); // pasted = editable + let minx = Infinity, miny = Infinity, maxx = -Infinity, maxy = -Infinity; + for (const b of clones) { + minx = Math.min(minx, b.x); miny = Math.min(miny, b.y); + maxx = Math.max(maxx, b.x + b.w); maxy = Math.max(maxy, b.y + b.h); + } + const gw = maxx - minx, gh = maxy - miny; + const m = node._lastMouseN; + let tx = m ? m.x - gw / 2 : minx + 0.03; // target group top-left + let ty = m ? m.y - gh / 2 : miny + 0.03; + tx = Math.max(0, Math.min(tx, 1 - gw)); // clamp group into the canvas + ty = Math.max(0, Math.min(ty, 1 - gh)); + const dx = tx - minx, dy = ty - miny; + const start = node._boxes.length; + for (const b of clones) { + b.x = clamp01(b.x + dx); b.y = clamp01(b.y + dy); + delete b.nobbox; // pasted boxes are placed + node._boxes.push(b); + } + node._selection = new Set(); + for (let i = start; i < node._boxes.length; i++) node._selection.add(i); + node._activeIdx = node._boxes.length - 1; + commit(); fitNode(); + } + // Keyboard: Delete removes; Ctrl/Cmd C/V/D copy/paste/duplicate the active region. + // Canvas must be focused; stop the event so LiteGraph doesn't act on the node. + canvasEl.addEventListener("keydown", (e) => { + if (node._placing) { + if (e.key === "Escape") { e.preventDefault(); e.stopPropagation(); cancelPlacing(); } + return; + } + if (node._drawing || node._hideBoxes) return; // view-only while boxes are hidden (H) + const ctrl = e.ctrlKey || e.metaKey; + if ((e.key === "Delete" || e.key === "Backspace") && node._activeIdx >= 0) { + e.preventDefault(); e.stopPropagation(); + removeSelected(); commit(); fitNode(); // removes all selected (or the active one) + } else if (ctrl && e.key === "c" && node._activeIdx >= 0) { + e.preventDefault(); e.stopPropagation(); + copiedBoxes = copySelection(); + } else if (ctrl && e.key === "v" && copiedBoxes) { + e.preventDefault(); e.stopPropagation(); + pasteBoxes(); + } else if (ctrl && e.key === "d" && node._activeIdx >= 0) { + e.preventDefault(); e.stopPropagation(); + copiedBoxes = copySelection(); + pasteBoxes(); + } + }); + + function onMove(e) { + if (!node._drawing) return; + const mN = mouseN(e); + const dN = { x: mN.x - node._dragStartN.x, y: mN.y - node._dragStartN.y }; + if (Math.abs(dN.x) + Math.abs(dN.y) > 0.001) node._pendingCollapse = -1; // it's a drag, not a click + if (node._dragMode === "move" && node._groupStart) { + let dx = dN.x, dy = dN.y; // clamp delta so the whole group stays in bounds + for (const i in node._groupStart) { + const s = node._groupStart[i]; + dx = Math.min(Math.max(dx, -s.x), 1 - s.w - s.x); + dy = Math.min(Math.max(dy, -s.y), 1 - s.h - s.y); + } + if (node.properties.snap) { // snap the group's movement to whole grid cells + const { sx, sy } = gridStep(); + dx = Math.round(dx / sx) * sx; + dy = Math.round(dy / sy) * sy; + } + for (const i in node._groupStart) { + const s = node._groupStart[i]; + node._boxes[i] = { ...s, x: s.x + dx, y: s.y + dy }; + delete node._boxes[i].nobbox; + } + drawCanvas(); updateBboxLabel(); + return; + } + if (node._groupStart && node._dragMode.startsWith("resize")) { + // Scale every selected box by the primary's resize, about the handle's fixed edge. + const a = node._boxAtStart, na = applyDrag(node._dragMode, a, dN); + const suf = node._dragMode.slice(7); // "tl"|"tr"|"bl"|"br"|"t"|"b"|"l"|"r" + const scaleX = (suf.includes("l") || suf.includes("r")) && a.w > 0 ? Math.max(0.02, na.w / a.w) : 1; + const scaleY = (suf.includes("t") || suf.includes("b")) && a.h > 0 ? Math.max(0.02, na.h / a.h) : 1; + const ax = suf.includes("l") ? a.x + a.w : a.x; // fixed (anchor) edges + const ay = suf.includes("t") ? a.y + a.h : a.y; + for (const i in node._groupStart) { + const s = node._groupStart[i]; + node._boxes[i] = normalizeBox({ + ...s, x: ax + (s.x - ax) * scaleX, y: ay + (s.y - ay) * scaleY, + w: s.w * scaleX, h: s.h * scaleY, + }); + delete node._boxes[i].nobbox; + } + drawCanvas(); updateBboxLabel(); + return; + } + const nb = snapBox(applyDrag(node._dragMode, node._boxAtStart, dN), node._dragMode); + delete nb.nobbox; // moving/resizing places the element (gives it a bbox) + node._boxes[node._activeIdx] = nb; + drawCanvas(); updateBboxLabel(); + } + function onUp() { + if (!node._drawing) return; + node._drawing = false; + canvasEl.removeEventListener("pointermove", onMove); + canvasEl.removeEventListener("pointerup", onUp); + canvasEl.removeEventListener("pointercancel", onUp); + // a click (no drag) on empty space drops the placeholder box and deselects everything + const b = node._boxes[node._activeIdx]; + if (b && (b.w < 0.005 || b.h < 0.005) && node._dragMode === "draw") { + node._boxes.splice(node._activeIdx, 1); + selectOnly(-1); + } else if (node._pendingCollapse >= 0) { // click (no drag) on a group member → keep only it + selectOnly(node._pendingCollapse); + } + node._pendingCollapse = -1; node._groupStart = null; + commit(); + } + + // ── marquee (rubber-band) selection: shift-drag ── + function marqueeRect() { + const m = node._marquee; + return { x: Math.min(m.x0, m.x), y: Math.min(m.y0, m.y), + w: Math.abs(m.x - m.x0), h: Math.abs(m.y - m.y0) }; + } + function rectsOverlap(r, b) { + return r.x < b.x + b.w && r.x + r.w > b.x && r.y < b.y + b.h && r.y + r.h > b.y; + } + function startMarquee(mN, startHit) { + node._marquee = { x0: mN.x, y0: mN.y, x: mN.x, y: mN.y }; + node._marqueeBase = new Set(node._selection); // additive: marquee unions with what's selected + node._marqueeStartHit = startHit; // for the shift-click (no drag) toggle fallback + node._marqueeActive = false; + canvasEl.focus(); + canvasEl.addEventListener("pointermove", onMarqueeMove); + canvasEl.addEventListener("pointerup", onMarqueeUp); + canvasEl.addEventListener("pointercancel", onMarqueeUp); + drawCanvas(); + } + function onMarqueeMove(e) { + if (!node._marquee) return; + const mN = mouseN(e); + node._marquee.x = mN.x; node._marquee.y = mN.y; + if (Math.abs(mN.x - node._marquee.x0) + Math.abs(mN.y - node._marquee.y0) > 0.01) node._marqueeActive = true; + if (node._marqueeActive) { + const r = marqueeRect(); + const sel = new Set(node._marqueeBase); + node._boxes.forEach((b, i) => { if (!b.locked && rectsOverlap(r, b)) sel.add(i); }); // marquee skips locked + node._selection = sel; + if (node._activeIdx < 0 || !sel.has(node._activeIdx)) node._activeIdx = sel.size ? [...sel][0] : node._activeIdx; + } + drawCanvas(); + } + function onMarqueeUp() { + canvasEl.removeEventListener("pointermove", onMarqueeMove); + canvasEl.removeEventListener("pointerup", onMarqueeUp); + canvasEl.removeEventListener("pointercancel", onMarqueeUp); + if (!node._marqueeActive && node._marqueeStartHit >= 0) { // shift-click on a box → toggle it + const idx = node._marqueeStartHit; + if (node._selection.has(idx) && node._selection.size > 1) { + node._selection.delete(idx); + if (node._activeIdx === idx) node._activeIdx = node._selection.values().next().value; + } else { + node._selection.add(idx); node._activeIdx = idx; + } + } + node._marquee = null; node._marqueeActive = false; + if (node._activeIdx >= 0 && !node._selection.has(node._activeIdx)) { + node._activeIdx = node._selection.size ? [...node._selection][0] : -1; + } + commit(); + } + + // ── duplicate placement: the new box follows the cursor until clicked ── + function placeFollower(mN) { + const b = node._boxes[node._activeIdx]; + if (!b) return; + b.x = clamp01(Math.min(mN.x - b.w / 2, 1 - b.w)); + b.y = clamp01(Math.min(mN.y - b.h / 2, 1 - b.h)); + delete b.nobbox; + drawCanvas(); updateBboxLabel(); + } + function startPlacing(srcIdx) { + const src = node._boxes[srcIdx]; + if (!src) return; + const nb = JSON.parse(JSON.stringify(src)); + delete nb.nobbox; delete nb.locked; // a duplicate is editable, not born locked + node._boxes.push(nb); + selectOnly(node._boxes.length - 1); + node._placing = true; + canvasEl.focus(); + canvasEl.style.cursor = "move"; + serialize(); renderPanel(); drawCanvas(); updateTokens(); + } + function finishPlacing() { + if (!node._placing) return; + node._placing = false; + canvasEl.style.cursor = "crosshair"; + commit(); fitNode(); + } + function cancelPlacing() { + if (!node._placing) return; + node._placing = false; + canvasEl.style.cursor = "crosshair"; + removeBox(node._activeIdx); + commit(); fitNode(); + } + + // ── right-click "layers" menu: list / select / delete / duplicate / reorder regions ── + function closeLayersMenu() { + if (node._layerMenu) { node._layerMenu.remove(); node._layerMenu = null; } + node._layerDismiss?.disarm(); node._layerDismiss = null; + } + function rowLabel(b) { + if (b.type === "text") { + const t = b.text ? `"${b.text}"` : ""; + return b.desc ? (t ? t + " — " + b.desc : b.desc) : t; + } + return b.desc || ""; + } + function openLayersMenu(clientX, clientY) { + closeLayersMenu(); + const menu = document.createElement("div"); + menu.className = "kjideo-menu"; + const hdr = document.createElement("div"); + hdr.className = "kjideo-mhdr"; + hdr.textContent = "Regions — top = front · click select · drag reorder"; + // 01 at the top of the list = front-most (drawn on top); see _draw / boxesAt. + menu.appendChild(hdr); + const list = document.createElement("div"); + menu.appendChild(list); + node._layerMenu = menu; + + const renumber = () => Array.from(list.querySelectorAll(".kjideo-lrow")).forEach((row, k) => { + row.querySelector(".kjideo-lnum").textContent = String(k + 1).padStart(2, "0"); + }); + function buildRows() { + list.innerHTML = ""; + if (!node._boxes.length) { + const empty = document.createElement("div"); + empty.className = "kjideo-mhdr"; empty.textContent = "No regions yet."; + list.appendChild(empty); + return; + } + node._boxes.forEach((b, i) => { + const row = document.createElement("div"); + row.className = "kjideo-lrow" + (i === node._activeIdx ? " active" : ""); + row._box = b; + const sw = document.createElement("div"); + sw.className = "kjideo-lsw"; + sw.style.background = (b.palette || []).find(Boolean) || "#8c8c8c"; + const num = document.createElement("span"); + num.className = "kjideo-lnum"; num.textContent = String(i + 1).padStart(2, "0"); + const txt = document.createElement("span"); + const label = rowLabel(b); + txt.className = "kjideo-ltext" + (label ? "" : " empty"); + txt.textContent = label || (b.type === "text" ? "(text)" : "(empty)"); + txt.title = label; + const lock = document.createElement("button"); + lock.className = "kjideo-lbtn kjideo-lock" + (b.locked ? " on" : ""); + lock.textContent = b.locked ? "🔒" : "🔓"; + lock.title = b.locked ? "Unlock (allow moving/resizing)" : "Lock (freeze on the canvas)"; + const dup = document.createElement("button"); + dup.className = "kjideo-lbtn"; dup.textContent = "⧉"; + dup.title = "Duplicate, then click on the canvas to place"; + const del = document.createElement("button"); + del.className = "kjideo-lbtn del"; del.textContent = "✕"; + del.title = b.locked ? "Unlock to delete" : "Delete region"; + del.disabled = !!b.locked; + row.append(sw, num, txt, lock, dup, del); + list.appendChild(row); + + lock.addEventListener("click", (e) => { + e.stopPropagation(); + b.locked = !b.locked; + if (b.locked) { // a locked box drops out of the selection + const idx = node._boxes.indexOf(b); + node._selection.delete(idx); + if (node._activeIdx === idx) node._activeIdx = node._selection.size ? [...node._selection][0] : -1; + } + commit(); buildRows(); + }); + + row.addEventListener("click", () => { + if (row._dragged) { row._dragged = false; return; } + selectOnly(node._boxes.indexOf(b)); + commit(); + for (const r of list.querySelectorAll(".kjideo-lrow")) r.classList.toggle("active", r._box === b); + }); + dup.addEventListener("click", (e) => { + e.stopPropagation(); + const idx = node._boxes.indexOf(b); + closeLayersMenu(); + startPlacing(idx); + }); + del.addEventListener("click", (e) => { + e.stopPropagation(); + const idx = node._boxes.indexOf(b); + if (idx < 0) return; + removeBox(idx); commit(); fitNode(); + if (!node._boxes.length) { closeLayersMenu(); return; } + buildRows(); + }); + // drag-reorder (vertical FLIP, mirrors the palette swatch reorder) + row.addEventListener("pointerdown", (e) => { + if (e.button !== 0 || e.target === lock || e.target === dup || e.target === del) return; + e.preventDefault(); e.stopPropagation(); + const sx = e.clientX, sy = e.clientY; + let dragging = false; + const move = (me) => { + if (!dragging) { + if (Math.abs(me.clientX - sx) + Math.abs(me.clientY - sy) < 4) return; + dragging = true; row.classList.add("dragging"); document.body.classList.add("kjideo-dragging"); + } + for (const other of list.querySelectorAll(".kjideo-lrow")) { + if (other === row) continue; + const r = other.getBoundingClientRect(); + if (me.clientY >= r.top && me.clientY <= r.bottom) { + const ref = me.clientY > r.top + r.height / 2 ? other.nextSibling : other; + if (ref === row || ref === row.nextSibling) break; + const els = Array.from(list.querySelectorAll(".kjideo-lrow")); + const prev = els.map((el) => el.getBoundingClientRect().top); + list.insertBefore(row, ref); + els.forEach((el, k) => { // FLIP: slide to new positions + const dy = prev[k] - el.getBoundingClientRect().top; + if (!dy) return; + el.style.transition = "none"; + el.style.transform = `translateY(${dy}px)`; + el.getBoundingClientRect(); // flush + el.style.transition = ""; el.style.transform = ""; + }); + break; + } + } + }; + const up = () => { + document.removeEventListener("pointermove", move); + document.removeEventListener("pointerup", up); + document.removeEventListener("pointercancel", up); + document.body.classList.remove("kjideo-dragging"); + if (dragging) { + row.classList.remove("dragging"); + row._dragged = true; // suppress the trailing click + const active = node._boxes[node._activeIdx]; + const order = Array.from(list.querySelectorAll(".kjideo-lrow")).map((el) => el._box); + if (order.length === node._boxes.length) node._boxes = order; + selectOnly(active ? node._boxes.indexOf(active) : -1); // reorder invalidates multi-select indices + renumber(); + commit(); + } + }; + document.addEventListener("pointermove", move); + document.addEventListener("pointerup", up); + document.addEventListener("pointercancel", up); + }); + }); + } + buildRows(); + + document.body.appendChild(menu); + const r = menu.getBoundingClientRect(); // clamp into the viewport + let left = clientX, top = clientY; + if (left + r.width > window.innerWidth) left = window.innerWidth - r.width - 4; + if (top + r.height > window.innerHeight) top = window.innerHeight - r.height - 4; + menu.style.left = Math.max(4, left) + "px"; + menu.style.top = Math.max(4, top) + "px"; + + node._layerDismiss = outsideDismiss(menu, () => closeLayersMenu()); + node._layerDismiss.arm(); + } + + canvasEl.addEventListener("contextmenu", (e) => { + e.preventDefault(); e.stopPropagation(); + if (node._placing || node._hideBoxes) return; + closeInlineEditor(); + openLayersMenu(e.clientX, e.clientY); + }); + stopProp(clearBtn); + clearBtn.addEventListener("click", () => { + closeInlineEditor(); + node._boxes = []; node._activeIdx = -1; node._selection = new Set(); node._stylePalette = []; + node._lastImported = ""; + commit(); rebuildStylePalette(); fitNode(); // a wired import re-seeds on the next run (serializeValue cache-busts) + }); + + // ── build caption JSON (mirrors Python key order) ── + // pyJson: matches Python _dumps — indent=4, but scalar arrays stay on one line. + function pyJson(v, lvl = 0) { + if (v === null) return "null"; + if (typeof v === "number" || typeof v === "boolean") return String(v); + if (typeof v === "string") return JSON.stringify(v); + const pad = " ".repeat(lvl + 1), end = " ".repeat(lvl); + if (Array.isArray(v)) { + if (!v.length) return "[]"; + if (v.every((x) => x === null || typeof x !== "object")) // scalar array → inline + return "[" + v.map((x) => pyJson(x, lvl)).join(", ") + "]"; + return "[\n" + v.map((x) => pad + pyJson(x, lvl + 1)).join(",\n") + "\n" + end + "]"; + } + const keys = Object.keys(v); + if (!keys.length) return "{}"; + return "{\n" + keys.map((k) => pad + JSON.stringify(k) + ": " + pyJson(v[k], lvl + 1)).join(",\n") + "\n" + end + "}"; + } + function getW(name) { const w = findW(name); return w ? w.value : ""; } + function cleanPalette(arr) { return (arr || []).filter((c) => c).map((c) => c.toUpperCase()); } + function normBboxJS(b) { + const c = (v) => Math.max(0, Math.min(1000, Math.round(v * 1000))); + let ymin = c(b.y), xmin = c(b.x), ymax = c(b.y + b.h), xmax = c(b.x + b.w); + if (ymin > ymax) [ymin, ymax] = [ymax, ymin]; + if (xmin > xmax) [xmin, xmax] = [xmax, xmin]; + return [ymin, xmin, ymax, xmax]; + } + function buildCaption() { + const cap = {}; + if ((getW("high_level_description") || "").trim()) cap.high_level_description = getW("high_level_description"); + const styleW = findW("style"); + const kind = styleW ? styleW.value : "none"; + if (kind !== "none") { + const sd = { aesthetics: getW("aesthetics"), lighting: getW("lighting") }; + if (kind === "photo") { sd.photo = getW("style.photo") || ""; sd.medium = getW("medium"); } + else { sd.medium = getW("medium"); sd.art_style = getW("style.art_style") || ""; } + const pal = cleanPalette(node._stylePalette); + if (pal.length) sd.color_palette = pal; + cap.style_description = sd; + } + const elements = node._boxes.map((b) => { + const etype = b.type === "text" ? "text" : "obj"; + const el = { type: etype }; + if (!b.nobbox) el.bbox = normBboxJS(b); // unplaced elements omit bbox + if (etype === "text") el.text = b.text || ""; + el.desc = b.desc || ""; + const pal = cleanPalette(b.palette).slice(0, MAX_ELEM_COLORS); + if (pal.length) el.color_palette = pal; + return el; + }); + cap.compositional_deconstruction = { background: getW("background"), elements }; + return (formatWidget?.value) === "pretty" ? pyJson(cap) : JSON.stringify(cap); // compact by default; matches the node output + } + // Rough token estimate (~chars/4); exact count needs the Qwen tokenizer. + function updateTokens() { + const n = Math.ceil(buildCaption().length / 4); + tokenSpan.textContent = "~" + n + " tok"; + // grey <256 (sparse) · green healthy · orange nearing · red ≥2048 (model hard cap) + tokenSpan.style.color = n >= 2048 ? "#e05555" : n >= 1792 ? "#e6a23c" : n >= 256 ? "#6cc06c" : "#888"; + } + async function doCopy() { + const txt = buildCaption(); + try { await navigator.clipboard.writeText(txt); copyBtn.textContent = "Copied"; setTimeout(() => (copyBtn.textContent = "Copy"), 900); } + catch (e) { window.prompt("Copy the caption JSON:", txt); } + } + stopProp(copyBtn); + copyBtn.addEventListener("click", doCopy); + + // ── import a caption JSON and populate the node ── + function setWidgetVal(name, val) { + const w = findW(name); + if (w) { w.value = val; w.callback?.(val); } + } + function bboxElemToBox(el, idx) { + if (!el || typeof el !== "object") return null; + const box = { type: el.type === "text" ? "text" : "obj", + text: el.text || "", desc: el.desc || "", + palette: Array.isArray(el.color_palette) ? el.color_palette.slice() : [] }; + const bb = el.bbox; + if (Array.isArray(bb) && bb.length === 4) { + const [ymin, xmin, ymax, xmax] = bb; + box.x = xmin / 1000; box.y = ymin / 1000; box.w = (xmax - xmin) / 1000; box.h = (ymax - ymin) / 1000; + } else { + // No bbox: "unplaced" element — small placeholder, flagged so export omits bbox. + const k = (idx || 0) % 6; + box.x = 0.03 + k * 0.035; box.y = 0.03 + k * 0.035; box.w = 0.22; box.h = 0.14; + box.nobbox = true; + } + return box; + } + function applyCaption(cap) { + const cd = (cap && cap.compositional_deconstruction) || {}; + const els = Array.isArray(cd.elements) ? cd.elements : []; + node._boxes = els.map((el, i) => bboxElemToBox(el, i)).filter(Boolean); + selectOnly(node._boxes.length ? 0 : -1); + setWidgetVal("high_level_description", cap.high_level_description || ""); + setWidgetVal("background", cd.background || ""); + const sd = cap.style_description || {}; + let kind = "none"; + if (typeof sd.photo === "string") kind = "photo"; + else if (typeof sd.art_style === "string") kind = "art_style"; + const styleW = findW("style"); + if (styleW) styleW.value = kind; // setter synchronously rebuilds sub-widgets + if (kind === "photo") setWidgetVal("style.photo", sd.photo || ""); + else if (kind === "art_style") setWidgetVal("style.art_style", sd.art_style || ""); + setWidgetVal("aesthetics", sd.aesthetics || ""); + setWidgetVal("lighting", sd.lighting || ""); + setWidgetVal("medium", sd.medium || ""); + node._stylePalette = Array.isArray(sd.color_palette) ? sd.color_palette.slice() : []; + } + function tryParseCaption(t) { + if (!t) return null; + try { const o = JSON.parse(t); return (o && typeof o === "object" && o.compositional_deconstruction) ? o : null; } + catch (e) { return null; } + } + // Apply a parsed caption to the editor and refresh everything. + function loadCaption(cap) { + closeInlineEditor(); + applyCaption(cap); + syncCanvasToDims(); commit(); rebuildStylePalette(); fitNode(); + } + // Append a caption's regions to the current canvas (keeps existing boxes + caption fields). + function insertCaptionBoxes(cap) { + closeInlineEditor(); + const cd = (cap && cap.compositional_deconstruction) || {}; + const els = Array.isArray(cd.elements) ? cd.elements : []; + const added = els.map((el, i) => bboxElemToBox(el, i)).filter(Boolean); + if (!added.length) return; + const start = node._boxes.length; + node._boxes.push(...added); + node._selection = new Set(); // select the inserted regions + for (let i = start; i < node._boxes.length; i++) node._selection.add(i); + node._activeIdx = node._boxes.length - 1; + commit(); fitNode(); + } + async function doImport() { + let cap = null, txt = ""; + try { txt = (await navigator.clipboard.readText() || "").trim(); cap = tryParseCaption(txt); } catch (e) {} + if (!cap) { txt = (window.prompt("Paste Ideogram 4 caption JSON:", "") || "").trim(); cap = tryParseCaption(txt); } + if (!cap) { if (txt) alert("Not a valid Ideogram 4 caption JSON (needs 'compositional_deconstruction')."); return; } + loadCaption(cap); + } + stopProp(importBtn); + importBtn.addEventListener("click", doImport); + + // Populate the editor from a caption pushed back by execute() when import_json + // is connected (a connected socket can't be read in the frontend directly). + function applyImported(capStr) { + if (!capStr) return; + // "always" mode re-applies even an unchanged import so the editor snaps back to the + // authoritative JSON after edits; "when empty" keeps the guard so edits stick. + const always = findW("import_mode")?.value === "always"; + if (capStr === node._lastImported && !always) return; + const cap = tryParseCaption(capStr); + if (!cap) return; + node._lastImported = capStr; + loadCaption(cap); + } + chainCallback(node, "onExecuted", function (message) { + if (message?.caption) applyImported(message.caption[0]); + // Seed regions from the bboxes input only when the editor is empty, so user-drawn/edited + // regions are never overwritten. + if (message?.boxes && !node._boxes.length) { + const seeded = JSON.parse(message.boxes[0]); + if (Array.isArray(seeded) && seeded.length) { + node._boxes = seeded.filter((b) => b && typeof b.x === "number" && typeof b.w === "number"); + selectOnly(node._boxes.length ? 0 : -1); + commit(); fitNode(); + } + } + // Reflect resolved width/height (e.g. from connected inputs) in the canvas aspect. + // A connected background image governs the aspect itself, so skip then. + if (message?.dims && !node._bgImg) { + const [w, h] = message.dims; + if (wWidget && w) wWidget.value = w; + if (hWidget && h) hWidget.value = h; + syncCanvasToDims(); fitNode(); + } + }); + + // ── property panel ── + function stopProp(el) { + for (const ev of ["mousedown", "pointerdown", "wheel"]) el.addEventListener(ev, (e) => e.stopPropagation()); + } + // Color swatches: onEdit on change, onStruct on add/remove/reorder. Shared by both palettes. + // Pointer-based drag (HTML5 DnD is unreliable inside LiteGraph DOM widgets) with live reorder. + function buildSwatchRow(container, arr, max, onEdit, onStruct) { + arr.forEach((hex, i) => { + const sw = document.createElement("div"); + sw.className = "kjideo-sw"; + sw.style.background = hex; + sw.dataset.hex = hex; + sw.title = "Click edit · drag reorder · Ctrl+C/V copy/paste hex · right-click remove"; + const inp = document.createElement("input"); + inp.type = "color"; inp.value = hex; + sw.appendChild(inp); + container.appendChild(sw); + const setColor = (hex2) => { arr[i] = hex2; inp.value = hex2; sw.style.background = hex2; sw.dataset.hex = hex2; onEdit(); }; + inp.addEventListener("input", () => setColor(inp.value)); + sw.addEventListener("pointerenter", () => { hoveredSwatch = { sw, setColor }; }); + sw.addEventListener("pointerleave", () => { if (hoveredSwatch && hoveredSwatch.sw === sw) hoveredSwatch = null; }); + sw.addEventListener("wheel", (e) => e.stopPropagation()); + sw.addEventListener("contextmenu", (e) => { e.preventDefault(); e.stopPropagation(); arr.splice(i, 1); onStruct(); }); + sw.addEventListener("pointerdown", (e) => { + if (e.button !== 0) return; + e.preventDefault(); e.stopPropagation(); + try { sw.setPointerCapture(e.pointerId); } catch (e2) {} // capture so Nodes 2.0's WidgetDOM .stop can't swallow the drag + const sx = e.clientX, sy = e.clientY; + let dragging = false; + const move = (me) => { + if (!dragging) { + if (Math.abs(me.clientX - sx) + Math.abs(me.clientY - sy) < 4) return; + dragging = true; sw.classList.add("dragging"); document.body.classList.add("kjideo-dragging"); + } + for (const other of container.querySelectorAll(".kjideo-sw")) { + if (other === sw) continue; + const r = other.getBoundingClientRect(); + if (me.clientX >= r.left && me.clientX <= r.right && me.clientY >= r.top - 6 && me.clientY <= r.bottom + 6) { + const ref = me.clientX > r.left + r.width / 2 ? other.nextSibling : other; + if (ref === sw || ref === sw.nextSibling) break; // already there + const els = Array.from(container.querySelectorAll(".kjideo-sw")); + const prev = els.map((el) => el.getBoundingClientRect().left); + container.insertBefore(sw, ref); + els.forEach((el, k) => { // FLIP: slide to new positions + const dx = prev[k] - el.getBoundingClientRect().left; + if (!dx) return; + el.style.transition = "none"; + el.style.transform = `translateX(${dx}px)`; + el.getBoundingClientRect(); // flush + el.style.transition = ""; el.style.transform = ""; + }); + break; + } + } + }; + const up = () => { + sw.removeEventListener("pointermove", move); + sw.removeEventListener("pointerup", up); + sw.removeEventListener("pointercancel", up); + document.body.classList.remove("kjideo-dragging"); + if (dragging) { + sw.classList.remove("dragging"); + const order = Array.from(container.querySelectorAll(".kjideo-sw")).map((el) => el.dataset.hex); + if (order.length === arr.length) { arr.length = 0; arr.push(...order); } + onStruct(); + } else { + inp.click(); // no drag → treat as click, open the picker + } + }; + sw.addEventListener("pointermove", move); + sw.addEventListener("pointerup", up); + sw.addEventListener("pointercancel", up); + }); + }); + if (arr.length < max) { + const add = document.createElement("button"); + add.className = "kjideo-btn"; add.textContent = "+"; + add.title = "Add a color (uses the clipboard color if it is one)"; + stopProp(add); + add.addEventListener("click", async () => { + let col = "#ffffff"; + try { const c = parseColorString(await navigator.clipboard.readText()); if (c) col = c; } catch (e) {} + arr.push(col); onStruct(); + }); + container.appendChild(add); + } + } + + // Swatch color changed (no add/remove): persist + repaint. + function swatchEdit() { serialize(); drawCanvas(); } + + function rebuildStylePalette() { + while (styleBar.children.length > 1) styleBar.removeChild(styleBar.lastChild); + buildSwatchRow(styleBar, node._stylePalette, MAX_STYLE_COLORS, + swatchEdit, + () => { swatchEdit(); rebuildStylePalette(); fitNode(); }); + } + + // Textarea that flexes to fill the prompt panel (whose height is set by the splitter). + function makeArea(field, value, placeholder, onInput) { + const ta = document.createElement("textarea"); + ta.className = "kjideo-area"; + ta.placeholder = placeholder; + ta.value = value || ""; + stopProp(ta); + ta.addEventListener("input", onInput); + return ta; + } + let bboxPx = null, bboxGrid = null; // editable bbox fields: pixels and 0–1000 grid + function dims() { return [wWidget ? wWidget.value : 1024, hWidget ? hWidget.value : 1024]; } + function boxToPx(b) { // same order as the grid: ymin, xmin, ymax, xmax + const [W, H] = dims(); + return [Math.round(b.y * H), Math.round(b.x * W), Math.round((b.y + b.h) * H), Math.round((b.x + b.w) * W)]; + } + function setField(inp, b, fn) { // skip while the field is being edited + if (!inp || document.activeElement === inp) return; + inp.value = (!b || b.nobbox) ? "" : fn(b).join(", "); + } + function updateBboxLabel() { + const b = node._boxes[node._activeIdx]; + setField(bboxPx, b, boxToPx); + setField(bboxGrid, b, normBboxJS); + } + function parse4(inp) { + const nums = inp.value.split(/[,\s]+/).map(Number).filter((n) => !isNaN(n)); + return nums.length === 4 ? nums : null; + } + function commitPxEdit() { + const b = node._boxes[node._activeIdx]; if (!b || !bboxPx) return; + const nums = parse4(bboxPx); if (!nums) { updateBboxLabel(); return; } + const [W, H] = dims(); + let [ymin, xmin, ymax, xmax] = nums; // ymin, xmin, ymax, xmax (matches grid) + ymin = Math.max(0, Math.min(H, ymin)); ymax = Math.max(0, Math.min(H, ymax)); + xmin = Math.max(0, Math.min(W, xmin)); xmax = Math.max(0, Math.min(W, xmax)); + if (ymin > ymax) [ymin, ymax] = [ymax, ymin]; + if (xmin > xmax) [xmin, xmax] = [xmax, xmin]; + b.y = ymin / H; b.x = xmin / W; b.h = (ymax - ymin) / H; b.w = (xmax - xmin) / W; + delete b.nobbox; commit(); fitNode(); + } + function commitGridEdit() { + const b = node._boxes[node._activeIdx]; if (!b || !bboxGrid) return; + const nums = parse4(bboxGrid); if (!nums) { updateBboxLabel(); return; } + let [ymin, xmin, ymax, xmax] = nums.map((n) => Math.max(0, Math.min(1000, n))); + if (ymin > ymax) [ymin, ymax] = [ymax, ymin]; + if (xmin > xmax) [xmin, xmax] = [xmax, xmin]; + b.y = ymin / 1000; b.x = xmin / 1000; b.h = (ymax - ymin) / 1000; b.w = (xmax - xmin) / 1000; + delete b.nobbox; commit(); fitNode(); + } + function makeBboxField(placeholder, title, onCommit) { + const inp = document.createElement("input"); + inp.type = "text"; inp.className = "kjideo-bbox"; + inp.placeholder = placeholder; inp.title = title; + stopProp(inp); + inp.addEventListener("keydown", (e) => { + e.stopPropagation(); + if (e.key === "Enter") inp.blur(); + else if (e.key === "Escape") { updateBboxLabel(); inp.blur(); } + }); + inp.addEventListener("change", onCommit); + return inp; + } + function renderPanel() { + for (const ro of node._areaObservers) ro.disconnect(); + node._areaObservers = []; + panel.innerHTML = ""; + const b = node._boxes[node._activeIdx]; + if (!b) { + hint.textContent = ""; + const p = document.createElement("div"); + p.style.color = "#888"; + p.textContent = node._boxes.length ? "Click a region to edit it." : "No regions yet."; + panel.appendChild(p); + return; + } + const col = (b.palette || []).find(Boolean) || "#bbb"; + const selN = node._selection.size; + // Build with DOM + style.color (a CSS value) — never innerHTML — since col comes from + // box data that may be loaded from an untrusted template/import (avoids HTML injection). + hint.textContent = ""; + const tag = document.createElement("b"); + tag.style.color = col; tag.textContent = "region " + (node._activeIdx + 1); + hint.appendChild(tag); + if (selN > 1) { + const s = document.createElement("span"); + s.style.color = "#888"; s.textContent = ` (${selN} selected)`; + hint.appendChild(s); + } + + // type toggle + const typeRow = document.createElement("div"); + typeRow.className = "kjideo-row"; + const lbl = document.createElement("span"); lbl.textContent = "type:"; typeRow.appendChild(lbl); + for (const t of ["obj", "text"]) { + const btn = document.createElement("button"); + btn.className = "kjideo-btn" + (b.type === t ? " active" : ""); + btn.textContent = t; + stopProp(btn); + btn.addEventListener("click", () => { b.type = t; commit(); }); + typeRow.appendChild(btn); + } + const pxLbl = document.createElement("span"); + pxLbl.textContent = "px:"; pxLbl.style.cssText = "margin-left:auto; color:#888;"; + typeRow.appendChild(pxLbl); + bboxPx = makeBboxField("ymin, xmin, ymax, xmax", "Pixel bbox (of the node's width/height): ymin, xmin, ymax, xmax — editable", commitPxEdit); + typeRow.appendChild(bboxPx); + const gl = document.createElement("span"); + gl.textContent = "out:"; gl.style.color = "#888"; + typeRow.appendChild(gl); + bboxGrid = makeBboxField("ymin, xmin, ymax, xmax", "Exported bbox on the 0–1000 grid: ymin, xmin, ymax, xmax — editable", commitGridEdit); + typeRow.appendChild(bboxGrid); + updateBboxLabel(); + panel.appendChild(typeRow); + + // text (only for text type) + if (b.type === "text") { + panel.appendChild(makeArea("text", b.text, "text to render (verbatim)", + function () { b.text = this.value; touch(); })); + } + + panel.appendChild(makeArea("desc", b.desc, "description of this region", + function () { b.desc = this.value; touch(); })); + + // palette + const palRow = document.createElement("div"); + palRow.className = "kjideo-row"; + const pl = document.createElement("span"); pl.textContent = "colors:"; palRow.appendChild(pl); + b.palette = b.palette || []; + buildSwatchRow(palRow, b.palette, MAX_ELEM_COLORS, swatchEdit, commit); + panel.appendChild(palRow); + } + + // ── width/height widget callbacks ── + for (const w of [wWidget, hWidget]) { + if (!w) continue; + chainCallback(w, "callback", () => { syncCanvasToDims(); drawCanvas(); fitNode(); }); + } + // Update the token estimate when the caption-level text widgets change. + for (const name of ["background", "high_level_description", "aesthetics", "lighting", "medium", "style"]) { + const w = findW(name); + if (w) chainCallback(w, "callback", () => updateTokens()); + } + + // ── node resizes freely; the canvas just re-fits the space the widget is given ── + let _resizing = false; + chainCallback(node, "onResize", function () { + if (_resizing || node._detached || wrap.offsetParent === null) return; + _resizing = true; + fitCanvas(); + requestAnimationFrame(fitCanvas); // re-fit after LiteGraph re-arranges widgets on the next draw + _resizing = false; + }); + + // Optional reference image as the canvas background (matches ImageTransformKJ). + function loadBg(src) { + const img = new Image(); + img.crossOrigin = "anonymous"; + img.onload = () => { + node._bgImg = img; + const r16 = (v) => Math.max(16, Math.round(v / 16) * 16); // model needs multiples of 16 + if (wWidget) wWidget.value = r16(img.naturalWidth); // match canvas aspect to the image + if (hWidget) hWidget.value = r16(img.naturalHeight); + syncCanvasToDims(); drawCanvas(); fitNode(); updateGrabBtn(); + }; + img.src = src; + } + watchImageInputs(node, "image", (sources) => { + if (!sources.length) { if (!node._bgManual) { node._bgImg = null; drawCanvas(); updateGrabBtn(); } return; } + node._bgManual = false; // a connected image takes over + const s = sources[0]; + if (s.isVideo && s.videoEl) captureVideoFrame(s.videoEl, (cv) => loadBg(cv.toDataURL("image/webp", 0.9))); + else if (s.url) loadBg(s.url); + }); + // "Grab BG" button: use the last generated image as the background, or clear it. + node._grabResultBg = () => { + if (!lastResultImage) { alert("No sampling result yet — run a generation first."); return; } + node._bgManual = true; + loadBg(resultViewUrl(lastResultImage)); + }; + node._clearBg = () => { + node._bgManual = false; node._bgImg = null; + if (node._liveBmp?.close) { try { node._liveBmp.close(); } catch (e) {} node._liveBmp = null; } + drawCanvas(); updateGrabBtn(); + }; + // Feed a live sampling-preview frame as the background (no width/height change). + node._ideoSetLiveBg = (bmp) => { + if (node._liveBmp?.close && node._liveBmp !== bmp) { try { node._liveBmp.close(); } catch (e) {} } + node._liveBmp = bmp; node._bgImg = bmp; node._bgManual = true; + drawCanvas(); updateGrabBtn(); + }; + // After generation, replace the live preview with the full-res final result. + node._ideoGrabFinal = () => { + if (!lastResultImage) return; + if (node._liveBmp?.close) { try { node._liveBmp.close(); } catch (e) {} node._liveBmp = null; } + node._bgManual = true; + loadBg(resultViewUrl(lastResultImage)); + }; + + // Active-box highlight only while the editor is focused or the node is selected. + wrap.addEventListener("focusin", () => { if (!node._focused) { node._focused = true; drawCanvas(); } }); + wrap.addEventListener("focusout", (e) => { + if (!wrap.contains(e.relatedTarget)) { node._focused = false; drawCanvas(); } + }); + chainCallback(node, "onSelected", function () { node._selected = true; drawCanvas(); }); + chainCallback(node, "onDeselected", function () { node._selected = false; finishPlacing(); closeLayersMenu(); drawCanvas(); }); + + chainCallback(node, "onRemoved", function () { + livePreviewNodes.delete(node); + if (hoveredCanvasNode === node) hoveredCanvasNode = null; + if (node._fullscreen) { // tear down the popup if open + document.removeEventListener("keydown", onFsEsc, true); + window.removeEventListener("resize", fitFsCanvas); + node._fsOverlay?.remove(); + } + pinnedDocks.delete(node); liveDocks.delete(node); + node._dockRO?.disconnect(); node._dockEl?.remove(); // tear down the floating dock if open + node._visObserver?.disconnect(); + if (node._liveBmp?.close) { try { node._liveBmp.close(); } catch (e) {} node._liveBmp = null; } // release GPU bitmap + closeBgMenu(); node._bgMenu?.remove(); + closeTxtMenu(); node._txtMenu?.remove(); + closeTplMenu(); node._tplMenu?.remove(); + closeInlineEditor(); + closeLayersMenu(); + for (const ro of node._areaObservers) ro.disconnect(); + node._areaObservers = []; + }); + + // ── restore on load ── + function _parseBoxes(s) { + try { + const p = JSON.parse(s); + if (Array.isArray(p) && p.some((b) => b && typeof b.x === "number" && typeof b.w === "number")) return p; + } catch (e) {} + return null; + } + // Persist editor data by name (robust to widget-order changes across versions). + chainCallback(node, "onSerialize", function (o) { + if (!o) return; + const p = node.properties; + o.ideo = { boxes: node._boxes, palette: node._stylePalette, importMode: findW("import_mode")?.value, + outputFormat: findW("output_format")?.value, + dock: { pinned: p.dockPinned, graph: p.dockGraph, rect: p.dockRect, panelH: node._panelH, min: p.dockMin } }; + }); + chainCallback(node, "onConfigure", function (o) { + const raw = o && Array.isArray(o.widgets_values) ? o.widgets_values : []; + // Restore dock geometry from the blob into node.properties BEFORE ensureDocked runs. + node._savedSize = node.size ? [node.size[0], node.size[1]] : null; // pre-collapse size (old aspect-locked height) + const d = o && o.ideo && o.ideo.dock; + if (d) { + node._dockGeomRestored = true; // honor it; old workflows w/o the blob get a default + if (d.graph) node.properties.dockGraph = d.graph; + if (d.rect) node.properties.dockRect = d.rect; + if (d.pinned != null) node.properties.dockPinned = d.pinned; + if (d.min != null) node.properties.dockMin = d.min; + if (typeof d.panelH === "number") { node._panelH = d.panelH; panel.style.height = d.panelH + "px"; } + } + // Recover regions: name-keyed blob → named widget → raw saved values (survives + // any widget reorder/remap across versions) → live widgets. + let boxes = (o && o.ideo && Array.isArray(o.ideo.boxes)) ? o.ideo.boxes : _parseBoxes(elementsWidget?.value || ""); + if (!boxes) { for (const v of raw) { const b = _parseBoxes(v); if (b) { boxes = b; break; } } } + if (!boxes) { for (const w of node.widgets || []) { const b = _parseBoxes(w?.value); if (b) { boxes = b; break; } } } + if (boxes) { + node._boxes = boxes.filter((b) => b && typeof b.x === "number"); + selectOnly(node._boxes.length ? 0 : -1); + } + const isPal = (p) => Array.isArray(p) && p.length && p.every((c) => typeof c === "string" && c[0] === "#"); + let pal = (o && o.ideo && isPal(o.ideo.palette)) ? o.ideo.palette : null; + if (!pal) { try { const p = JSON.parse(stylePaletteWidget?.value || ""); if (isPal(p)) pal = p; } catch (e) {} } + if (!pal) { for (const v of raw) { try { const p = JSON.parse(v); if (isPal(p)) { pal = p; break; } } catch (e) {} } } + if (pal) node._stylePalette = pal.slice(); + const imW = findW("import_mode"); // restore import_mode; coerce to a valid option so + if (imW) { // old workflows (saved before this widget) don't fail Combo validation + const im = o && o.ideo && o.ideo.importMode, opts = ["when empty", "always"]; + imW.value = opts.includes(im) ? im : (opts.includes(imW.value) ? imW.value : "when empty"); + } + if (formatWidget) { // restore output format from the blob; default compact unless explicitly pretty + formatWidget.value = (o && o.ideo && o.ideo.outputFormat) === "pretty" ? "pretty" : "compact"; + compactLbl._cb.checked = formatWidget.value === "compact"; + } + node._configured = true; // mark as loaded so initial layout keeps the restored size + hideDataWidgets(); + serialize(); // realign widget values for Python + future saves + if (bgBrightnessWidget) { + if (typeof bgBrightnessWidget.value !== "number") bgBrightnessWidget.value = 25; // old workflows may restore "" + bgSlider.value = bgBrightnessWidget.value; + } + // node.properties is restored after onNodeCreated, so resync the toolbar controls to it. + liveChk.checked = !!node.properties.liveBg; + if (liveChk.checked) livePreviewNodes.add(node); else livePreviewNodes.delete(node); + guideSel.value = node.properties.guide || "none"; + gridSlider.value = GRID_INV - (node.properties.gridSize || 10); + snapChk.checked = !!node.properties.snap; + guideColor.value = node.properties.guideColor || "#ffffff"; + opacitySlider.value = node.properties.guideOpacity == null ? 100 : node.properties.guideOpacity; + showLbl._cb.checked = node.properties.showBoxText !== false; + strokeLbl._cb.checked = node.properties.textStroke !== false; + autoLbl._cb.checked = node.properties.textAutoPlace !== false; + sizeSlider.value = node.properties.textSize || 12; + boxOpacSlider.value = node.properties.boxOpacity == null ? 14 : node.properties.boxOpacity; + syncCanvasToDims(); + rebuildStylePalette(); + renderPanel(); + drawCanvas(); + updateTokens(); + requestAnimationFrame(fitNode); + requestAnimationFrame(() => ensureDocked(false)); // reload: honor the saved dock geometry, never match node width + }); + + // initial layout (deferred so size/last_y are settled) + setTimeout(() => { + hideDataWidgets(); + if (!node._configured) node.setSize([Math.max(480, node.size[0]), node.size[1]]); // fresh node: a comfortable default width + else if (node.size[0] < 380) node.setSize([380, node.size[1]]); // loaded node: just enforce the min + syncCanvasToDims(); + rebuildStylePalette(); + renderPanel(); + drawCanvas(); + updateTokens(); + if (!node._configured) ensureDocked(true); // fresh node: pop the dock under it, matching its width + fitNode(); + }, 0); + }); + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/image_transform.js b/custom_nodes/ComfyUI-KJNodes/web/js/image_transform.js new file mode 100644 index 0000000000000000000000000000000000000000..0555a30378b774cbd68b8a0d64a7c0c2653e6f7a --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/image_transform.js @@ -0,0 +1,2025 @@ +import { chainCallback, addMiddleClickPan, addWheelPassthrough, captureVideoFrame, watchImageInputs, rectHitTest, cursorForBboxMode } from './utility.js'; +const { app } = window.comfyAPI.app; + +const BBOX_PALETTE = ["#46b4e6", "#e68246", "#82e646", "#e646b4", "#e6e646", "#46e6c8"]; +const MAX_CANVAS_W = 1024, MAX_CANVAS_H = 768; +const BBOX_COLORS = BBOX_PALETTE.map((hex) => { + const r = parseInt(hex.slice(1, 3), 16), g = parseInt(hex.slice(3, 5), 16), b = parseInt(hex.slice(5, 7), 16); + return { rgb: `rgb(${r}, ${g}, ${b})`, gridActive: `rgba(${r}, ${g}, ${b}, 0.35)`, + gridInactive: `rgba(${r}, ${g}, ${b}, 0.3)`, tint: `rgba(${r}, ${g}, ${b}, 0.2)` }; +}); + +function getBboxColor(index, active) { + const c = BBOX_COLORS[index % BBOX_COLORS.length]; + return { border: c.rgb, fill: c.rgb, handle: c.rgb, gridColor: active ? c.gridActive : c.gridInactive, tint: c.tint }; +} + +function clampToMaxCanvas(w, h) { + if (w > MAX_CANVAS_W || h > MAX_CANVAS_H) { + const s = Math.min(MAX_CANVAS_W / w, MAX_CANVAS_H / h); + return [Math.round(w * s), Math.round(h * s)]; + } + return [w, h]; +} + +function roundDown(val, divBy) { + return divBy > 1 ? val - (val % divBy) : val; +} + +app.registerExtension({ + name: "KJNodes.ImageTransformKJ", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== "ImageTransformKJ") return; + + chainCallback(nodeType.prototype, "onNodeCreated", function () { + const node = this; + + const findW = (n) => this.widgets.find((w) => w.name === n); + const bboxWidget = findW("bboxes"), twWidget = findW("target_width"), thWidget = findW("target_height"); + const kpWidget = findW("keep_proportion"), divWidget = findW("divisible_by"); + const epWidget = findW("extra_padding"), icWidget = findW("invert_crop"); + bboxWidget.hidden = true; + + const _subWidgetCache = {}; + function _refreshSubWidgetCache() { + if (!node.widgets) return; + for (const w of node.widgets) { + if (w.name?.includes(".")) _subWidgetCache[w.name] = w; + } + } + function _getSubWidget(name) { + let w = _subWidgetCache[name]; + if (w && node.widgets?.includes(w)) return w; + // Cache miss — rescan + w = node.widgets?.find(ww => ww.name === name); + if (w) _subWidgetCache[name] = w; + return w || null; + } + + const _tmpCanvas = document.createElement("canvas"), _tmpCtx = _tmpCanvas.getContext("2d"); + const _mirrorCanvas = document.createElement("canvas"), _mirrorCtx = _mirrorCanvas.getContext("2d"); + const _rotCanvas = document.createElement("canvas"), _rotCtx = _rotCanvas.getContext("2d"); + let _rotCacheKey = ""; + function getRotatedPreview() { + if (!node._previewImg) return null; + const nw = node._previewImg.naturalWidth; + const nh = node._previewImg.naturalHeight; + const { w: effW, h: effH } = getEffectiveImageDims(true); + // Scale to reasonable preview size + let [rw, rh] = clampToMaxCanvas(effW, effH); + const key = `${node._previewImg.src}|${node._rotation}|${rw}|${rh}`; + if (_rotCacheKey === key) return _rotCanvas; + _rotCanvas.width = rw; + _rotCanvas.height = rh; + _rotCtx.clearRect(0, 0, rw, rh); + if (node._rotation === 0) { + _rotCtx.drawImage(node._previewImg, 0, 0, rw, rh); + } else { + const rad = Math.abs(rotRad()); + const rotW = Math.abs(nw * Math.cos(rad)) + Math.abs(nh * Math.sin(rad)); + const rotH = Math.abs(nw * Math.sin(rad)) + Math.abs(nh * Math.cos(rad)); + const fitScale = Math.min(rw / rotW, rh / rotH); + _rotCtx.translate(rw / 2, rh / 2); + _rotCtx.rotate(rotRad()); + _rotCtx.drawImage(node._previewImg, -(nw * fitScale) / 2, -(nh * fitScale) / 2, nw * fitScale, nh * fitScale); + _rotCtx.setTransform(1, 0, 0, 1, 0, 0); + } + _rotCacheKey = key; + return _rotCanvas; + } + + let _frameEp = { top: 0, bottom: 0, left: 0, right: 0, mode: "disabled", color: null, edgeMode: "clamp" }; + let _framePadXY = { x: 0.5, y: 0.5 }; + + const wrapper = document.createElement("div"); + wrapper.style.cssText = "display:flex;flex-direction:column;overflow:hidden;position:relative;pointer-events:auto;"; + const canvasEl = document.createElement("canvas"); + const canvasCtx = canvasEl.getContext("2d"); + canvasEl.style.cssText = "cursor:crosshair;display:block;width:100%;height:auto;"; + addWheelPassthrough(wrapper); + addMiddleClickPan(canvasEl); + + // Inject slider thumb styles once + if (!document.getElementById("kjcrop-slider-style")) { + const style = document.createElement("style"); + style.id = "kjcrop-slider-style"; + style.textContent = ` + .kjcrop-slider { -webkit-appearance: none; appearance: none; background: #444; border-radius: 4px; outline: none; } + .kjcrop-slider::-webkit-slider-thumb { -webkit-appearance: none; width: 8px; height: 18px; background: #46b4e6; border-radius: 4px; cursor: pointer; border: none; } + .kjcrop-slider::-moz-range-thumb { width: 8px; height: 18px; background: #46b4e6; border-radius: 4px; cursor: pointer; border: none; } + .kjcrop-btn { background: #333; border: 1px solid #555; border-radius: 4px; color: #bbb; font: 11px sans-serif; cursor: pointer; padding: 2px 8px; line-height: 16px; white-space: nowrap; flex-shrink: 0; transition: border-color 0.15s, color 0.15s, background 0.15s; text-align: center; } + .kjcrop-btn:hover { border-color: #46b4e6; color: #fff; background: #3a3a3a; } + .kjcrop-btn.active { border-color: #46b4e6; color: #46b4e6; background: #2a3a42; } + `; + document.head.appendChild(style); + } + // Grid size slider overlay (purely visual aid, not serialized) + const gridBar = document.createElement("div"); + gridBar.style.cssText = "display:flex;align-items:center;gap:6px;padding:3px 6px;margin-bottom:4px;font:11px sans-serif;color:#aaa;user-select:none;box-sizing:border-box;width:100%;flex:0 0 auto;"; + const gridLabel = document.createElement("span"); + gridLabel.textContent = "Grid: off"; + gridLabel.style.cssText = "min-width:62px;text-align:right;"; + const gridSlider = document.createElement("input"); + gridSlider.type = "range"; + gridSlider.min = "0"; + gridSlider.max = "256"; + gridSlider.step = "1"; + gridSlider.value = node.properties.gridSize ?? "0"; + gridSlider.className = "kjcrop-slider"; + gridSlider.style.cssText = "flex:1;height:8px;cursor:pointer;"; + function updateGridLabel() { + const v = parseInt(gridSlider.value); + gridLabel.textContent = v > 0 ? `Grid: ${v}px` : "Grid: off"; + } + updateGridLabel(); + let _gridStops = null; + function getGridStops() { + if (!node._previewImg) return null; + const tw = node._previewImg.naturalWidth; + const th = node._previewImg.naturalHeight; + const max = parseInt(gridSlider.max); + const set = new Set([0]); + for (let n = 1; n <= Math.max(tw, th); n++) { + const gw = tw / n, gh = th / n; + if (gw >= 2 && gw <= max) set.add(Math.round(gw * 10) / 10); + if (gh >= 2 && gh <= max) set.add(Math.round(gh * 10) / 10); + if (gw < 2 && gh < 2) break; + } + return [...set].sort((a, b) => a - b); + } + function invalidateGridStops() { _gridStops = null; } + gridSlider.addEventListener("input", () => { + let v = parseInt(gridSlider.value); + // Snap to divisors of the longer dimension for whole-cell grid + if (v > 0 && node._previewImg) { + if (!_gridStops) _gridStops = getGridStops(); + if (_gridStops) { + let best = 0, bestDist = Infinity; + for (const s of _gridStops) { + const d = Math.abs(s - v); + if (d < bestDist) { bestDist = d; best = s; } + } + v = best; + } + } + gridSlider.value = Math.round(v); + node.properties.gridSize = Math.round(v); + updateGridLabel(); + drawCanvas(); + }); + // Prevent drag events from propagating to the LiteGraph canvas, + // but only when not actively dragging a bbox handle + for (const evt of ["mousedown", "mousemove", "mouseup"]) { + gridSlider.addEventListener(evt, (e) => { + if (!node._drawing) e.stopPropagation(); + }); + } + // Toggle button to disable image preview (draw bboxes on blank canvas) + const previewBtn = document.createElement("button"); + previewBtn.title = "Auto: canvas sized to input image. Manual: canvas sized to target dimensions."; + previewBtn.className = "kjcrop-btn"; + previewBtn.style.width = "56px"; + if (node.properties.previewEnabled === undefined) node.properties.previewEnabled = true; + node._previewEnabled = node.properties.previewEnabled; + + function updatePreviewBtn() { + previewBtn.textContent = node._previewEnabled ? "Auto" : "Manual"; + previewBtn.classList.toggle("active", !node._previewEnabled); + } + updatePreviewBtn(); + + previewBtn.addEventListener("click", () => { + node._previewEnabled = !node._previewEnabled; + node.properties.previewEnabled = node._previewEnabled; + updatePreviewBtn(); + if (!node._previewEnabled) { + updateCanvasFromTargetDims(); + } else if (node._previewImg) { + // Restore canvas to image dimensions + const [cw, ch] = clampToMaxCanvas(node._previewImg.naturalWidth, node._previewImg.naturalHeight); + setCanvasSize(cw, ch); + } + drawCanvas(); + }); + previewBtn.addEventListener("mousedown", (e) => e.stopPropagation()); + + const colorSwatch = document.createElement("div"); + const colorInput = document.createElement("input"); + colorInput.type = "color"; + colorInput.value = node.properties.fillColor || "#000000"; + colorInput.style.cssText = "position:absolute;opacity:0;width:0;height:0;pointer-events:none;"; + colorSwatch.style.cssText = "width:18px;height:18px;border:1px solid #666;border-radius:3px;cursor:pointer;flex-shrink:0;background:" + (node.properties.fillColor || "#000000") + ";"; + colorSwatch.title = "Fill color (used by pad, invert crop, rotation)"; + colorSwatch.appendChild(colorInput); + colorSwatch.addEventListener("click", () => colorInput.click()); + colorSwatch.addEventListener("mousedown", (e) => e.stopPropagation()); + colorInput.addEventListener("input", () => { + node.properties.fillColor = colorInput.value; colorSwatch.style.background = colorInput.value; + updateBboxWidgets(); drawCanvas(); + }); + + function getEffectiveImageDims(expand = true) { + const w = node._previewImg.naturalWidth, h = node._previewImg.naturalHeight; + if (node._rotation === 0 || !expand) return { w, h }; + const rad = Math.abs(rotRad()); + const cosR = Math.cos(rad), sinR = Math.sin(rad); + return { w: Math.round(Math.abs(w * cosR) + Math.abs(h * sinR)), + h: Math.round(Math.abs(w * sinR) + Math.abs(h * cosR)) }; + } + + function renderReplicateRotated(destW, destH) { + const img = node._previewImg, nw = img.naturalWidth, nh = img.naturalHeight; + const pad = Math.max(nw, nh), pw = nw + pad * 2, ph = nh + pad * 2; + // Build replicate-padded image on _mirrorCanvas + _mirrorCanvas.width = pw; _mirrorCanvas.height = ph; + _mirrorCtx.drawImage(img, pad, pad, nw, nh); + // Edges (top, bottom, left, right) + _mirrorCtx.drawImage(img, 0, 0, nw, 1, pad, 0, nw, pad); + _mirrorCtx.drawImage(img, 0, nh - 1, nw, 1, pad, pad + nh, nw, pad); + _mirrorCtx.drawImage(img, 0, 0, 1, nh, 0, pad, pad, nh); + _mirrorCtx.drawImage(img, nw - 1, 0, 1, nh, pad + nw, pad, pad, nh); + // Corners (TL, TR, BL, BR) + _mirrorCtx.drawImage(img, 0, 0, 1, 1, 0, 0, pad, pad); + _mirrorCtx.drawImage(img, nw - 1, 0, 1, 1, pad + nw, 0, pad, pad); + _mirrorCtx.drawImage(img, 0, nh - 1, 1, 1, 0, pad + nh, pad, pad); + _mirrorCtx.drawImage(img, nw - 1, nh - 1, 1, 1, pad + nw, pad + nh, pad, pad); + // Rotate padded image, crop to original size, scale to dest + _tmpCanvas.width = destW; _tmpCanvas.height = destH; + _tmpCtx.save(); + _tmpCtx.scale(destW / nw, destH / nh); + _tmpCtx.translate(nw / 2, nh / 2); + _tmpCtx.rotate(rotRad()); + _tmpCtx.drawImage(_mirrorCanvas, -(pw / 2), -(ph / 2)); + _tmpCtx.restore(); + } + + // Resize node to fit the current canvas/image. If onlyGrow is true, + // the node will only get larger, never shrink (preserves user resize). + // If false, the node is set to the exact size needed. + function fitNodeToImage(onlyGrow) { + if (!node._previewEnabled) { + // Manual mode: only grow to fit current canvas, never change width + if (onlyGrow && canvasEl.width > 0 && canvasEl.height > 0) { + _resizing = true; + const displayedH = Math.round(canvasEl.height * ((node.size[0] - 30) / canvasEl.width)); + node._widgetHeight = displayedH + GRID_BAR_HEIGHT; + const aboveH = node.cropEditor?.last_y || 350; + const neededH = aboveH + node._widgetHeight + 20; + if (node.size[1] < neededH) node.setSize([node.size[0], neededH]); + _resizing = false; + } + return; + } + if (!node._previewImg) return; + _resizing = true; + const { w: effW, h: effH } = getEffectiveImageDims(true); + const ar = effH / effW; + const nodeW = onlyGrow ? node.size[0] : canvasEl.width + 30; + const canvasDisplayH = Math.round((nodeW - 30) * ar); + node._widgetHeight = canvasDisplayH + GRID_BAR_HEIGHT; + const aboveH = node.cropEditor?.last_y || 350; + // +20 accounts for DOM widget margin (10px top + 10px bottom) + const neededH = aboveH + node._widgetHeight + 20; + const newW = onlyGrow ? Math.max(node.size[0], nodeW) : nodeW; + const newH = onlyGrow ? Math.max(node.size[1], neededH) : neededH; + if (newW !== node.size[0] || newH !== node.size[1]) { + node.setSize([newW, newH]); + } + _resizing = false; + // Force reflow and redraw so the layout updates immediately + void canvasEl.offsetHeight; + if (node.graph) { + node.graph.setDirtyCanvas(true, true); + requestAnimationFrame(() => { + drawCanvas(); + if (node.graph) node.graph.setDirtyCanvas(true, true); + }); + } + } + + function resetRotation() { + if (node._rotation === 0) return; + node._rotation = 0; + node.properties.rotation = 0; + restoreCanvasToImage(); + fitNodeToImage(false); + updateBboxWidgets(); + drawCanvas(); + } + + const rotateBtn = document.createElement("button"); + rotateBtn.title = "Toggle rotation cross (drag to rotate image)"; + rotateBtn.className = "kjcrop-btn"; + if (node.properties.showRotationCross === undefined) node.properties.showRotationCross = false; + node._showRotationCross = node.properties.showRotationCross; + + function updateRotateBtn() { + rotateBtn.textContent = "Rotate"; + rotateBtn.classList.toggle("active", node._showRotationCross); + } + updateRotateBtn(); + + rotateBtn.addEventListener("click", () => { + node._showRotationCross = !node._showRotationCross; + node.properties.showRotationCross = node._showRotationCross; + updateRotateBtn(); + drawCanvas(); + }); + rotateBtn.addEventListener("mousedown", (e) => e.stopPropagation()); + rotateBtn.addEventListener("contextmenu", (e) => { + e.preventDefault(); + resetRotation(); + }); + + const gridColorInput = document.createElement("input"); + gridColorInput.type = "color"; + gridColorInput.value = node.properties.gridColor || "#ffffff"; + gridColorInput.style.cssText = "position:absolute;opacity:0;width:0;height:0;pointer-events:none;"; + const gridColorSwatch = document.createElement("div"); + gridColorSwatch.style.cssText = "width:18px;height:18px;border:1px solid #666;border-radius:3px;cursor:pointer;flex-shrink:0;background:" + (node.properties.gridColor || "#ffffff") + ";"; + gridColorSwatch.title = "Grid color"; + gridColorSwatch.appendChild(gridColorInput); + gridColorSwatch.addEventListener("click", () => gridColorInput.click()); + gridColorSwatch.addEventListener("mousedown", (e) => e.stopPropagation()); + let _gridColorRGB = null; + function getGridColorRGB() { + const gc = node.properties.gridColor || "#ffffff"; + if (!_gridColorRGB || _gridColorRGB.hex !== gc) { + _gridColorRGB = { hex: gc, r: parseInt(gc.slice(1, 3), 16), g: parseInt(gc.slice(3, 5), 16), b: parseInt(gc.slice(5, 7), 16) }; + } + return _gridColorRGB; + } + gridColorInput.addEventListener("input", () => { + node.properties.gridColor = gridColorInput.value; gridColorSwatch.style.background = gridColorInput.value; + _gridColorRGB = null; + drawCanvas(); + }); + + for (const el of [previewBtn, rotateBtn, colorSwatch, gridLabel, gridColorSwatch, gridSlider]) gridBar.appendChild(el); + wrapper.appendChild(gridBar); wrapper.appendChild(canvasEl); + + const GRID_BAR_HEIGHT = 46; + node._widgetHeight = 300 + GRID_BAR_HEIGHT; + + node.cropEditor = this.addDOMWidget("crop_preview", "CropPreviewWidget", wrapper, { + serialize: false, hideOnZoom: false, + getMinHeight: () => GRID_BAR_HEIGHT + 50, + }); + + this.resizable = true; + // Set default width wider than LiteGraph's default — deferred so onConfigure can set the flag first + setTimeout(() => { + if (!node._cropConfigured) node.setSize([450, node.size[1]]); + }, 0); + + Object.assign(node, { _previewImg: null, _bboxes: [], _activeIdx: -1, + _drawing: false, _dragMode: null, _dragStart: null, _bboxAtDragStart: null, + _rotation: node.properties.rotation || 0 }); + + for (const w of [twWidget, thWidget, kpWidget, divWidget, epWidget, icWidget]) { + if (!w) continue; + const origCb = w.callback; + w.callback = function (...args) { + if (origCb) origCb.apply(this, args); + updateCanvasFromTargetDims(); + drawCanvas(); + fitNodeToCanvas(); + fitNodeToImage(true); + if (w === kpWidget || w === icWidget || w === epWidget) { + // Restore values with increasing delays to catch ComfyUI's own default-setting + const restoreSaved = () => { + hookSubWidgets(); + _refreshSubWidgetCache(); + if (w === epWidget && node._savedEpValues && epWidget.value !== "disabled") { + for (const [n, v] of Object.entries(node._savedEpValues)) { + const sw = _getSubWidget(`extra_padding.${n}`); + if (sw) sw.value = v; + } + } else if (w === kpWidget && node._savedKpValues) { + for (const [n, v] of Object.entries(node._savedKpValues)) { + const sw = _getSubWidget(`keep_proportion.${n}`); + if (sw) sw.value = v; + } + } + drawCanvas(); + fitNodeToCanvas(); + fitNodeToImage(true); + }; + for (const delay of [50, 100, 200, 300, 500, 1000]) { + setTimeout(restoreSaved, delay); + } + } + }; + } + + const hookSubWidgets = () => { + _refreshSubWidgetCache(); + for (const subName of ["keep_proportion.edge_mode", "keep_proportion.pad_x", "keep_proportion.pad_y", "keep_proportion.width_mult", "keep_proportion.height_mult", "extra_padding.pad_top", "extra_padding.pad_bottom", "extra_padding.pad_left", "extra_padding.pad_right", "extra_padding.edge_mode", ]) { + const sw = _getSubWidget(subName); + if (sw && !sw._cropPreviewHooked) { + sw._cropPreviewHooked = true; + // Restore saved values from mode switch + const shortName = subName.split(".")[1]; + if (subName.startsWith("extra_padding.") && node._savedEpValues && shortName in node._savedEpValues) { + sw.value = node._savedEpValues[shortName]; + } else if (subName.startsWith("keep_proportion.") && node._savedKpValues && shortName in node._savedKpValues) { + sw.value = node._savedKpValues[shortName]; + } + // Hook callback — save values on change so they persist across mode switches + const swOrig = sw.callback; + sw.callback = function (...a) { + if (_suppressSubCallbacks) return; + if (subName.startsWith("extra_padding.")) node._savedEpValues[shortName] = sw.value; + else if (subName.startsWith("keep_proportion.")) node._savedKpValues[shortName] = sw.value; + if (swOrig) swOrig.apply(this, a); + drawCanvas(); + }; + } + } + }; + for (const delay of [50, 150, 300, 500]) { + setTimeout(hookSubWidgets, delay); + } + + let _lastWidgetSnapshot = ""; + const _snapshotNames = [ + "extra_padding", "extra_padding.pad_top", "extra_padding.pad_bottom", + "extra_padding.pad_left", "extra_padding.pad_right", + "extra_padding.edge_mode", + "keep_proportion.pad_x", "keep_proportion.pad_y", "keep_proportion.color", + "keep_proportion.edge_mode", "keep_proportion.width_mult", "keep_proportion.height_mult", + ]; + node._savedEpValues = node._savedEpValues || {}; + node._savedKpValues = node._savedKpValues || {}; + const origDrawFg = node.onDrawForeground; + node.onDrawForeground = function (...args) { + if (origDrawFg) origDrawFg.apply(this, args); + const snap = _snapshotNames.map(n => _getSubWidget(n)?.value ?? (n === "extra_padding" ? epWidget?.value : "") ?? "").join("|"); + if (snap !== _lastWidgetSnapshot) { + _lastWidgetSnapshot = snap; + drawCanvas(); + } + }; + + // Restore bboxes from widget + restoreBboxesFromWidget(); + + function restoreBboxesFromWidget() { + if (bboxWidget.value) { + try { + const parsed = JSON.parse(bboxWidget.value); + // New format: { bboxes: [...], rotation: N } + if (parsed && !Array.isArray(parsed) && typeof parsed === "object") { + if (parsed.bboxes && Array.isArray(parsed.bboxes)) { + node._bboxes = parsed.bboxes.filter((b) => b?.startX != null); + node._activeIdx = node._bboxes.length > 0 ? 0 : -1; + } + if (parsed.rotation !== undefined) { + node._rotation = parsed.rotation; + node.properties.rotation = parsed.rotation; + } + if (parsed.fillColor) { + node.properties.fillColor = parsed.fillColor; + colorInput.value = parsed.fillColor; + colorSwatch.style.background = parsed.fillColor; + } + } + // Legacy format: [bbox, bbox, ...] + else if (Array.isArray(parsed) && parsed.length > 0) { + node._bboxes = parsed.filter((b) => b?.startX != null); + node._activeIdx = node._bboxes.length > 0 ? 0 : -1; + } + } catch (e) {} + } + } + + function normBbox(bbox) { + return { x1: Math.min(bbox.startX, bbox.endX), y1: Math.min(bbox.startY, bbox.endY), + x2: Math.max(bbox.startX, bbox.endX), y2: Math.max(bbox.startY, bbox.endY) }; + } + + function getCanvasMouse(e) { + const rect = canvasEl.getBoundingClientRect(); + return { x: (e.clientX - rect.left) * (canvasEl.width / rect.width), + y: (e.clientY - rect.top) * (canvasEl.height / rect.height) }; + } + + function rotRad() { return node._rotation * Math.PI / 180; } + function getDivBy() { return divWidget ? divWidget.value : 0; } + function getGridSize() { return parseInt(gridSlider.value) || 0; } + // Subtract extra padding from target dims for non-pad keep_proportion modes + function adjustTargetForPadding(rw, rh, tgtW, tgtH) { + const ep = getExtraPadding(); + if (hasExtraPad(ep) && !ep.isFirst) { + if (tgtW > 0) rw = Math.max(1, rw - ep.left - ep.right); + if (tgtH > 0) rh = Math.max(1, rh - ep.top - ep.bottom); + } + return { rw, rh }; + } + function getEffectiveTargetDims(fallbackW, fallbackH) { + const divBy = getDivBy(); + const { w: tgtW, h: tgtH } = getTgt(); + let { rw, rh } = adjustTargetForPadding(tgtW > 0 ? tgtW : fallbackW, tgtH > 0 ? tgtH : fallbackH, tgtW, tgtH); + return { rw: roundDown(rw, divBy), rh: roundDown(rh, divBy) }; + } + + function getFillColor() { return node.properties.fillColor || "#000000"; } + function getTgt() { return { w: twWidget ? twWidget.value : 0, h: thWidget ? thWidget.value : 0 }; } + function getKp() { return kpWidget ? kpWidget.value : "stretch"; } + function isPadMode(kp) { return kp === "pad_color" || kp === "pad_edge"; } + function hasExtraPad(ep) { return ep.top + ep.bottom + ep.left + ep.right > 0; } + + function hitTestRotationCross(mx, my) { + if (!node._showRotationCross) return false; + const cx = canvasEl.width / 2; + const cy = canvasEl.height / 2; + const dx = mx - cx; + const dy = my - cy; + const rad = rotRad(); + const cosR = Math.cos(rad); + const sinR = Math.sin(rad); + const perpH = -dx * sinR + dy * cosR; + const perpV = dx * cosR + dy * sinR; + const dist = Math.sqrt(dx * dx + dy * dy); + return dist < 15 || Math.abs(perpH) < 8 || Math.abs(perpV) < 8; + } + + function restoreCanvasToImage() { + if (!node._previewImg || !node._previewEnabled) return; + setCanvasSize(...clampToMaxCanvas(...Object.values(getEffectiveImageDims(true)))); + } + + function drawPreviewRegion(ctx, rx, ry, rw, rh, dx, dy, dw, dh) { + const rp = getRotatedPreview(); + const sx = _rotCanvas.width / canvasEl.width, sy = _rotCanvas.height / canvasEl.height; + ctx.drawImage(rp, rx * sx, ry * sy, rw * sx, rh * sy, dx, dy, dw, dh); + } + + function drawStretchOverlay(ctx, ar, rx, ry, rw, rh) { + const { pW, pH, pX, pY } = fitAspect(ar, rw, rh, rx, ry); + ctx.fillStyle = "rgba(0, 0, 0, 0.5)"; + ctx.fillRect(rx, ry, rw, rh); + if (node._previewEnabled && node._previewImg) drawPreviewRegion(ctx, rx, ry, rw, rh, pX, pY, pW, pH); + drawDashedBorder(ctx, pX, pY, pW, pH); + } + + // Compute content dimensions after keep_proportion resize, before extra padding + function computeContentDims(cropW, cropH) { + if (cropW <= 0 || cropH <= 0) return null; + const { w: tgtW, h: tgtH } = getTgt(); + let outW = cropW, outH = cropH; + const kp = getKp(); + // For pad-first non-pad modes, the source is padded before keep_proportion + const ep = getExtraPadding(); + if (ep.isFirst && hasExtraPad(ep) && !isPadMode(kp)) { + cropW += ep.left + ep.right; + cropH += ep.top + ep.bottom; + } + if (tgtW > 0 || tgtH > 0) { + let { rw, rh } = adjustTargetForPadding(tgtW > 0 ? tgtW : cropW, tgtH > 0 ? tgtH : cropH, tgtW, tgtH); + if (kp === "keep_long_edge") { + const ratio = Math.min(rw / cropW, rh / cropH); + rw = Math.round(cropW * ratio); + rh = Math.round(cropH * ratio); + } else if (kp === "keep_short_edge") { + const ratio = Math.max(rw / cropW, rh / cropH); + rw = Math.round(cropW * ratio); + rh = Math.round(cropH * ratio); + } else if (kp === "total_pixels") { + const totalPx = rw * rh; + const ar = cropW / cropH; + rh = Math.round(Math.sqrt(totalPx / ar)); + rw = Math.round(Math.sqrt(totalPx * ar)); + } + outW = rw; + outH = rh; + } + if (kp === "multiplier") { + const wm = _getSubWidget("keep_proportion.width_mult")?.value || 1.0; + const hm = _getSubWidget("keep_proportion.height_mult")?.value || 1.0; + outW = Math.round(cropW * wm); + outH = Math.round(cropH * hm); + } + return { w: outW, h: outH }; + } + + function computeOutputDims(cropW, cropH) { + const content = computeContentDims(cropW, cropH); + if (!content) return null; + let outW = content.w, outH = content.h; + // Add extra padding (skip for pad-first — handled on the source image) + const ep = getExtraPadding(); + if (!ep.isFirst) { + outW += ep.left + ep.right; + outH += ep.top + ep.bottom; + } + const divBy = getDivBy(); + outW = roundDown(outW, divBy); + outH = roundDown(outH, divBy); + return { w: outW, h: outH }; + } + + function getPadXY() { return _framePadXY; } + function getExtraPadding() { return _frameEp; } + + function _readPadXY() { + const px = _getSubWidget("keep_proportion.pad_x"), py = _getSubWidget("keep_proportion.pad_y"); + return { x: px ? px.value : 0.5, y: py ? py.value : 0.5 }; + } + + let _suppressSubCallbacks = false; + function setPadXY(x, y, defer) { + x = Math.max(0, Math.min(1, x)); + y = Math.max(0, Math.min(1, y)); + _framePadXY = { x, y }; + if (!defer) { + const px = _getSubWidget("keep_proportion.pad_x"), py = _getSubWidget("keep_proportion.pad_y"); + _suppressSubCallbacks = true; + if (px) px.value = x; + if (py) py.value = y; + _suppressSubCallbacks = false; + } + } + + function _readExtraPadding() { + const get = (name) => _getSubWidget(`extra_padding.${name}`)?.value || 0; + const mode = epWidget?.value || "disabled"; + // pad_color/pad_edge = pad first (before crop), pad_crop_color/pad_crop_edge = pad after crop + const isCrop = mode.startsWith("pad_crop_"); + const baseMode = isCrop ? mode.replace("pad_crop_", "pad_") : mode; + return { + top: get("pad_top"), bottom: get("pad_bottom"), + left: get("pad_left"), right: get("pad_right"), + mode: baseMode, + isFirst: !isCrop && mode !== "disabled", + edgeMode: _getSubWidget("extra_padding.edge_mode")?.value || "clamp", + }; + } + + // Draw extra padding fill around content area (cX,cY,cW,cH) within pad area (pX,pY,pW,pH) + // srcImg is the source image, srcX/Y/W/H define the source region + // Draw edge/color fill around content area. srcImg must be pre-rendered at its full dimensions. + function drawExtraPadFill(ctx, ep, pX, pY, pW, pH, cX, cY, cW, cH, srcImg) { + if (ep.mode === "pad_color") { + ctx.fillStyle = getFillColor(); + ctx.fillRect(pX, pY, pW, pH); + } else if (ep.mode === "pad_edge") { + const em = ep.edgeMode; + const sw = srcImg.width, sh = srcImg.height; + if (em === "clamp") { + const l = cX - pX, r = (pX + pW) - (cX + cW), t = cY - pY, b = (pY + pH) - (cY + cH); + if (t > 0) ctx.drawImage(srcImg, 0, 0, sw, 1, cX, pY, cW, t); + if (b > 0) ctx.drawImage(srcImg, 0, sh - 1, sw, 1, cX, cY + cH, cW, b); + if (l > 0) ctx.drawImage(srcImg, 0, 0, 1, sh, pX, cY, l, cH); + if (r > 0) ctx.drawImage(srcImg, sw - 1, 0, 1, sh, cX + cW, cY, r, cH); + if (l > 0 && t > 0) ctx.drawImage(srcImg, 0, 0, 1, 1, pX, pY, l, t); + if (r > 0 && t > 0) ctx.drawImage(srcImg, sw - 1, 0, 1, 1, cX + cW, pY, r, t); + if (l > 0 && b > 0) ctx.drawImage(srcImg, 0, sh - 1, 1, 1, pX, cY + cH, l, b); + if (r > 0 && b > 0) ctx.drawImage(srcImg, sw - 1, sh - 1, 1, 1, cX + cW, cY + cH, r, b); + } else if (em === "repeat" || em === "mirror") { + // Copy srcImg to _mirrorCanvas first — srcImg may be _tmpCanvas + const tileW = Math.max(1, Math.round(cW)); + const tileH = Math.max(1, Math.round(cH)); + _mirrorCanvas.width = tileW; + _mirrorCanvas.height = tileH; + _mirrorCtx.drawImage(srcImg, 0, 0, tileW, tileH); + let tileCanvas; + if (em === "mirror") { + _tmpCanvas.width = tileW * 2; _tmpCanvas.height = tileH * 2; + _tmpCtx.drawImage(_mirrorCanvas, 0, 0); + _tmpCtx.save(); _tmpCtx.translate(_tmpCanvas.width, 0); _tmpCtx.scale(-1, 1); _tmpCtx.drawImage(_mirrorCanvas, 0, 0); _tmpCtx.restore(); + _tmpCtx.save(); _tmpCtx.translate(0, _tmpCanvas.height); _tmpCtx.scale(1, -1); _tmpCtx.drawImage(_mirrorCanvas, 0, 0); _tmpCtx.restore(); + _tmpCtx.save(); _tmpCtx.translate(_tmpCanvas.width, _tmpCanvas.height); _tmpCtx.scale(-1, -1); _tmpCtx.drawImage(_mirrorCanvas, 0, 0); _tmpCtx.restore(); + tileCanvas = _tmpCanvas; + } else { + tileCanvas = _mirrorCanvas; + } + const pattern = ctx.createPattern(tileCanvas, "repeat"); + ctx.save(); + ctx.beginPath(); ctx.rect(pX, pY, pW, pH); ctx.clip(); + ctx.translate(cX, cY); + ctx.fillStyle = pattern; + ctx.fillRect(pX - cX, pY - cY, pW, pH); + ctx.restore(); + } + } + } + + function fitAspect(targetAR, cw, ch, cx = 0, cy = 0) { + if (targetAR > cw / ch) { const pH = cw / targetAR; return { pW: cw, pH, pX: cx, pY: cy + (ch - pH) / 2 }; } + const pW = ch * targetAR; return { pW, pH: ch, pX: cx + (cw - pW) / 2, pY: cy }; + } + + function drawDashedBorder(ctx, x, y, w, h, color = "rgba(255, 180, 60, 0.8)") { + ctx.setLineDash([4, 4]); ctx.strokeStyle = color; ctx.lineWidth = 1; + ctx.strokeRect(x, y, w, h); ctx.setLineDash([]); + } + + // Dim four sides around an inner rect within an outer rect + function dimSides(ctx, ox, oy, ow, oh, ix, iy, iw, ih, color = "rgba(255, 80, 80, 0.3)") { + ctx.fillStyle = color; + ctx.fillRect(ox, oy, ow, iy - oy); // top + ctx.fillRect(ox, iy + ih, ow, (oy + oh) - (iy + ih)); // bottom + ctx.fillRect(ox, iy, ix - ox, ih); // left + ctx.fillRect(ix + iw, iy, (ox + ow) - (ix + iw), ih); // right + } + + // Compute pad mode layout: fit target AR, compute content rect inside + function computePadLayout(rw, rh, cropW, cropH, fitW, fitH, fitX, fitY) { + const { pW, pH, pX, pY } = fitAspect(rw / rh, fitW, fitH, fitX, fitY); + const imgRatio = Math.min(rw / cropW, rh / cropH); + const cW = pW * (cropW * imgRatio / rw), cH = pH * (cropH * imgRatio / rh); + const { cX, cY } = computePadContentPos(pX, pY, pW, pH, cW, cH); + return { pW, pH, pX, pY, cW, cH, cX, cY }; + } + + // Render source to _tmpCanvas at given size. If srcImg + source rect provided, draws that region. + // Otherwise draws the rotated preview (with replicate rotation for pad_edge mode). + // replicate: force replicate-pad rotation (for keep_proportion pad_edge mode) + function renderToTmp(w, h, srcImg, sx, sy, sw, sh, replicate) { + const ep = getExtraPadding(); + if (!srcImg && node._rotation !== 0 && (replicate || (ep.mode === "pad_edge" && ep.isFirst))) { + renderReplicateRotated(w, h); return; + } + _tmpCanvas.width = w; _tmpCanvas.height = h; + if (srcImg) _tmpCtx.drawImage(srcImg, sx, sy, sw, sh, 0, 0, w, h); + else _tmpCtx.drawImage(getRotatedPreview(), 0, 0, _rotCanvas.width, _rotCanvas.height, 0, 0, w, h); + } + + function getEdgeEp() { + return { mode: "pad_edge", edgeMode: _getSubWidget("keep_proportion.edge_mode")?.value || "clamp" }; + } + + function computePadContentPos(pX, pY, pW, pH, cW, cH) { + const { x, y } = getPadXY(); + return { cX: pX + (pW - cW) * x, cY: pY + (pH - cH) * y }; + } + + function getOtherBboxEdges(axis) { + const edges = []; + for (let i = 0; i < node._bboxes.length; i++) { + if (i === node._activeIdx) continue; + const { x1, y1, x2, y2 } = normBbox(node._bboxes[i]); + edges.push(axis === "x" ? x1 : y1, axis === "x" ? x2 : y2); + } + return edges; + } + + function getEffectiveSnap() { + const g = getGridParams(); + if (g) return { x: g.effGridW, y: g.effGridH }; + const d = getDivBy(); + return d > 1 ? { x: d, y: d } : { x: 0, y: 0 }; + } + + function snapCoord(val, snapSize, scale, canvasDim, otherEdges) { + let best = val, bestDist = Infinity; + const consider = (c) => { const d = Math.abs(val - c); if (d < bestDist) { bestDist = d; best = c; } }; + consider(0); + consider(canvasDim); + if (snapSize > 1) consider(Math.round(val * scale / snapSize) * snapSize / scale); + if (otherEdges) for (const edge of otherEdges) consider(edge); + return best; + } + + // Snap the moving corner(s) of a bbox to grid / edges / other bboxes. + function snapBbox(bbox) { + if (!bbox) return bbox; + const snap = getEffectiveSnap(), scale = getCanvasScale(); + const xEdges = getOtherBboxEdges("x"), yEdges = getOtherBboxEdges("y"); + return { startX: snapCoord(bbox.startX, snap.x, scale.x, canvasEl.width, xEdges), + startY: snapCoord(bbox.startY, snap.y, scale.y, canvasEl.height, yEdges), + endX: snapCoord(bbox.endX, snap.x, scale.x, canvasEl.width, xEdges), + endY: snapCoord(bbox.endY, snap.y, scale.y, canvasEl.height, yEdges) }; + } + + function snapBboxPosition(bbox) { + if (!bbox) return bbox; + const snap = getEffectiveSnap(); + const scale = getCanvasScale(); + const xEdges = getOtherBboxEdges("x"); + const yEdges = getOtherBboxEdges("y"); + const { x1, y1 } = normBbox(bbox); + const w = Math.abs(bbox.endX - bbox.startX); + const h = Math.abs(bbox.endY - bbox.startY); + // Snap both the left/top and right/bottom edges, pick whichever is closer + const snappedX1 = snapCoord(x1, snap.x, scale.x, canvasEl.width, xEdges); + const snappedX2 = snapCoord(x1 + w, snap.x, scale.x, canvasEl.width, xEdges) - w; + const snappedX = Math.abs(snappedX1 - x1) <= Math.abs(snappedX2 - x1) ? snappedX1 : snappedX2; + const snappedY1 = snapCoord(y1, snap.y, scale.y, canvasEl.height, yEdges); + const snappedY2 = snapCoord(y1 + h, snap.y, scale.y, canvasEl.height, yEdges) - h; + const snappedY = Math.abs(snappedY1 - y1) <= Math.abs(snappedY2 - y1) ? snappedY1 : snappedY2; + return { startX: snappedX, startY: snappedY, endX: snappedX + w, endY: snappedY + h }; + } + + function constrainAspect(bbox, dragMode) { + if (!bbox) return bbox; + const { w: tw, h: th } = getTgt(); + // Target aspect in source-image pixels + const srcAR = (tw > 0 && th > 0) ? tw / th : 1; + // Convert to preview-pixel aspect by accounting for non-uniform scaling + const scale = getCanvasScale(); + const ar = srcAR * scale.y / scale.x; + + const { x1, y1, x2, y2 } = normBbox(bbox); + let bw = x2 - x1; + let bh = y2 - y1; + if (bw <= 0 || bh <= 0) return bbox; + + // Adjust height to match aspect, keep width + const hFromW = bw / ar; + const wFromH = bh * ar; + // Pick the smaller fit so we don't exceed the drawn extent + if (hFromW <= bh) { + bh = hFromW; + } else { + bw = wFromH; + } + + // Anchor based on drag mode + if (dragMode === "resize-tl") { + return { startX: x2 - bw, startY: y2 - bh, endX: x2, endY: y2 }; + } else if (dragMode === "resize-tr") { + return { startX: x1, startY: y2 - bh, endX: x1 + bw, endY: y2 }; + } else if (dragMode === "resize-bl") { + return { startX: x2 - bw, startY: y1, endX: x2, endY: y1 + bh }; + } else { + // resize-br or draw (anchor top-left) + return { startX: x1, startY: y1, endX: x1 + bw, endY: y1 + bh }; + } + } + + function removeBbox(index) { + node._bboxes.splice(index, 1); + if (node._bboxes.length === 0) node._activeIdx = -1; + else if (index <= node._activeIdx) node._activeIdx = Math.min(node._activeIdx - (index < node._activeIdx ? 1 : 0), node._bboxes.length - 1); + } + + function hitTestBboxes(mx, my) { + const order = node._activeIdx >= 0 ? [node._activeIdx] : []; + for (let i = 0; i < node._bboxes.length; i++) { if (i !== node._activeIdx) order.push(i); } + for (const idx of order) { + const { x1, y1, x2, y2 } = normBbox(node._bboxes[idx]); + const mode = rectHitTest(mx, my, x1, y1, x2, y2, 10); + if (mode) return { index: idx, mode }; + } + return null; + } + + canvasEl.addEventListener("mousedown", (e) => { + canvasEl.focus(); + if (e.button === 2) { + e.preventDefault(); + const m = getCanvasMouse(e); + // Right-click on rotation cross: reset rotation + if (node._rotation !== 0 && hitTestRotationCross(m.x, m.y)) { + resetRotation(); + return; + } + // Right-click: delete the bbox under cursor, or active bbox + const hit = hitTestBboxes(m.x, m.y); + if (hit) { + removeBbox(hit.index); + } else if (node._activeIdx >= 0) { + removeBbox(node._activeIdx); + } + updateBboxWidgets(); + drawCanvas(); + return; + } + if (e.button !== 0) return; + const m = getCanvasMouse(e); + + // Check for rotation cross drag (click near the center or any arm) + if (hitTestRotationCross(m.x, m.y)) { + node._dragMode = "rotate"; + node._rotateStart = node._rotation; + const rect0 = canvasEl.getBoundingClientRect(); + // Capture screen center once — immune to canvas resize + node._rotateCenterX = rect0.left + rect0.width / 2; + node._rotateCenterY = rect0.top + rect0.height / 2; + node._rotateScreenStartAngle = Math.atan2( + e.clientY - node._rotateCenterY, + e.clientX - node._rotateCenterX + ) * 180 / Math.PI; + node._drawing = true; + node._dragStart = m; + document.addEventListener("mousemove", onDragMove); + document.addEventListener("mouseup", onDragEnd); + e.preventDefault(); + e.stopPropagation(); + return; + } + + // Shift+drag: drag content position within padded area + if (e.shiftKey) { + const kpVal = getKp(); + const { w: tW2, h: tH2 } = getTgt(); + const ep = getExtraPadding(); + const startDrag = (mode, extra) => { + node._dragMode = mode; Object.assign(node, extra); + node._drawing = true; node._dragStart = m; + document.addEventListener("mousemove", onDragMove); + document.addEventListener("mouseup", onDragEnd); + e.preventDefault(); e.stopPropagation(); + }; + if ((tW2 > 0 || tH2 > 0) && isPadMode(kpVal)) { + startDrag("pad_drag", { _padDragStart: getPadXY(), _bboxAtDragStart: null }); return; + } else if (hasExtraPad(ep)) { + startDrag("extra_pad_drag", { _extraPadStart: { ...ep }, _bboxAtDragStart: null }); return; + } + } + + const hit = hitTestBboxes(m.x, m.y); + if (hit) { + // Select and interact with existing bbox + node._activeIdx = hit.index; + node._dragMode = hit.mode; + const bbox = node._bboxes[hit.index]; + const { x1, y1, x2, y2 } = normBbox(bbox); + node._bboxAtDragStart = { startX: x1, startY: y1, endX: x2, endY: y2 }; + } else { + // Draw new bbox + node._dragMode = "draw"; + const newBbox = { startX: m.x, startY: m.y, endX: m.x, endY: m.y }; + node._bboxes.push(newBbox); + node._activeIdx = node._bboxes.length - 1; + node._bboxAtDragStart = null; + } + node._drawing = true; + node._dragStart = m; + + document.addEventListener("mousemove", onDragMove); + document.addEventListener("mouseup", onDragEnd); + e.preventDefault(); + e.stopPropagation(); + }); + + // Hover cursor updates + canvasEl.addEventListener("mousemove", (e) => { + if (node._drawing) return; + const m = getCanvasMouse(e); + const hit = hitTestBboxes(m.x, m.y); + if (!hit) { + if (!node._showRotationCross) { + canvasEl.style.cursor = e.shiftKey ? "grab" : "crosshair"; + } else if (hitTestRotationCross(m.x, m.y)) { + canvasEl.style.cursor = "alias"; + } else if (e.shiftKey) { + const kpVal = getKp(); + const { w: tW, h: tH } = getTgt(); + const epHover = getExtraPadding(); + if (((tW > 0 || tH > 0) && isPadMode(kpVal)) || hasExtraPad(epHover)) { + canvasEl.style.cursor = "grab"; + } else { + canvasEl.style.cursor = "crosshair"; + } + } else { + canvasEl.style.cursor = "crosshair"; + } + } else { + canvasEl.style.cursor = cursorForBboxMode(hit.mode) || "crosshair"; + } + }); + + function onDragMove(e) { + if (!node._drawing) return; + if (node._dragMode === "rotate") { + // Screen-space atan2 with fixed center (captured at drag start) + const sdx = e.clientX - node._rotateCenterX; + const sdy = e.clientY - node._rotateCenterY; + const screenDist = Math.sqrt(sdx * sdx + sdy * sdy); + if (screenDist < 15) { + node._rotateStart = node._rotation; + node._rotateScreenStartAngle = Math.atan2(sdy, sdx) * 180 / Math.PI; + return; + } + const angleCur = Math.atan2(sdy, sdx) * 180 / Math.PI; + let angleDelta = angleCur - node._rotateScreenStartAngle; + // Normalize delta to -180..180 to handle atan2 wrap-around + while (angleDelta > 180) angleDelta -= 360; + while (angleDelta < -180) angleDelta += 360; + let newRot = node._rotateStart + angleDelta; + if (e.ctrlKey) { + newRot = Math.round(newRot / 15) * 15; + } else { + for (const snap of [0, 90, 180, 270, -90, -180, -270, 360]) { + if (Math.abs(newRot - snap) < 3) { newRot = snap; break; } + } + } + while (newRot > 180) newRot -= 360; + while (newRot < -180) newRot += 360; + node._rotation = newRot; + node.properties.rotation = newRot; + // Resize canvas to match rotated image dimensions + if (node._previewImg && node._previewEnabled) { + const { w: ew, h: eh } = getEffectiveImageDims(true); + const [rw, rh] = clampToMaxCanvas(ew, eh); + if (Math.abs(rw - canvasEl.width) > 3 || Math.abs(rh - canvasEl.height) > 3) { + setCanvasSize(rw, rh); + } + } + drawCanvas(); + return; + } + if (node._dragMode === "pad_drag") { + const m = getCanvasMouse(e); + const dx = m.x - node._dragStart.x; + const dy = m.y - node._dragStart.y; + // Compute the pad preview layout to find how many canvas pixels = full 0-1 range + const ep = getExtraPadding(); + const { w: imgW, h: imgH } = node._previewImg + ? getEffectiveImageDims(getKp() !== "pad_edge") + : { w: canvasEl.width, h: canvasEl.height }; + const { rw: baseTw, rh: baseTh } = getEffectiveTargetDims(imgW, imgH); + const totalW = baseTw + ep.left + ep.right; + const totalH = baseTh + ep.top + ep.bottom; + const { pW, pH } = fitAspect(totalW / totalH, canvasEl.width, canvasEl.height); + const imgRatio = Math.min(baseTw / imgW, baseTh / imgH); + const cW = pW * (imgW * imgRatio / totalW), cH = pH * (imgH * imgRatio / totalH); + const rangeX = pW - cW, rangeY = pH - cH; + const newX = rangeX > 0 ? node._padDragStart.x + dx / rangeX : 0.5; + const newY = rangeY > 0 ? node._padDragStart.y + dy / rangeY : 0.5; + setPadXY(newX, newY, true); + drawCanvas(); + return; + } + if (node._dragMode === "extra_pad_drag") { + const m = getCanvasMouse(e); + const dx = m.x - node._dragStart.x; + const dy = m.y - node._dragStart.y; + const s = node._extraPadStart; + const totalH = s.top + s.bottom, totalW = s.left + s.right; + const fullW = (node._previewImg ? node._previewImg.naturalWidth : canvasEl.width) + totalW; + const fullH = (node._previewImg ? node._previewImg.naturalHeight : canvasEl.height) + totalH; + const scale = (fullW / fullH > canvasEl.width / canvasEl.height) + ? fullW / canvasEl.width : fullH / canvasEl.height; + const pixDx = Math.round(dx * scale); + const pixDy = Math.round(dy * scale); + const newLeft = Math.max(0, Math.min(totalW, s.left + pixDx)); + const newTop = Math.max(0, Math.min(totalH, s.top + pixDy)); + for (const [name, val] of [["pad_left", newLeft], ["pad_right", totalW - newLeft], + ["pad_top", newTop], ["pad_bottom", totalH - newTop]]) { + const w = _getSubWidget(`extra_padding.${name}`); + if (w) w.value = val; + } + drawCanvas(); + return; + } + if (node._activeIdx < 0) return; + const m = getCanvasMouse(e); + const cw = canvasEl.width; + const ch = canvasEl.height; + let bbox = node._bboxes[node._activeIdx]; + + if (node._dragMode === "draw") { + bbox.endX = Math.max(0, Math.min(cw, m.x)); + bbox.endY = Math.max(0, Math.min(ch, m.y)); + } else if (node._dragMode === "move" && node._bboxAtDragStart) { + const dx = m.x - node._dragStart.x; + const dy = m.y - node._dragStart.y; + const bs = node._bboxAtDragStart; + const w = bs.endX - bs.startX; + const h = bs.endY - bs.startY; + const nx = Math.max(0, Math.min(cw - w, bs.startX + dx)); + const ny = Math.max(0, Math.min(ch - h, bs.startY + dy)); + bbox = { startX: nx, startY: ny, endX: nx + w, endY: ny + h }; + } else if (node._dragMode?.startsWith("resize") && node._bboxAtDragStart) { + const bs = node._bboxAtDragStart; + const dm = node._dragMode; + // Edge resize: only one axis moves + if (dm === "resize-t" || dm === "resize-b" || dm === "resize-l" || dm === "resize-r") { + const edgeCoord = dm === "resize-t" ? bs.startY : dm === "resize-b" ? bs.endY : dm === "resize-l" ? bs.startX : bs.endX; + const isVert = dm === "resize-t" || dm === "resize-b"; + const off = (isVert ? node._dragStart.y : node._dragStart.x) - edgeCoord; + const val = Math.max(0, Math.min(isVert ? ch : cw, (isVert ? m.y : m.x) - off)); + if (dm === "resize-t") bbox = { startX: bs.startX, startY: val, endX: bs.endX, endY: bs.endY }; + else if (dm === "resize-b") bbox = { startX: bs.startX, startY: bs.startY, endX: bs.endX, endY: val }; + else if (dm === "resize-l") bbox = { startX: val, startY: bs.startY, endX: bs.endX, endY: bs.endY }; + else bbox = { startX: bs.startX, startY: bs.startY, endX: val, endY: bs.endY }; + } else { + // Corner resize + let cornerX, cornerY; + if (dm === "resize-tl") { cornerX = bs.startX; cornerY = bs.startY; } + else if (dm === "resize-tr") { cornerX = bs.endX; cornerY = bs.startY; } + else if (dm === "resize-bl") { cornerX = bs.startX; cornerY = bs.endY; } + else { cornerX = bs.endX; cornerY = bs.endY; } + const offX = node._dragStart.x - cornerX; + const offY = node._dragStart.y - cornerY; + const cx = Math.max(0, Math.min(cw, m.x - offX)); + const cy = Math.max(0, Math.min(ch, m.y - offY)); + if (dm === "resize-tl") bbox = { startX: bs.endX, startY: bs.endY, endX: cx, endY: cy }; + else if (dm === "resize-tr") bbox = { startX: bs.startX, startY: bs.endY, endX: cx, endY: cy }; + else if (dm === "resize-bl") bbox = { startX: bs.endX, startY: bs.startY, endX: cx, endY: cy }; + else if (dm === "resize-br") bbox = { startX: bs.startX, startY: bs.startY, endX: cx, endY: cy }; + } + } + + // Alt: resize symmetrically — for draw, the click point is center; + // for resize, the center of the original bbox is the pivot + if (e.altKey && node._dragMode !== "move") { + let cx, cy; + if (node._dragMode === "draw") { + cx = node._dragStart.x; + cy = node._dragStart.y; + } else if (node._bboxAtDragStart) { + const bs = node._bboxAtDragStart; + cx = (bs.startX + bs.endX) / 2; + cy = (bs.startY + bs.endY) / 2; + } else { + cx = (bbox.startX + bbox.endX) / 2; + cy = (bbox.startY + bbox.endY) / 2; + } + const dx = Math.abs(bbox.endX - cx); + const dy = Math.abs(bbox.endY - cy); + bbox = { + startX: Math.max(0, cx - dx), + startY: Math.max(0, cy - dy), + endX: Math.min(cw, cx + dx), + endY: Math.min(ch, cy + dy), + }; + } + + // Shift: constrain to target aspect ratio (not applicable to edge resize) + const isEdgeResize = /^resize-[tblr]$/.test(node._dragMode); + if (e.shiftKey && node._dragMode !== "move" && !isEdgeResize) { + bbox = constrainAspect(bbox, node._dragMode); + } + + // Ctrl: snap to grid in source image space + if (e.ctrlKey) { + if (node._dragMode === "move") { + bbox = snapBboxPosition(bbox); + } else if (e.altKey) { + // Snap both edges, keep centered + const snap = getEffectiveSnap(); + const scale = getCanvasScale(); + const xEdges = getOtherBboxEdges("x"); + const yEdges = getOtherBboxEdges("y"); + const sx1 = snapCoord(bbox.startX, snap.x, scale.x, cw, xEdges); + const sx2 = snapCoord(bbox.endX, snap.x, scale.x, cw, xEdges); + const sy1 = snapCoord(bbox.startY, snap.y, scale.y, ch, yEdges); + const sy2 = snapCoord(bbox.endY, snap.y, scale.y, ch, yEdges); + bbox = { startX: sx1, startY: sy1, endX: sx2, endY: sy2 }; + } else if (isEdgeResize) { + // Only snap the edge being dragged + const snap = getEffectiveSnap(); + const scale = getCanvasScale(); + const xEdges = getOtherBboxEdges("x"), yEdges = getOtherBboxEdges("y"); + const dm = node._dragMode; + bbox = { ...bbox, + ...(dm === "resize-t" && { startY: snapCoord(bbox.startY, snap.y, scale.y, ch, yEdges) }), + ...(dm === "resize-b" && { endY: snapCoord(bbox.endY, snap.y, scale.y, ch, yEdges) }), + ...(dm === "resize-l" && { startX: snapCoord(bbox.startX, snap.x, scale.x, cw, xEdges) }), + ...(dm === "resize-r" && { endX: snapCoord(bbox.endX, snap.x, scale.x, cw, xEdges) }), + }; + } else { + bbox = snapBbox(bbox); + } + } + + node._bboxes[node._activeIdx] = bbox; + drawCanvas(); + e.preventDefault(); + } + + function onDragEnd(e) { + if (!node._drawing) return; + const wasRotate = node._dragMode === "rotate"; + const wasPadDrag = node._dragMode === "pad_drag"; + Object.assign(node, { _drawing: false, _dragMode: null, _dragStart: null, + _bboxAtDragStart: null, _padDragStart: null, _extraPadStart: null }); + // Flush deferred pad_x/pad_y widget values + if (wasPadDrag) setPadXY(_framePadXY.x, _framePadXY.y); + document.removeEventListener("mousemove", onDragMove); + document.removeEventListener("mouseup", onDragEnd); + + // Snap canvas to exact rotated dimensions on release (during drag we use a threshold to avoid jitter) + if (wasRotate) { + restoreCanvasToImage(); + fitNodeToImage(false); + } + + if (node._activeIdx >= 0) { + const bbox = node._bboxes[node._activeIdx]; + const { x1, y1, x2, y2 } = normBbox(bbox); + if (x2 - x1 < 3 && y2 - y1 < 3) { + // Too small — remove it + removeBbox(node._activeIdx); + } else { + node._bboxes[node._activeIdx] = { startX: x1, startY: y1, endX: x2, endY: y2 }; + } + } + updateBboxWidgets(); + drawCanvas(); + } + + canvasEl.addEventListener("contextmenu", (e) => { e.preventDefault(); e.stopPropagation(); }); + + canvasEl.tabIndex = 0; + canvasEl.style.outline = "none"; + canvasEl.addEventListener("keydown", (e) => { + if (node._activeIdx < 0) return; + const arrows = { ArrowLeft: [-1, 0], ArrowRight: [1, 0], ArrowUp: [0, -1], ArrowDown: [0, 1] }; + const dir = arrows[e.key]; + if (!dir) { + // Delete key removes active bbox + if (e.key === "Delete" || e.key === "Backspace") { + e.preventDefault(); + e.stopPropagation(); + removeBbox(node._activeIdx); + updateBboxWidgets(); + drawCanvas(); + } + // Tab cycles through bboxes + if (e.key === "Tab" && node._bboxes.length > 1) { + e.preventDefault(); + e.stopPropagation(); + node._activeIdx = (node._activeIdx + (e.shiftKey ? -1 : 1) + node._bboxes.length) % node._bboxes.length; + drawCanvas(); + } + return; + } + e.preventDefault(); + e.stopPropagation(); + + const snap = getEffectiveSnap(); + const step = e.ctrlKey ? (snap.x > 1 ? snap.x / getCanvasScale().x : 10) : 1; + const stepY = e.ctrlKey ? (snap.y > 1 ? snap.y / getCanvasScale().y : 10) : 1; + const dx = dir[0] * step; + const dy = dir[1] * stepY; + const bbox = node._bboxes[node._activeIdx]; + const { x1, y1, x2, y2 } = normBbox(bbox); + const w = x2 - x1, h = y2 - y1; + const nx = Math.max(0, Math.min(canvasEl.width - w, x1 + dx)); + const ny = Math.max(0, Math.min(canvasEl.height - h, y1 + dy)); + node._bboxes[node._activeIdx] = { startX: nx, startY: ny, endX: nx + w, endY: ny + h }; + updateBboxWidgets(); + drawCanvas(); + }); + + function updateBboxWidgets() { + const data = {}; + if (node._rotation !== 0) data.rotation = node._rotation; + data.fillColor = getFillColor(); + if (node._bboxes.length > 0) { + data.bboxes = node._bboxes.map((bbox) => { + const b = Object.assign({}, bbox); + if (node._previewImg) { + b.previewWidth = canvasEl.width; + b.previewHeight = canvasEl.height; + } + return b; + }); + } + bboxWidget.value = Object.keys(data).length > 0 ? JSON.stringify(data) : ""; + } + + let _resizing = false; + function fitNodeToCanvas() { + if (canvasEl.width <= 0 || canvasEl.height <= 0) return; + // If in a no-bbox mode with target dims, resize canvas to output AR + const kp = getKp(); + const { w: tgtW, h: tgtH } = getTgt(); + const hasBboxes = node._bboxes.length > 0; + if (!hasBboxes && node._previewImg && node._previewEnabled && node._rotation === 0 && (tgtW > 0 || tgtH > 0) && + (kp === "stretch" || isPadMode(kp))) { + const { w: imgW, h: imgH } = getEffectiveImageDims(kp !== "pad_edge"); + let rw = tgtW > 0 ? tgtW : imgW; + let rh = tgtH > 0 ? tgtH : imgH; + const outAR = rw / rh; + const newH = Math.round(canvasEl.width / outAR); + if (Math.abs(newH - canvasEl.height) > 2) { + canvasEl.height = newH; + canvasEl.style.aspectRatio = `${canvasEl.width} / ${newH}`; + _rotCacheKey = ""; + } + } else if (node._previewImg && node._previewEnabled) { + // Restore to image AR (including pad-first extra padding if active) + let { w: iw, h: ih } = getEffectiveImageDims(true); + const ep = getExtraPadding(); + if (ep.isFirst && (ep.top + ep.bottom + ep.left + ep.right) > 0) { + iw += ep.left + ep.right; + ih += ep.top + ep.bottom; + } + const imgAR = iw / ih; + const newH = Math.round(canvasEl.width / imgAR); + if (Math.abs(newH - canvasEl.height) > 2) { + canvasEl.height = newH; + canvasEl.style.aspectRatio = `${canvasEl.width} / ${newH}`; + _rotCacheKey = ""; + } + } + // Update widget height but don't force node size — only grow if needed + const nodeW = node.size[0]; + const displayedW = nodeW - 30; + const displayedH = Math.round(canvasEl.height * (displayedW / canvasEl.width)); + node._widgetHeight = displayedH + GRID_BAR_HEIGHT; + } + + function setCanvasSize(cw, ch) { + canvasEl.width = cw; + canvasEl.height = ch; + canvasEl.style.aspectRatio = `${cw} / ${ch}`; + node.properties.canvasSize = [cw, ch]; + node._widgetHeight = ch + GRID_BAR_HEIGHT; + if (node.graph) node.graph.setDirtyCanvas(true, true); + fitNodeToCanvas(); + } + + function updateCanvasFromTargetDims() { + if (node._previewImg && node._previewEnabled) return; + const divBy = getDivBy(); + const { w: twV, h: thV } = getTgt(); + let tw = roundDown(twV, divBy); + let th = roundDown(thV, divBy); + const [cw, ch] = clampToMaxCanvas(tw > 0 ? Math.max(tw, 200) : 400, th > 0 ? Math.max(th, 150) : 300); + if (cw !== canvasEl.width || ch !== canvasEl.height) { + setCanvasSize(cw, ch); + } + } + + function getCanvasScale() { + if (node._previewImg && node._previewEnabled) { + if (node._rotation !== 0) { + // When rotated, canvas is sized to rotated bounds — uniform scale + const { w: rotW } = getEffectiveImageDims(true); + const s = rotW / canvasEl.width; // same as rotH / canvasEl.height + return { x: s, y: s }; + } + let nw = node._previewImg.naturalWidth; + let nh = node._previewImg.naturalHeight; + // For "pad first" modes, bbox coords map to padded image dimensions + const ep = getExtraPadding(); + if (ep.isFirst && hasExtraPad(ep)) { + nw += ep.left + ep.right; + nh += ep.top + ep.bottom; + } + return { x: nw / canvasEl.width, y: nh / canvasEl.height }; + } + const { w: tw, h: th } = getTgt(); + return { x: tw > 0 ? tw / canvasEl.width : 1, y: th > 0 ? th / canvasEl.height : 1 }; + } + + function getOutputDims(bbox) { + if (!bbox) return null; + const { x: sX, y: sY } = getCanvasScale(); + const { x1, y1, x2, y2 } = normBbox(bbox); + const cropW = Math.round((x2 - x1) * sX); + const cropH = Math.round((y2 - y1) * sY); + const dims = computeOutputDims(cropW, cropH); + return dims ? { w: Math.max(1, dims.w), h: Math.max(1, dims.h) } : null; + } + + function getGridParams() { + const gridSize = getGridSize(); + if (gridSize <= 1) return null; + + let tw, th; + if (node._previewImg && node._previewEnabled) { + const dims = getEffectiveImageDims(true); + tw = dims.w; + th = dims.h; + } else { + tw = twWidget ? twWidget.value : canvasEl.width; + th = thWidget ? thWidget.value : canvasEl.height; + } + const scaleX = canvasEl.width / tw; + const scaleY = canvasEl.height / th; + if (gridSize * scaleX < 4 && gridSize * scaleY < 4) return null; + // Per-dimension effective grid: no partial cells, cells are near-square + const cellsX = Math.max(1, Math.round(tw / gridSize)); + const cellsY = Math.max(1, Math.round(th / gridSize)); + const effGridW = tw / cellsX; + const effGridH = th / cellsY; + return { gridSize, tw, th, scaleX, scaleY, effGridW, effGridH }; + } + + function drawGrid(ctx, alpha, clipBbox, color) { + const g = getGridParams(); + if (!g) return; + const { tw, th, scaleX, scaleY, effGridW, effGridH } = g; + + ctx.save(); + if (clipBbox) { + const { x1, y1, x2, y2 } = normBbox(clipBbox); + ctx.beginPath(); + ctx.rect(x1, y1, x2 - x1, y2 - y1); + ctx.clip(); + } + + const showPreview = node._previewImg && node._previewEnabled; + const defaultAlpha = showPreview ? 0.25 : 0.1; + if (color) { + ctx.strokeStyle = color; + } else { + const { r, g, b } = getGridColorRGB(); + ctx.strokeStyle = `rgba(${r}, ${g}, ${b}, ${alpha ?? defaultAlpha})`; + } + ctx.lineWidth = 1; + ctx.beginPath(); + for (let rx = effGridW; rx < tw - 0.5; rx += effGridW) { + const px = Math.round(rx * scaleX) + 0.5; + ctx.moveTo(px, 0); + ctx.lineTo(px, canvasEl.height); + } + for (let ry = effGridH; ry < th - 0.5; ry += effGridH) { + const py = Math.round(ry * scaleY) + 0.5; + ctx.moveTo(0, py); + ctx.lineTo(canvasEl.width, py); + } + ctx.stroke(); + ctx.restore(); + } + + function drawRotationCross(ctx) { + if (!node._showRotationCross) return; + const { width: cw, height: ch } = canvasEl; + const armLen = Math.sqrt(cw * cw + ch * ch) / 2; + const rad = rotRad(); + const isActive = node._rotation !== 0; + ctx.save(); + ctx.translate(cw / 2, ch / 2); + ctx.rotate(rad); + ctx.beginPath(); + ctx.moveTo(-armLen, 0); ctx.lineTo(armLen, 0); + ctx.moveTo(0, -armLen); ctx.lineTo(0, armLen); + // Black outline + Object.assign(ctx, { strokeStyle: "rgba(0, 0, 0, 0.6)", lineWidth: isActive ? 5 : 4 }); + ctx.stroke(); + // Orange fill + Object.assign(ctx, { strokeStyle: "rgba(255, 160, 30, 0.8)", lineWidth: isActive ? 3 : 2 }); + ctx.stroke(); + if (isActive) { + ctx.rotate(-rad); + const label = `${node._rotation.toFixed(1)}°`; + Object.assign(ctx, { font: "12px sans-serif", textAlign: "center" }); + const lw = ctx.measureText(label).width + 8; + ctx.fillStyle = "rgba(0, 0, 0, 0.6)"; + ctx.beginPath(); ctx.roundRect(-lw / 2, -34, lw, 16, 3); ctx.fill(); + ctx.fillStyle = "rgba(255, 160, 30, 1)"; + ctx.fillText(label, 0, -20); + } + ctx.restore(); + } + + function drawBbox(ctx, bbox, index, isActive) { + const { x1, y1, x2, y2 } = normBbox(bbox); + const w = x2 - x1, h = y2 - y1; + const color = getBboxColor(index, isActive); + const invertEnabled = icWidget ? icWidget.value === "enabled" : false; + if (isActive) { + if (invertEnabled) { + ctx.fillStyle = getFillColor(); ctx.fillRect(x1, y1, w, h); + } else { + ctx.fillStyle = "rgba(0, 0, 0, 0.5)"; ctx.beginPath(); + ctx.rect(0, 0, canvasEl.width, canvasEl.height); ctx.rect(x1, y1, w, h); ctx.fill("evenodd"); + } + } + drawGrid(ctx, null, bbox, color.gridColor); + if (!isActive) { ctx.fillStyle = color.tint; ctx.fillRect(x1, y1, w, h); } + + if (isActive && !invertEnabled) { + const { x: sX, y: sY } = getCanvasScale(); + const cropW = Math.round(w * sX); + const cropH = Math.round(h * sY); + const { w: tgtW, h: tgtH } = getTgt(); + const kp = getKp(); + + if (kp === "multiplier") { + const wmv = _getSubWidget("keep_proportion.width_mult")?.value || 1.0; + const hmv = _getSubWidget("keep_proportion.height_mult")?.value || 1.0; + if (wmv !== 1.0 || hmv !== 1.0) { + drawStretchOverlay(ctx, (cropW * wmv) / (cropH * hmv), x1, y1, w, h); + } + } + + if ((tgtW > 0 || tgtH > 0) && kp === "stretch") { + const { rw, rh } = getEffectiveTargetDims(cropW, cropH); + drawStretchOverlay(ctx, rw / rh, x1, y1, w, h); + } + + if ((tgtW > 0 || tgtH > 0) && (kp === "crop" || isPadMode(kp))) { + const { rw, rh } = getEffectiveTargetDims(cropW, cropH); + + if (kp === "crop") { + const ratio = Math.max(rw / cropW, rh / cropH); + const visW = rw / ratio, visH = rh / ratio; + const offX = (cropW - visW) / 2, offY = (cropH - visH) / 2; + const vx1 = x1 + offX / sX, vy1 = y1 + offY / sY; + const vx2 = x1 + (offX + visW) / sX, vy2 = y1 + (offY + visH) / sY; + dimSides(ctx, x1, y1, w, h, vx1, vy1, vx2 - vx1, vy2 - vy1); + drawDashedBorder(ctx, vx1, vy1, vx2 - vx1, vy2 - vy1, "rgba(255, 255, 100, 0.8)"); + } else if (isPadMode(kp)) { + const { pW, pH, pX, pY, cW, cH, cX, cY } = computePadLayout(rw, rh, cropW, cropH, w, h, x1, y1); + ctx.fillStyle = "rgba(0, 0, 0, 0.7)"; + ctx.fillRect(x1, y1, w, h); + if (kp === "pad_edge" && node._previewEnabled && node._previewImg) { + renderToTmp(Math.max(1, Math.round(cW)), Math.max(1, Math.round(cH)), + node._previewImg, x1 * sX, y1 * sY, w * sX, h * sY); + drawExtraPadFill(ctx, getEdgeEp(), pX, pY, pW, pH, cX, cY, cW, cH, _tmpCanvas); + } else if (kp === "pad_color") { + ctx.fillStyle = getFillColor(); + ctx.fillRect(pX, pY, pW, pH); + } + if (node._previewEnabled) { + drawPreviewRegion(ctx, x1, y1, w, h, cX, cY, cW, cH); + } + drawDashedBorder(ctx, cX, cY, cW, cH, "rgba(100, 200, 255, 0.8)"); + } + } + + const epBbox = getExtraPadding(); + const padModeHandledBbox = isPadMode(kp) && (tgtW > 0 || tgtH > 0); + if (hasExtraPad(epBbox) && !padModeHandledBbox && !invertEnabled && !epBbox.isFirst && node._previewImg) { + const epContent = computeContentDims(cropW, cropH); + const epCW = epContent ? epContent.w : cropW, epCH = epContent ? epContent.h : cropH; + const totalW2 = epCW + epBbox.left + epBbox.right, totalH2 = epCH + epBbox.top + epBbox.bottom; + const { pW: pW2, pH: pH2, pX: pX2, pY: pY2 } = fitAspect(totalW2 / totalH2, w, h, x1, y1); + ctx.fillStyle = "rgba(0, 0, 0, 0.7)"; ctx.fillRect(x1, y1, w, h); + const cW2 = pW2 * (epCW / totalW2), cH2 = pH2 * (epCH / totalH2); + const cX2 = pX2 + pW2 * (epBbox.left / totalW2), cY2 = pY2 + pH2 * (epBbox.top / totalH2); + // Draw crop fitted within content area (letterbox, don't stretch) + const cropAR = cropW / cropH; + const { pW: iW2, pH: iH2, pX: iX2, pY: iY2 } = fitAspect(cropAR, cW2, cH2, cX2, cY2); + const riW2 = Math.max(1, Math.round(iW2)), riH2 = Math.max(1, Math.round(iH2)); + renderToTmp(riW2, riH2, node._previewImg, x1 * sX, y1 * sY, w * sX, h * sY); + drawExtraPadFill(ctx, epBbox, pX2, pY2, pW2, pH2, cX2, cY2, cW2, cH2, _tmpCanvas); + drawPreviewRegion(ctx, x1, y1, w, h, iX2, iY2, iW2, iH2); + drawDashedBorder(ctx, cX2, cY2, cW2, cH2, "rgba(100, 200, 255, 0.8)"); + } + + const { w: outW, h: outH } = getOutputDims(bbox) || { w: cropW, h: cropH }; + const boxLabel = node._bboxes.length > 1 ? `[${index + 1}] ` : ""; + const label = (tgtW > 0 || tgtH > 0) + ? `${boxLabel}${cropW}×${cropH} → ${outW}×${outH}` + : `${boxLabel}${cropW} × ${cropH}`; + ctx.font = "12px sans-serif"; + const lw = ctx.measureText(label).width + 12; + const lh = 20; + const ly = y1 > lh + 4 ? y1 - lh - 2 : y1 + 4; + ctx.fillStyle = "rgba(0, 0, 0, 0.7)"; + ctx.fillRect(x1, ly, lw, lh); + ctx.fillStyle = "#fff"; + ctx.textAlign = "left"; + ctx.fillText(label, x1 + 6, ly + 14); + } else if (!isActive) { + // Inactive: show index label centered in bbox + const boxLabel = `${index + 1}`; + ctx.font = "bold 20px sans-serif"; + ctx.fillStyle = color.border; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText(boxLabel, x1 + w / 2, y1 + h / 2); + ctx.textBaseline = "alphabetic"; + } + + if (isActive) { + const cx = (x1 + x2) / 2, cy = (y1 + y2) / 2; + Object.assign(ctx, { strokeStyle: color.border, lineWidth: 1 }); + ctx.beginPath(); + ctx.moveTo(cx - 5, cy); ctx.lineTo(cx + 5, cy); + ctx.moveTo(cx, cy - 5); ctx.lineTo(cx, cy + 5); + ctx.stroke(); + } + if (!isActive) { ctx.strokeStyle = "rgba(0, 0, 0, 0.6)"; ctx.lineWidth = 4; ctx.strokeRect(x1, y1, w, h); } + ctx.strokeStyle = color.border; ctx.lineWidth = 2; ctx.strokeRect(x1, y1, w, h); + + ctx.fillStyle = color.handle; + const hs = isActive ? 5 : 4, s2 = hs * 2; + const cx1 = Math.max(hs, Math.min(canvasEl.width - hs, x1)); + const cx2 = Math.max(hs, Math.min(canvasEl.width - hs, x2)); + const cy1 = Math.max(hs, Math.min(canvasEl.height - hs, y1)); + const cy2 = Math.max(hs, Math.min(canvasEl.height - hs, y2)); + for (const [hx, hy] of [[cx1, cy1], [cx2, cy1], [cx1, cy2], [cx2, cy2]]) + ctx.fillRect(hx - hs, hy - hs, s2, s2); + } + + function drawCanvas() { + const ctx = canvasCtx; + ctx.clearRect(0, 0, canvasEl.width, canvasEl.height); + _frameEp = _readExtraPadding(); + if (node._dragMode !== "pad_drag") _framePadXY = _readPadXY(); + const _kpVal = getKp(), { w: _tgtW, h: _tgtH } = getTgt(); + const hasBboxes = node._bboxes.length > 0; + const showPreview = node._previewImg && node._previewEnabled; + const _skipFullImage = !hasBboxes && showPreview && isPadMode(_kpVal) && (_tgtW > 0 || _tgtH > 0); + + const _epBase = getExtraPadding(); + const _isFirst = _epBase.isFirst, _hasFirstPad = _isFirst && hasExtraPad(_epBase); + const maybeResize = (cw, ch) => { + if (cw !== canvasEl.width || ch !== canvasEl.height) { setCanvasSize(cw, ch); ctx.clearRect(0, 0, cw, ch); } + }; + if (showPreview) { + if (_hasFirstPad) { + const { w: eiw, h: eih } = getEffectiveImageDims(_epBase.mode !== "pad_edge"); + maybeResize(...clampToMaxCanvas(eiw + _epBase.left + _epBase.right, eih + _epBase.top + _epBase.bottom)); + } else if (node._previewImg) { + maybeResize(...clampToMaxCanvas(...Object.values(getEffectiveImageDims(true)))); + } + } + + if (showPreview && !_skipFullImage) { + if (_hasFirstPad) { + const { w: epImgW, h: epImgH } = getEffectiveImageDims(_epBase.mode !== "pad_edge"); + const totalW = epImgW + _epBase.left + _epBase.right, totalH = epImgH + _epBase.top + _epBase.bottom; + const cX = canvasEl.width * (_epBase.left / totalW), cY = canvasEl.height * (_epBase.top / totalH); + const cW = canvasEl.width * (epImgW / totalW), cH = canvasEl.height * (epImgH / totalH); + const rcW = Math.max(1, Math.round(cW)), rcH = Math.max(1, Math.round(cH)); + // Render content (with replicate-rotate for pad_edge + rotation) + renderToTmp(rcW, rcH); + drawExtraPadFill(ctx, _epBase, 0, 0, canvasEl.width, canvasEl.height, cX, cY, cW, cH, _tmpCanvas); + drawPreviewRegion(ctx, 0, 0, canvasEl.width, canvasEl.height, cX, cY, cW, cH); + } else { + if (node._rotation !== 0) { ctx.fillStyle = getFillColor(); ctx.fillRect(0, 0, canvasEl.width, canvasEl.height); } + const { pW, pH, pX, pY } = fitAspect(_rotCanvas.width / _rotCanvas.height, canvasEl.width, canvasEl.height); + ctx.drawImage(getRotatedPreview(), 0, 0, _rotCanvas.width, _rotCanvas.height, pX, pY, pW, pH); + } + } else if (!showPreview) { + ctx.fillStyle = "#383838"; ctx.fillRect(0, 0, canvasEl.width, canvasEl.height); + const divBy = getDivBy(), { w: lw0, h: lh0 } = getTgt(); + Object.assign(ctx, { fillStyle: "#888", font: "14px sans-serif", textAlign: "center" }); + ctx.fillText(`${roundDown(lw0, divBy) || canvasEl.width} × ${roundDown(lh0, divBy) || canvasEl.height}`, + canvasEl.width / 2, canvasEl.height / 2); + } + + // Manual mode: visualize extra padding on the placeholder + if (!hasBboxes && !showPreview) { + const _ep = getExtraPadding(); + if (hasExtraPad(_ep)) { + const cW0 = canvasEl.width, cH0 = canvasEl.height; + const { w: tw0, h: th0 } = getTgt(); + const rawW = tw0 > 0 ? tw0 : cW0; + const rawH = th0 > 0 ? th0 : cH0; + const content = computeContentDims(rawW, rawH); + const contentW = content ? content.w : rawW; + const contentH = content ? content.h : rawH; + const totalW = contentW + _ep.left + _ep.right; + const totalH = contentH + _ep.top + _ep.bottom; + const { pW, pH, pX, pY } = fitAspect(totalW / totalH, cW0, cH0); + const cW = pW * (contentW / totalW), cH = pH * (contentH / totalH); + const cX = pX + pW * (_ep.left / totalW), cY = pY + pH * (_ep.top / totalH); + ctx.fillStyle = "rgba(0, 0, 0, 0.5)"; + ctx.fillRect(0, 0, cW0, cH0); + if (_ep.mode === "pad_color" || _ep.mode === "pad_crop_color") { + ctx.fillStyle = getFillColor(); + } else { + ctx.fillStyle = "#555"; + } + ctx.fillRect(pX, pY, pW, pH); + ctx.fillStyle = "#383838"; + ctx.fillRect(cX, cY, cW, cH); + drawDashedBorder(ctx, cX, cY, cW, cH, "rgba(100, 200, 255, 0.8)"); + // Redraw dimension text on top + const divBy = getDivBy(), { w: lw0, h: lh0 } = getTgt(); + Object.assign(ctx, { fillStyle: "#888", font: "14px sans-serif", textAlign: "center" }); + ctx.fillText(`${roundDown(lw0, divBy) || cW0} × ${roundDown(lh0, divBy) || cH0}`, cX + cW / 2, cY + cH / 2); + } + } + + // When no bboxes are drawn, show pad/crop preview for the full image (auto mode only) + if (!hasBboxes && node._previewEnabled) { + const cW0 = canvasEl.width, cH0 = canvasEl.height; + if (node._previewImg) { + const { w: tgtW, h: tgtH } = getTgt(); + const kp = getKp(); + // Multiplier mode: compute effective target from crop dims × multiplier + if (kp === "multiplier") { + const wmv = _getSubWidget("keep_proportion.width_mult")?.value || 1.0; + const hmv = _getSubWidget("keep_proportion.height_mult")?.value || 1.0; + if (wmv !== 1.0 || hmv !== 1.0) { + const nw = node._previewImg.naturalWidth; + const nh = node._previewImg.naturalHeight; + drawStretchOverlay(ctx, (nw * wmv) / (nh * hmv), 0, 0, cW0, cH0); + } + } + if ((tgtW > 0 || tgtH > 0) && (kp === "stretch" || isPadMode(kp) || kp === "crop")) { + const { w: imgW, h: imgH } = getEffectiveImageDims(kp !== "pad_edge"); + const { rw, rh } = getEffectiveTargetDims(imgW, imgH); + + if (kp === "stretch") { + drawStretchOverlay(ctx, rw / rh, 0, 0, cW0, cH0); + } else if (isPadMode(kp)) { + const epPad = getExtraPadding(); + const hasEp = hasExtraPad(epPad); + // Full output = target. Content = target - extra padding. + const fullTgtW = tgtW > 0 ? tgtW : imgW; + const fullTgtH = tgtH > 0 ? tgtH : imgH; + const contentTgtW = hasEp ? Math.max(1, fullTgtW - epPad.left - epPad.right) : fullTgtW; + const contentTgtH = hasEp ? Math.max(1, fullTgtH - epPad.top - epPad.bottom) : fullTgtH; + // Fit the full output to the canvas + const { pW: outerW, pH: outerH, pX: outerX, pY: outerY } = fitAspect(fullTgtW / fullTgtH, cW0, cH0); + // Content area within the output (offset by padding) + const contentW = outerW * (contentTgtW / fullTgtW); + const contentH = outerH * (contentTgtH / fullTgtH); + const contentX = outerX + outerW * (epPad.left / fullTgtW); + const contentY = outerY + outerH * (epPad.top / fullTgtH); + // Image is sized to fit the content area, but positioned across the full output via pad_x/pad_y + const imgRatio = Math.min(contentTgtW / imgW, contentTgtH / imgH); + const scaledW = imgW * imgRatio, scaledH = imgH * imgRatio; + const pW = outerW, pH = outerH, pX = outerX, pY = outerY; + const cW = outerW * (scaledW / fullTgtW), cH = outerH * (scaledH / fullTgtH); + const { cX, cY } = computePadContentPos(outerX, outerY, outerW, outerH, cW, cH); + ctx.clearRect(0, 0, cW0, cH0); + // Fill outer area (extra padding region) + if (hasEp) { + ctx.fillStyle = getFillColor(); + ctx.fillRect(outerX, outerY, outerW, outerH); + } + if (kp === "pad_edge") { + renderToTmp(Math.max(1, Math.round(cW)), Math.max(1, Math.round(cH)), + null, 0, 0, 0, 0, true); + ctx.drawImage(_tmpCanvas, 0, 0, _tmpCanvas.width, _tmpCanvas.height, cX, cY, cW, cH); + drawExtraPadFill(ctx, getEdgeEp(), pX, pY, pW, pH, cX, cY, cW, cH, _tmpCanvas); + } else { + ctx.fillStyle = getFillColor(); + ctx.fillRect(pX, pY, pW, pH); + drawPreviewRegion(ctx, 0, 0, cW0, cH0, cX, cY, cW, cH); + } + drawDashedBorder(ctx, cX, cY, cW, cH, "rgba(100, 200, 255, 0.8)"); + if (hasEp) drawDashedBorder(ctx, contentX, contentY, contentW, contentH, "rgba(255, 200, 100, 0.6)"); + } else if (kp === "crop") { + const ratio = Math.max(rw / imgW, rh / imgH); + const visW = rw / ratio, visH = rh / ratio; + const offX = (imgW - visW) / 2, offY = (imgH - visH) / 2; + const sc = Math.min(cW0 / imgW, cH0 / imgH); + const ox = (cW0 - imgW * sc) / 2, oy = (cH0 - imgH * sc) / 2; + const vx1 = ox + offX * sc, vy1 = oy + offY * sc; + const vw = visW * sc, vh = visH * sc; + dimSides(ctx, 0, 0, cW0, cH0, vx1, vy1, vw, vh); + drawDashedBorder(ctx, vx1, vy1, vw, vh, "rgba(255, 255, 100, 0.8)"); + } + } + } + const _ep = getExtraPadding(); + if (hasExtraPad(_ep) && node._previewImg && !_isFirst && !(isPadMode(_kpVal) && (_tgtW > 0 || _tgtH > 0))) { + const { w: rawImgW, h: rawImgH } = getEffectiveImageDims(_ep.mode !== "pad_edge"); + const content = computeContentDims(rawImgW, rawImgH); + const epImgW = content ? content.w : rawImgW; + const epImgH = content ? content.h : rawImgH; + const totalW = epImgW + _ep.left + _ep.right, totalH = epImgH + _ep.top + _ep.bottom; + const { pW, pH, pX, pY } = fitAspect(totalW / totalH, cW0, cH0); + ctx.clearRect(0, 0, cW0, cH0); + const cW = pW * (epImgW / totalW), cH = pH * (epImgH / totalH); + const cX = pX + pW * (_ep.left / totalW), cY = pY + pH * (_ep.top / totalH); + // Draw padding fill, then image in content area (same approach as keep_proportion pad mode) + if (_ep.mode === "pad_color") { + ctx.fillStyle = getFillColor(); + ctx.fillRect(pX, pY, pW, pH); + drawPreviewRegion(ctx, 0, 0, cW0, cH0, cX, cY, cW, cH); + } else { + // pad_edge: render content for edge fill source + const rcW2 = Math.max(1, Math.round(cW)), rcH2 = Math.max(1, Math.round(cH)); + renderToTmp(rcW2, rcH2, null, 0, 0, 0, 0, _ep.isFirst); + drawExtraPadFill(ctx, _ep, pX, pY, pW, pH, cX, cY, cW, cH, _tmpCanvas); + drawPreviewRegion(ctx, 0, 0, cW0, cH0, cX, cY, cW, cH); + } + drawDashedBorder(ctx, cX, cY, cW, cH, "rgba(100, 200, 255, 0.8)"); + } + + drawGrid(ctx); + drawRotationCross(ctx); + if (!node._drawing) updateOutputSlotLabels(); + return; + } + + if (node._activeIdx >= 0 && node._activeIdx < node._bboxes.length) { + drawBbox(ctx, node._bboxes[node._activeIdx], node._activeIdx, true); + } + for (let i = 0; i < node._bboxes.length; i++) { + if (i === node._activeIdx) continue; + drawBbox(ctx, node._bboxes[i], i, false); + } + + drawGrid(ctx); + drawRotationCross(ctx); + if (!node._drawing) updateOutputSlotLabels(); + } + + let _lastSlotLabelKey = ""; + function updateOutputSlotLabels() { + const activeBbox = (node._activeIdx >= 0 && node._activeIdx < node._bboxes.length) + ? node._bboxes[node._activeIdx] : null; + let dims = activeBbox ? getOutputDims(activeBbox) : null; + if (!dims) { + const { w: tW, h: tH } = getTgt(); + const imgW = node._previewImg ? node._previewImg.naturalWidth + : (tW > 0 ? tW : canvasEl.width); + const imgH = node._previewImg ? node._previewImg.naturalHeight + : (tH > 0 ? tH : canvasEl.height); + const computed = computeOutputDims(imgW, imgH); + if (computed && computed.w > 0 && computed.h > 0) { + dims = computed; + } + } + // Skip expensive DOM/canvas updates if nothing changed + const key = `${dims?.w || 0},${dims?.h || 0},${node._bboxes.length}`; + if (key === _lastSlotLabelKey) return; + _lastSlotLabelKey = key; + if (node.outputs) { + const hasBboxes = node._bboxes.length > 0; + for (let i = 0; i < node.outputs.length; i++) { + const out = node.outputs[i]; + if (!out._origName) { + if (out.name?.includes("width") || out.label?.includes("width")) out._origName = "width"; + else if (out.name?.includes("height") || out.label?.includes("height")) out._origName = "height"; + else if (out.name === "bbox") out._origName = "bbox"; + else if (out.name === "bbox_mask") out._origName = "bbox_mask"; + } + if (out._origName === "width") { + out.label = dims ? `${dims.w} width` : "width"; + } else if (out._origName === "height") { + out.label = dims ? `${dims.h} height` : "height"; + } + // Identify output/output_mask slots + if (!out._origName) { + if (out.name === "cropped" || out.name === "output") out._origName = "output"; + else if (out.name === "cropped_mask" || out.name === "output_mask") out._origName = "output_mask"; + } + // Toggle list icon on output/output_mask + if (out._origName === "output" || out._origName === "output_mask") { + out.shape = hasBboxes && node._bboxes.length > 1 ? 6 : undefined; + } + // Visually hide bbox/bbox_mask slots when no bboxes, show with list icon when multiple + if (out._origName === "bbox" || out._origName === "bbox_mask") { + if (!hasBboxes) { + out.label = " "; + out.color_off = "transparent"; + out.color_on = "transparent"; + out.shape = 6; + out.dir = undefined; + } else { + out.label = out._origName; + out.color_off = undefined; + out.color_on = undefined; + out.shape = node._bboxes.length > 1 ? 6 : undefined; + } + } + } + node.setDirtyCanvas(true, true); + // Vue nodes mode: hide/show bbox slots via DOM since canvas properties don't apply + if (LiteGraph.vueNodesMode && node.id != null) { + const nodeEl = document.querySelector(`[data-node-id="${node.id}"]`); + if (nodeEl) { + const outputSlots = nodeEl.querySelectorAll(".lg-slot--output"); + for (let i = 0; i < node.outputs.length; i++) { + const out = node.outputs[i]; + if ((out._origName === "bbox" || out._origName === "bbox_mask") && outputSlots[i]) { + outputSlots[i].style.visibility = hasBboxes ? "" : "hidden"; + // Update label text directly since Vue may not observe raw property changes + if (hasBboxes) { + const labelEl = outputSlots[i].querySelector("span"); + if (labelEl) labelEl.textContent = out._origName; + } + } + } + } + } + } + } + + function loadPreviewImage(src, preserveBbox) { + const img = new Image(); + img.crossOrigin = "anonymous"; + img.onload = () => { + node._previewImg = img; + _rotCacheKey = ""; + invalidateGridStops(); + restoreCanvasToImage(); + if (node._previewEnabled && !preserveBbox) { + node._bboxes = []; + node._activeIdx = -1; + } + updateBboxWidgets(); + drawCanvas(); + fitNodeToImage(preserveBbox); + }; + img.onerror = () => { + console.warn("ImageTransformKJ: failed to load preview:", src); + }; + img.src = src; + } + + function loadPreviewFromTempFile(filename) { + if (!filename) return; + loadPreviewImage(`/view?filename=${encodeURIComponent(filename)}&type=temp`, true); + } + + function captureVideoFrameToPreview(videoEl, preserveBbox) { + captureVideoFrame(videoEl, (canvas) => { + loadPreviewImage(canvas.toDataURL("image/webp", 0.9), preserveBbox); + }); + } + + function handleSourceChange(sources, preserveBbox) { + if (sources.length === 0) { + // Disconnected + node._previewImg = null; + delete node.properties.previewFile; + if (node._previewEnabled) { + node._bboxes = []; + node._activeIdx = -1; + updateBboxWidgets(); + } + updateCanvasFromTargetDims(); + drawCanvas(); + return; + } + const src = sources[0]; + if (src.isVideo && src.videoEl) { + captureVideoFrameToPreview(src.videoEl, preserveBbox); + } else if (src.url) { + loadPreviewImage(src.url, preserveBbox); + } + } + + let _configuring = false, _sourceResolved = false; + watchImageInputs(node, "image", (sources) => { + _sourceResolved = true; + handleSourceChange(sources, _configuring); + }); + + chainCallback(this, "onResize", function () { + if (_resizing) return; + const availW = Math.max(100, node.size[0] - 30); + _resizing = true; + if (node._previewEnabled) { + // Auto mode: resize canvas pixel buffer to match node width + const ar = canvasEl.height / canvasEl.width; + const newW = Math.round(availW); + const newH = Math.round(availW * ar); + if (newW !== canvasEl.width || newH !== canvasEl.height) { + canvasEl.width = newW; + canvasEl.height = newH; + canvasEl.style.aspectRatio = `${newW} / ${newH}`; + node.properties.canvasSize = [newW, newH]; + } + node._widgetHeight = newH + GRID_BAR_HEIGHT; + } else { + // Manual mode: keep pixel buffer fixed, CSS handles display scaling + const displayedH = Math.round(canvasEl.height * (availW / canvasEl.width)); + node._widgetHeight = displayedH + GRID_BAR_HEIGHT; + } + drawCanvas(); + _resizing = false; + }); + + chainCallback(this, "onConfigure", function () { + node._cropConfigured = true; + if (node.properties.gridSize !== undefined) { + gridSlider.value = node.properties.gridSize; + updateGridLabel(); + } + if (node.properties.previewEnabled !== undefined) { + node._previewEnabled = node.properties.previewEnabled; + updatePreviewBtn(); + } + restoreBboxesFromWidget(); + if (node.properties.rotation !== undefined) node._rotation = node.properties.rotation; + if (node.properties.fillColor) { + colorInput.value = node.properties.fillColor; + colorSwatch.style.background = node.properties.fillColor; + } + if (node.properties.gridColor) { + gridColorInput.value = node.properties.gridColor; + gridColorSwatch.style.background = node.properties.gridColor; + } + + if (node.properties.canvasSize) { + setCanvasSize(node.properties.canvasSize[0], node.properties.canvasSize[1]); + } + + _configuring = true; + _sourceResolved = false; + setTimeout(() => { + if (node.properties.gridSize !== undefined) { + gridSlider.value = node.properties.gridSize; + updateGridLabel(); + } + _configuring = false; + if (!_sourceResolved) { + if (node.properties.previewFile) { + loadPreviewFromTempFile(node.properties.previewFile); + } else { + updateCanvasFromTargetDims(); + drawCanvas(); + } + } + }, 200); + }); + + chainCallback(this, "onExecuted", function (message) { + const filename = message?.preview_filename?.[0]; + if (filename) { + node.properties.previewFile = filename; + loadPreviewFromTempFile(filename); + } + }); + + // Initial draw + setCanvasSize(400, 300); + drawCanvas(); + }); + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/jsnodes.js b/custom_nodes/ComfyUI-KJNodes/web/js/jsnodes.js new file mode 100644 index 0000000000000000000000000000000000000000..a38556a37824c7ce6b70c31c26f6e14fffafba4e --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/jsnodes.js @@ -0,0 +1,365 @@ +const { app } = window.comfyAPI.app; +const { applyTextReplacements } = window.comfyAPI.utils; + +// ─── Dynamic input slot management for *Multi nodes ─── +// Adds an "Update inputs" button that rebuilds prefix_1..prefix_N slots to match the count widget. +// Also wired to the count callback so API-format reload (bare callback) restores the slots; the +// canvas arg from interactive edits is the tell, so scrubbing the count doesn't reflow the node. +function setupDynamicInputs(node, { type, prefix, countWidget = "inputcount", slotOptions } = {}) { + const rebuild = () => { + if (!node.inputs) node.inputs = []; + const countW = node.widgets?.find(w => w.name === countWidget); + if (!countW) return; + const target = countW.value; + const current = node.inputs.filter(i => i.name?.startsWith(prefix)).length; + if (target === current) return; + if (target < current) { + for (let i = 0; i < current - target; i++) node.removeInput(node.inputs.length - 1); + } else { + for (let i = current + 1; i <= target; i++) node.addInput(`${prefix}${i}`, type, slotOptions); + } + }; + node.addWidget("button", "Update inputs", null, rebuild); + const countW = node.widgets?.find(w => w.name === countWidget); + if (countW) { + const origCb = countW.callback; // guard: may be nullish + countW.callback = function (value, canvas) { + const r = origCb ? origCb.apply(this, arguments) : undefined; + if (!canvas) rebuild(); // bare = API reload; skip interactive scrub + return r; + }; + } + return rebuild; +} + +app.registerExtension({ + name: "KJNodes.jsnodes", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if(!nodeData?.category?.startsWith("KJNodes")) { + return; + } + switch (nodeData.name) { + case "ConditioningMultiCombine": + nodeType.prototype.onNodeCreated = function () { + this.inputs_offset = nodeData.name.includes("selective")?1:0 + setupDynamicInputs(this, { type: "CONDITIONING", prefix: "conditioning_" }); + } + break; + case "ImageBatchMulti": + case "ImageAddMulti": + case "CrossFadeImagesMulti": + case "TransitionImagesMulti": + nodeType.prototype.onNodeCreated = function () { + setupDynamicInputs(this, { type: "IMAGE", prefix: "image_", slotOptions: {shape: 7} }); + } + break; + case "ImageConcatMulti": + // Dynamic slots accept MASK too; name-prefix counting handles the mixed types. + nodeType.prototype.onNodeCreated = function () { + setupDynamicInputs(this, { type: "IMAGE,MASK", prefix: "image_", slotOptions: {shape: 7} }); + } + break; + case "MaskBatchMulti": + nodeType.prototype.onNodeCreated = function () { + setupDynamicInputs(this, { type: "MASK", prefix: "mask_" }); + } + break; + + case "FluxBlockLoraSelect": + case "HunyuanVideoBlockLoraSelect": + case "Wan21BlockLoraSelect": + case "LTX2BlockLoraSelect": + nodeType.prototype.onNodeCreated = function () { + this.addWidget("button", "Set all", null, () => { + const userInput = prompt("Enter the values to set for widgets (e.g., s0,1,2-7=2.0, d0,1,2-7=2.0, or 1.0):", ""); + if (userInput) { + const regex = /([sd])?(\d+(?:,\d+|-?\d+)*?)?=(\d+(\.\d+)?)/; + const match = userInput.match(regex); + if (match) { + const type = match[1]; + const indicesPart = match[2]; + const value = parseFloat(match[3]); + + let targetWidgets = []; + if (type === 's') { + targetWidgets = this.widgets.filter(widget => widget.name.includes("single")); + } else if (type === 'd') { + targetWidgets = this.widgets.filter(widget => widget.name.includes("double")); + } else { + targetWidgets = this.widgets; // No type specified, all widgets + } + + if (indicesPart) { + const indices = indicesPart.split(',').flatMap(part => { + if (part.includes('-')) { + const [start, end] = part.split('-').map(Number); + return Array.from({ length: end - start + 1 }, (_, i) => start + i); + } + return Number(part); + }); + + for (const index of indices) { + if (index < targetWidgets.length) { + targetWidgets[index].value = value; + } + } + } else { + // No indices provided, set value for all target widgets + for (const widget of targetWidgets) { + widget.value = value; + } + } + } else if (!isNaN(parseFloat(userInput))) { + // Single value provided, set it for all widgets + const value = parseFloat(userInput); + for (const widget of this.widgets) { + widget.value = value; + } + } else { + alert("Invalid input format. Please use the format s0,1,2-7=2.0, d0,1,2-7=2.0, or 1.0"); + } + } else { + alert("Invalid input. Please enter a value."); + } + }); + }; + break; + + case "GetMaskSizeAndCount": + const onGetMaskSizeConnectInput = nodeType.prototype.onConnectInput; + nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { + const v = onGetMaskSizeConnectInput? onGetMaskSizeConnectInput.apply(this, arguments): undefined + this.outputs[1]["label"] = "width" + this.outputs[2]["label"] = "height" + this.outputs[3]["label"] = "count" + return v; + } + const onGetMaskSizeExecuted = nodeType.prototype.onAfterExecuteNode; + nodeType.prototype.onExecuted = function(message) { + const r = onGetMaskSizeExecuted? onGetMaskSizeExecuted.apply(this,arguments): undefined + let values = message["text"].toString().split('x').map(Number); + this.outputs[1]["label"] = values[1] + " width" + this.outputs[2]["label"] = values[2] + " height" + this.outputs[3]["label"] = values[0] + " count" + return r + } + break; + + case "GetImageSizeAndCount": + const onGetImageSizeConnectInput = nodeType.prototype.onConnectInput; + nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { + console.log(this) + const v = onGetImageSizeConnectInput? onGetImageSizeConnectInput.apply(this, arguments): undefined + //console.log(this) + this.outputs[1]["label"] = "width" + this.outputs[2]["label"] = "height" + this.outputs[3]["label"] = "count" + return v; + } + //const onGetImageSizeExecuted = nodeType.prototype.onExecuted; + const onGetImageSizeExecuted = nodeType.prototype.onAfterExecuteNode; + nodeType.prototype.onExecuted = function(message) { + console.log(this) + const r = onGetImageSizeExecuted? onGetImageSizeExecuted.apply(this,arguments): undefined + let values = message["text"].toString().split('x').map(Number); + console.log(values) + this.outputs[1]["label"] = values[1] + " width" + this.outputs[2]["label"] = values[2] + " height" + this.outputs[3]["label"] = values[0] + " count" + return r + } + break; + + case "GetLatentSizeAndCount": + const onGetLatentConnectInput = nodeType.prototype.onConnectInput; + nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { + console.log(this) + const v = onGetLatentConnectInput? onGetLatentConnectInput.apply(this, arguments): undefined + //console.log(this) + this.outputs[1]["label"] = "batch_size" + this.outputs[2]["label"] = "channels" + this.outputs[3]["label"] = "frames" + this.outputs[4]["label"] = "height" + this.outputs[5]["label"] = "width" + return v; + } + //const onGetImageSizeExecuted = nodeType.prototype.onExecuted; + const onGetLatentSizeExecuted = nodeType.prototype.onAfterExecuteNode; + nodeType.prototype.onExecuted = function(message) { + console.log(this) + const r = onGetLatentSizeExecuted? onGetLatentSizeExecuted.apply(this,arguments): undefined + let values = message["text"].toString().split('x').map(Number); + console.log(values) + this.outputs[1]["label"] = values[0] + " batch" + this.outputs[2]["label"] = values[1] + " channels" + this.outputs[3]["label"] = values[2] + " frames" + this.outputs[4]["label"] = values[3] + " height" + this.outputs[5]["label"] = values[4] + " width" + return r + } + break; + + case "PreviewAnimation": + const onPreviewAnimationConnectInput = nodeType.prototype.onConnectInput; + nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { + const v = onPreviewAnimationConnectInput? onPreviewAnimationConnectInput.apply(this, arguments): undefined + this.title = "Preview Animation" + return v; + } + const onPreviewAnimationExecuted = nodeType.prototype.onAfterExecuteNode; + nodeType.prototype.onExecuted = function(message) { + const r = onPreviewAnimationExecuted? onPreviewAnimationExecuted.apply(this,arguments): undefined + let values = message["text"].toString(); + this.title = "Preview Animation " + values + return r + } + break; + + case "VRAM_Debug": + const onVRAM_DebugConnectInput = nodeType.prototype.onConnectInput; + nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) { + const v = onVRAM_DebugConnectInput? onVRAM_DebugConnectInput.apply(this, arguments): undefined + this.outputs[3]["label"] = "freemem_before" + this.outputs[4]["label"] = "freemem_after" + return v; + } + const onVRAM_DebugExecuted = nodeType.prototype.onAfterExecuteNode; + nodeType.prototype.onExecuted = function(message) { + const r = onVRAM_DebugExecuted? onVRAM_DebugExecuted.apply(this,arguments): undefined + let values = message["text"].toString().split('x'); + this.outputs[3]["label"] = values[0] + " freemem_before" + this.outputs[4]["label"] = values[1] + " freemem_after" + return r + } + break; + + case "JoinStringMulti": + const originalOnNodeCreated = nodeType.prototype.onNodeCreated || function() {}; + nodeType.prototype.onNodeCreated = function () { + originalOnNodeCreated.apply(this, arguments); + setupDynamicInputs(this, { type: "STRING", prefix: "string_", slotOptions: {shape: 7} }); + } + break; + case "SoundReactive": + nodeType.prototype.onNodeCreated = function () { + let audioContext; + let microphoneStream; + let animationFrameId; + let analyser; + let dataArray; + let startRangeHz; + let endRangeHz; + let smoothingFactor = 0.5; + let smoothedSoundLevel = 0; + + // Function to update the widget value in real-time + const updateWidgetValueInRealTime = () => { + // Ensure analyser and dataArray are defined before using them + if (analyser && dataArray) { + analyser.getByteFrequencyData(dataArray); + + const startRangeHzWidget = this.widgets.find(w => w.name === "start_range_hz"); + if (startRangeHzWidget) startRangeHz = startRangeHzWidget.value; + const endRangeHzWidget = this.widgets.find(w => w.name === "end_range_hz"); + if (endRangeHzWidget) endRangeHz = endRangeHzWidget.value; + const smoothingFactorWidget = this.widgets.find(w => w.name === "smoothing_factor"); + if (smoothingFactorWidget) smoothingFactor = smoothingFactorWidget.value; + + // Calculate frequency bin width (frequency resolution) + const frequencyBinWidth = audioContext.sampleRate / analyser.fftSize; + // Convert the widget values from Hz to indices + const startRangeIndex = Math.floor(startRangeHz / frequencyBinWidth); + const endRangeIndex = Math.floor(endRangeHz / frequencyBinWidth); + + // Function to calculate the average value for a frequency range + const calculateAverage = (start, end) => { + const sum = dataArray.slice(start, end).reduce((acc, val) => acc + val, 0); + const average = sum / (end - start); + + // Apply exponential moving average smoothing + smoothedSoundLevel = (average * (1 - smoothingFactor)) + (smoothedSoundLevel * smoothingFactor); + return smoothedSoundLevel; + }; + // Calculate the average levels for each frequency range + const soundLevel = calculateAverage(startRangeIndex, endRangeIndex); + + // Update the widget values + + const lowLevelWidget = this.widgets.find(w => w.name === "sound_level"); + if (lowLevelWidget) lowLevelWidget.value = soundLevel; + + animationFrameId = requestAnimationFrame(updateWidgetValueInRealTime); + } + }; + + // Function to start capturing audio from the microphone + const startMicrophoneCapture = () => { + // Only create the audio context and analyser once + if (!audioContext) { + audioContext = new (window.AudioContext || window.webkitAudioContext)(); + // Access the sample rate of the audio context + console.log(`Sample rate: ${audioContext.sampleRate}Hz`); + analyser = audioContext.createAnalyser(); + analyser.fftSize = 2048; + dataArray = new Uint8Array(analyser.frequencyBinCount); + // Get the range values from widgets (assumed to be in Hz) + const lowRangeWidget = this.widgets.find(w => w.name === "low_range_hz"); + if (lowRangeWidget) startRangeHz = lowRangeWidget.value; + + const midRangeWidget = this.widgets.find(w => w.name === "mid_range_hz"); + if (midRangeWidget) endRangeHz = midRangeWidget.value; + } + + navigator.mediaDevices.getUserMedia({ audio: true }).then(stream => { + microphoneStream = stream; + const microphone = audioContext.createMediaStreamSource(stream); + microphone.connect(analyser); + updateWidgetValueInRealTime(); + }).catch(error => { + console.error('Access to microphone was denied or an error occurred:', error); + }); + }; + + // Function to stop capturing audio from the microphone + const stopMicrophoneCapture = () => { + if (animationFrameId) { + cancelAnimationFrame(animationFrameId); + } + if (microphoneStream) { + microphoneStream.getTracks().forEach(track => track.stop()); + } + if (audioContext) { + audioContext.close(); + // Reset audioContext to ensure it can be created again when starting + audioContext = null; + } + }; + + // Add start button + this.addWidget("button", "Start mic capture", null, startMicrophoneCapture); + + // Add stop button + this.addWidget("button", "Stop mic capture", null, stopMicrophoneCapture); + }; + break; + case "SaveImageKJ": + case "SaveImageWithAlpha": + case "SaveStringKJ": + case "DecodeAndSaveVideo": + case "ModelSaveKJ": + case "LoraExtractKJ": + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function() { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : void 0; + const widget = this.widgets.find((w) => w.name === "filename_prefix"); + widget.serializeValue = () => { + return applyTextReplacements(app, widget.value); + }; + return r; + }; + break; + + } + + }, +}); \ No newline at end of file diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/node_insert.js b/custom_nodes/ComfyUI-KJNodes/web/js/node_insert.js new file mode 100644 index 0000000000000000000000000000000000000000..468c20265aff1612db6c59339d9db476c50847be --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/node_insert.js @@ -0,0 +1,449 @@ +const { app } = window.comfyAPI.app; +import { getSlotPos, clientToCanvas, getNodeAtPoint, typesCompatible, chainCallback } from "./utility.js"; + +// Max age of the last keydown for it to count as the command's trigger. +// Beyond this, treat the command as menu-fired (no release tracking). +const KEYDOWN_MAX_AGE_MS = 100; + +const state = { + pointerDown: false, + hasMoved: false, + insertKeyDown: false, + activationKey: null, + draggedNode: null, + startPos: null, + insertTargetLink: null, + insertSlots: null, + insertOriginNode: null, + insertDestNode: null, + lastScanTime: 0, + animating: false, +}; + +let lastKeyDown = null; +let lastKeyDownTime = 0; + +function bezierAt(p0, p1, p2, p3, t) { + const u = 1 - t; + const uu = u * u, uuu = uu * u; + const tt = t * t, ttt = tt * t; + return [ + uuu * p0[0] + 3 * uu * t * p1[0] + 3 * u * tt * p2[0] + ttt * p3[0], + uuu * p0[1] + 3 * uu * t * p1[1] + 3 * u * tt * p2[1] + ttt * p3[1], + ]; +} + +// Horizontal control-point offset for the link bezier; must match what +// LiteGraph draws so the hit-test region matches the visible curve. +const bezierOffsetX = (from, to) => Math.max(Math.abs(to[0] - from[0]) * 0.5, 50); + +function findLinkUnderNode(graph, draggedNode) { + const bounds = draggedNode.getBounding?.(); + const nodeX = bounds ? bounds[0] : draggedNode.pos[0]; + const nodeY = bounds ? bounds[1] : draggedNode.pos[1]; + const nodeW = bounds ? bounds[2] : (draggedNode.size[0] || 100); + const nodeH = bounds ? bounds[3] : (draggedNode.size[1] || 60); + const nodeCx = nodeX + nodeW / 2; + const nodeCy = nodeY + nodeH / 2; + + let bestLink = null; + let bestDist = Infinity; + + const links = graph.links; + if (!links) return null; + + for (const link of links.values()) { + if (!link) continue; + if (link.origin_id === draggedNode.id || link.target_id === draggedNode.id) continue; + + const originNode = graph.getNodeById(link.origin_id); + const targetNode = graph.getNodeById(link.target_id); + if (!originNode || !targetNode) continue; + + // Reject on node bounds before resolving slot positions — the latter + // hits the DOM in Vue mode and is the expensive step. + const oW = originNode.size?.[0] || 100, oH = originNode.size?.[1] || 60; + const tW = targetNode.size?.[0] || 100, tH = targetNode.size?.[1] || 60; + const cMinX = Math.min(originNode.pos[0], targetNode.pos[0]); + const cMaxX = Math.max(originNode.pos[0] + oW, targetNode.pos[0] + tW); + const cMinY = Math.min(originNode.pos[1], targetNode.pos[1]) - 50; + const cMaxY = Math.max(originNode.pos[1] + oH, targetNode.pos[1] + tH) + 50; + if (cMaxX < nodeX || cMinX > nodeX + nodeW || cMaxY < nodeY || cMinY > nodeY + nodeH) continue; + + const outPos = getSlotPos(originNode, false, link.origin_slot); + const inPos = getSlotPos(targetNode, true, link.target_slot); + + const lMinX = Math.min(outPos[0], inPos[0]); + const lMaxX = Math.max(outPos[0], inPos[0]); + const lMinY = Math.min(outPos[1], inPos[1]) - 50; // bezier can sag + const lMaxY = Math.max(outPos[1], inPos[1]) + 50; + if (lMaxX < nodeX || lMinX > nodeX + nodeW || lMaxY < nodeY || lMinY > nodeY + nodeH) continue; + + const offsetX = bezierOffsetX(outPos, inPos); + const p0 = outPos; + const p1 = [outPos[0] + offsetX, outPos[1]]; + const p2 = [inPos[0] - offsetX, inPos[1]]; + const p3 = inPos; + + for (let i = 0; i <= 20; i++) { + const pt = bezierAt(p0, p1, p2, p3, i / 20); + if (pt[0] >= nodeX && pt[0] <= nodeX + nodeW && pt[1] >= nodeY && pt[1] <= nodeY + nodeH) { + const d = Math.hypot(pt[0] - nodeCx, pt[1] - nodeCy); + if (d < bestDist) { + bestDist = d; + bestLink = link; + } + break; + } + } + } + return bestLink; +} + +function findInsertSlots(node, linkType) { + const inputSlot = node.inputs?.findIndex(i => i.link == null && typesCompatible(linkType, i.type)) ?? -1; + const outputSlot = node.outputs?.findIndex(o => typesCompatible(linkType, o.type)) ?? -1; + if (inputSlot === -1 || outputSlot === -1) return null; + return { inputSlot, outputSlot }; +} + +function executeNodeInsert(canvas, node, link) { + const graph = canvas.graph || app.graph; + + const originNode = graph.getNodeById(link.origin_id); + const targetNode = graph.getNodeById(link.target_id); + if (!originNode || !targetNode) return; + + const slots = findInsertSlots(node, link.type); + if (!slots) return; + + const originSlot = link.origin_slot; + const targetSlot = link.target_slot; + + targetNode.disconnectInput(targetSlot); + originNode.connect(originSlot, node, slots.inputSlot); + node.connect(slots.outputSlot, targetNode, targetSlot); + + graph.change(); + canvas.setDirty(true, true); +} + +function setInsertTarget(link, graph, slots = null) { + if (link === state.insertTargetLink) return; + // `_dragging` is LiteGraph's own runtime flag for hiding a link mid-rewire. + // Not serialized, so save/copy/undo stay safe. + if (state.insertTargetLink) delete state.insertTargetLink._dragging; + if (link) { + link._dragging = true; + state.insertSlots = slots ?? findInsertSlots(state.draggedNode, link.type); + state.insertOriginNode = graph.getNodeById(link.origin_id); + state.insertDestNode = graph.getNodeById(link.target_id); + } else { + state.insertSlots = null; + state.insertOriginNode = null; + state.insertDestNode = null; + } + state.insertTargetLink = link; +} + +function getTypeColor(lgCanvas, slotType) { + if (slotType && LGraphCanvas.link_type_colors && LGraphCanvas.link_type_colors[slotType]) { + return LGraphCanvas.link_type_colors[slotType]; + } + return lgCanvas.default_link_color || "#AAA"; +} + +function drawGhostLink(ctx, from, to, color, alpha, drawTime) { + const offsetX = bezierOffsetX(from, to); + + ctx.beginPath(); + ctx.moveTo(from[0], from[1]); + ctx.bezierCurveTo( + from[0] + offsetX, from[1], + to[0] - offsetX, to[1], + to[0], to[1], + ); + + ctx.strokeStyle = color; + ctx.globalAlpha = alpha; + ctx.lineWidth = 2.5; + ctx.setLineDash([8, 4]); + ctx.lineDashOffset = -(drawTime ?? performance.now()) / 50; + + ctx.stroke(); + ctx.setLineDash([]); + ctx.globalAlpha = 1; +} + +function drawSlotHighlight(ctx, pos, color, alpha) { + ctx.beginPath(); + ctx.arc(pos[0], pos[1], 6, 0, Math.PI * 2); + ctx.globalAlpha = alpha; + ctx.fillStyle = color; + ctx.fill(); + ctx.globalAlpha = 1; +} + +function drawGhostSegment(ctx, from, to, color, now) { + drawGhostLink(ctx, from, to, color, 0.8, now); + drawSlotHighlight(ctx, from, color, 0.6); + drawSlotHighlight(ctx, to, color, 0.6); +} + +function drawInsertPreview(ctx, lgCanvas) { + const { insertTargetLink: link, insertSlots: slots, insertOriginNode, insertDestNode, draggedNode } = state; + if (!link || !slots || !insertOriginNode || !insertDestNode || !draggedNode) return; + + const now = performance.now(); + const color = getTypeColor(lgCanvas, link.type); + + drawGhostSegment(ctx, + getSlotPos(insertOriginNode, false, link.origin_slot), + getSlotPos(draggedNode, true, slots.inputSlot), + color, now); + drawGhostSegment(ctx, + getSlotPos(draggedNode, false, slots.outputSlot), + getSlotPos(insertDestNode, true, link.target_slot), + color, now); +} + +function startAnimLoop(lgCanvas) { + if (state.animating) return; + state.animating = true; + + function tick() { + const link = state.insertTargetLink; + if (!link) { + state.animating = false; + lgCanvas.setDirty(true, true); + return; + } + // If the link was removed externally (e.g. undo) mid-preview, our + // `_dragging` flag has nothing to land on — drop the preview. + const graph = lgCanvas.graph || app.graph; + const stored = graph?.links?.get(link.id); + if (stored !== link) { + setInsertTarget(null); + state.animating = false; + lgCanvas.setDirty(true, true); + return; + } + // Steady-state: only the dash offset changes — front canvas only. + lgCanvas.setDirty(true, false); + requestAnimationFrame(tick); + } + requestAnimationFrame(tick); +} + +function clearState() { + setInsertTarget(null); + state.pointerDown = false; + state.hasMoved = false; + state.draggedNode = null; + state.startPos = null; +} + +app.registerExtension({ + name: "KJNodes.NodeInsert", + + settings: [ + { + id: "KJNodes.nodeInsertMode", + name: "Node insert activation", + category: ["KJNodes", "Node Insert", "Activation mode"], + tooltip: "Always: dragging a compatible node onto a link previews insertion. Hotkey: only while the hotkey (default: D) is held. Disabled: feature off.", + type: "combo", + defaultValue: "hotkey", + options: ["always", "hotkey", "disabled"], + }, + ], + + commands: [ + { + id: "KJNodes.NodeInsertMode", + label: "Node insert mode (hold to activate)", + active: () => state.insertKeyDown, + function: () => { + state.insertKeyDown = true; + // Snapshot the key that triggered this so any rebind works. + state.activationKey = (performance.now() - lastKeyDownTime < KEYDOWN_MAX_AGE_MS) + ? lastKeyDown + : null; + if (state.draggedNode) { + const graph = app.canvas?.graph || app.graph; + if (graph) { + const link = findLinkUnderNode(graph, state.draggedNode); + if (link && findInsertSlots(state.draggedNode, link.type)) { + setInsertTarget(link, graph); + startAnimLoop(app.canvas); + } + } + } + }, + }, + ], + + keybindings: [ + { + commandId: "KJNodes.NodeInsertMode", + // Default only — release detection captures the actually-triggering + // key at command-fire time, so user rebinds work. + combo: { key: "d" }, + targetElementId: "graph-canvas", + }, + ], + + async setup() { + const lgCanvas = app.canvas; + const canvasEl = lgCanvas.canvas; + + // Capture phase so `lastKeyDown` is set before ComfyUI fires the command. + document.addEventListener("keydown", (e) => { + if (e.repeat) return; + lastKeyDown = e.key?.toLowerCase() ?? null; + lastKeyDownTime = performance.now(); + }, true); + + document.addEventListener("keyup", (e) => { + const key = e.key?.toLowerCase() ?? null; + // Released key can't be a future trigger — invalidate so a stale + // keydown can't be picked up by a later menu activation. + if (key && key === lastKeyDown) lastKeyDownTime = 0; + + if (!state.insertKeyDown) return; + // No activation key (menu-fired) → blur/pointercancel will clean up. + if (!state.activationKey || key !== state.activationKey) return; + state.insertKeyDown = false; + state.activationKey = null; + if (state.insertTargetLink) { + setInsertTarget(null); + lgCanvas.setDirty(false, true); + } + }); + + // Releases outside the window never reach our keyup/pointerup listeners, + // so without these the `_dragging` flag and hotkey state can stick. + const dropTransient = () => { + state.insertKeyDown = false; + state.activationKey = null; + clearState(); + lgCanvas.setDirty(false, true); + }; + window.addEventListener("blur", dropTransient); + document.addEventListener("visibilitychange", () => { + if (document.hidden) dropTransient(); + }); + document.addEventListener("pointercancel", dropTransient, true); + + // Scope the Vue-mode node lookup to the graph canvas's container so + // stray `[data-node-id]` elements elsewhere can't activate this flow. + const graphRoot = canvasEl.parentElement; + document.addEventListener("pointerdown", (e) => { + if (e.button !== 0) return; + const onCanvas = e.target === canvasEl; + const vueNodeEl = graphRoot?.contains(e.target) + ? e.target?.closest?.("[data-node-id]") + : null; + if (!onCanvas && !vueNodeEl) return; + + state.pointerDown = true; + state.hasMoved = false; + + const graph = lgCanvas.graph || app.graph; + if (graph) { + let node = null; + if (vueNodeEl) { + const nodeId = parseInt(vueNodeEl.getAttribute("data-node-id")); + if (Number.isFinite(nodeId)) node = graph.getNodeById(nodeId); + } else { + const [cx, cy] = clientToCanvas(lgCanvas, e.clientX, e.clientY); + node = getNodeAtPoint(graph, cx, cy); + } + if (node && (node.inputs || node.outputs)) { + state.draggedNode = node; + state.startPos = node.pos ? [node.pos[0], node.pos[1]] : null; + } + } + }, true); + + document.addEventListener("pointermove", () => { + if (!state.pointerDown) return; + if (lgCanvas.connecting_links?.length) return; + if (lgCanvas.resizing_node) return; + if (lgCanvas.node_widget) return; + + const mode = app.ui.settings.getSettingValue("KJNodes.nodeInsertMode") ?? "always"; + if (mode === "disabled") return; + + state.hasMoved = true; + + const active = mode === "always" || state.insertKeyDown; + if (!active) { + setInsertTarget(null); + return; + } + + if (!state.draggedNode) return; + + // Multi-node drags move several nodes together — inserting one of + // them into a link is almost never what the user wants. + const selected = lgCanvas.selected_nodes; + if (selected) { + const count = selected instanceof Map ? selected.size : Object.keys(selected).length; + if (count > 1) { + setInsertTarget(null); + return; + } + } + + // Vue mode: pointermove fires on the document during widget drags + // (sliders etc.) too — require the node itself to have moved. + const np = state.draggedNode.pos; + if (!np || !state.startPos + || (np[0] === state.startPos[0] && np[1] === state.startPos[1])) { + setInsertTarget(null); + return; + } + + const graph = lgCanvas.graph || app.graph; + if (!graph) return; + + const now = performance.now(); + const nodeCount = graph._nodes?.length ?? 0; + const throttle = nodeCount > 200 ? 50 : nodeCount > 100 ? 32 : 16; + if (now - state.lastScanTime < throttle) return; + state.lastScanTime = now; + + const link = findLinkUnderNode(graph, state.draggedNode); + const slots = link ? findInsertSlots(state.draggedNode, link.type) : null; + const wasIdle = !state.insertTargetLink; + + if (link && slots) { + setInsertTarget(link, graph, slots); + if (wasIdle) startAnimLoop(lgCanvas); + } else { + setInsertTarget(null); + } + }, true); + + document.addEventListener("pointerup", (e) => { + if (e.button !== 0) return; + if (!state.pointerDown) return; + + if (state.insertTargetLink && state.draggedNode && state.hasMoved + && !state.draggedNode.flags?.pinned) { + delete state.insertTargetLink._dragging; + executeNodeInsert(lgCanvas, state.draggedNode, state.insertTargetLink); + } + + clearState(); + }, true); + + // onDrawForeground: ghost paints over the existing (hidden) link. + chainCallback(lgCanvas, "onDrawForeground", function (ctx) { + if (state.insertTargetLink) { + drawInsertPreview(ctx, lgCanvas); + } + }); + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/nodeswap.js b/custom_nodes/ComfyUI-KJNodes/web/js/nodeswap.js new file mode 100644 index 0000000000000000000000000000000000000000..d588fed5321cb4caf3fbb12da342732abe58c93e --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/nodeswap.js @@ -0,0 +1,434 @@ +const { app } = window.comfyAPI.app; +import { typesCompatible, clientToCanvas, getNodeAtPoint } from "./utility.js"; + +let swapTargetNode = null; +let swapDraggedNode = null; +let swapDragStartPos = null; +let swapAnimating = false; +let swapHasMoved = false; +let swapKeyDown = false; + +/** Find the topmost node whose bounding box overlaps with draggedNode */ +function getOverlappingNode(graph, draggedNode) { + const ax = draggedNode.pos[0]; + const ay = draggedNode.pos[1]; + const aw = draggedNode.size[0] || 100; + const ah = draggedNode.size[1] || 60; + + for (let i = graph._nodes.length - 1; i >= 0; i--) { + const n = graph._nodes[i]; + if (n === draggedNode) continue; + const bx = n.pos[0]; + const by = n.pos[1]; + if (ax < bx + (n.size[0] || 100) && ax + aw > bx && + ay < by + (n.size[1] || 60) && ay + ah > by) { + return n; + } + } + return null; +} + + +function startHighlightAnim(lgCanvas) { + if (swapAnimating) return; + swapAnimating = true; + + function tick() { + if (!swapTargetNode) { + swapAnimating = false; + return; + } + lgCanvas.setDirty(false, true); + requestAnimationFrame(tick); + } + requestAnimationFrame(tick); +} + +function clearSwapState() { + swapTargetNode = null; + swapDraggedNode = null; + swapDragStartPos = null; + swapHasMoved = false; +} + +function snapshotConnections(graph, node) { + const inputs = []; + for (let i = 0; i < (node.inputs?.length ?? 0); i++) { + const inp = node.inputs[i]; + if (inp.link == null) continue; + const link = graph.getLink(inp.link); + if (!link) continue; + inputs.push({ + slotIndex: i, + type: inp.type, + name: inp.name, + originNodeId: link.origin_id, + originSlot: link.origin_slot, + }); + } + + const outputs = []; + for (let o = 0; o < (node.outputs?.length ?? 0); o++) { + const out = node.outputs[o]; + if (!out.links?.length) continue; + const targets = out.links + .map((linkId) => graph.getLink(linkId)) + .filter(Boolean) + .map((link) => ({ targetNodeId: link.target_id, targetSlot: link.target_slot })); + if (targets.length > 0) { + outputs.push({ slotIndex: o, type: out.type, name: out.name, targets }); + } + } + + return { inputs, outputs }; +} + +function disconnectAll(node) { + if (node.inputs) { + for (let i = node.inputs.length - 1; i >= 0; i--) { + if (node.inputs[i].link != null) node.disconnectInput(i); + } + } + if (node.outputs) { + for (let o = node.outputs.length - 1; o >= 0; o--) { + if (node.outputs[o].links?.length) node.disconnectOutput(o); + } + } +} + +// Map each snapshot entry → a new-node slot. Name matches win first, then sameIndex, +// then first compatible. Each slot is claimed at most once so fallbacks can't stomp matches. +function resolveAssignments(snapshots, slots, isFree) { + const assignments = new Map(); + if (!slots || !slots.length) return assignments; + const used = new Set(); + + for (const snap of snapshots) { + if (!snap.name) continue; + const nameLower = String(snap.name).toLowerCase(); + for (let s = 0; s < slots.length; s++) { + if (used.has(s) || !isFree(s)) continue; + const slotName = slots[s].name; + if (slotName && String(slotName).toLowerCase() === nameLower && + typesCompatible(snap.type, slots[s].type)) { + assignments.set(snap, s); + used.add(s); + break; + } + } + } + + for (const snap of snapshots) { + if (assignments.has(snap)) continue; + const i = snap.slotIndex; + if (i < slots.length && !used.has(i) && isFree(i) && + typesCompatible(snap.type, slots[i].type)) { + assignments.set(snap, i); + used.add(i); + continue; + } + for (let s = 0; s < slots.length; s++) { + if (used.has(s) || !isFree(s)) continue; + if (typesCompatible(snap.type, slots[s].type)) { + assignments.set(snap, s); + used.add(s); + break; + } + } + } + + return assignments; +} + +// Reconnect external connections from a snapshot onto newNode. +function reconnectExternal(graph, snapshot, newNode, otherNodeId) { + const inputAssignments = resolveAssignments( + snapshot.inputs, newNode.inputs || [], + (s) => newNode.inputs[s].link == null, + ); + for (const inp of snapshot.inputs) { + if (inp.originNodeId === otherNodeId) continue; + if (!inputAssignments.has(inp)) continue; + const originNode = graph.getNodeById(inp.originNodeId); + if (!originNode) continue; + originNode.connect(inp.originSlot, newNode, inputAssignments.get(inp)); + } + + const outputAssignments = resolveAssignments( + snapshot.outputs, newNode.outputs || [], + () => true, + ); + for (const out of snapshot.outputs) { + if (!outputAssignments.has(out)) continue; + const bestSlot = outputAssignments.get(out); + for (const tgt of out.targets) { + if (tgt.targetNodeId === otherNodeId) continue; + const targetNode = graph.getNodeById(tgt.targetNodeId); + if (!targetNode) continue; + newNode.connect(bestSlot, targetNode, tgt.targetSlot); + } + } +} + +// Reconnect links that were between the two swapped nodes in one direction. +// A→B becomes B→A, if slot types are compatible. +function reconnectInternalOneDirection(snap, fromNodeId, srcNode, dstNode) { + for (const out of snap.outputs) { + for (const tgt of out.targets) { + if (tgt.targetNodeId !== fromNodeId) continue; + if (!srcNode.outputs || out.slotIndex >= srcNode.outputs.length) continue; + if (!dstNode.inputs || tgt.targetSlot >= dstNode.inputs.length) continue; + if (!typesCompatible(srcNode.outputs[out.slotIndex].type, dstNode.inputs[tgt.targetSlot].type)) continue; + srcNode.connect(out.slotIndex, dstNode, tgt.targetSlot); + } + } +} + +function getNodeElement(nodeId) { + return document.querySelector(`[data-node-id="${nodeId}"]`); +} + +function animateSwapVue(nodeA, toA, nodeB, toB, duration) { + const elA = getNodeElement(nodeA.id); + const elB = getNodeElement(nodeB.id); + const transition = `transform ${duration}ms cubic-bezier(0.65, 0, 0.35, 1)`; + + if (elA) elA.style.transition = transition; + if (elB) elB.style.transition = transition; + + nodeA.pos = [toA[0], toA[1]]; + nodeB.pos = [toB[0], toB[1]]; + app.canvas.setDirty(true, true); + + setTimeout(() => { + if (elA) elA.style.transition = ""; + if (elB) elB.style.transition = ""; + }, duration + 16); +} + +function animateSwapCanvas(nodeA, fromA, toA, nodeB, fromB, toB, duration) { + const start = performance.now(); + + function ease(t) { + return t < 0.5 ? 4 * t * t * t : 1 - Math.pow(-2 * t + 2, 3) / 2; + } + + function frame(now) { + const t = Math.min((now - start) / duration, 1); + const e = ease(t); + nodeA.pos = [ + fromA[0] + (toA[0] - fromA[0]) * e, + fromA[1] + (toA[1] - fromA[1]) * e, + ]; + nodeB.pos = [ + fromB[0] + (toB[0] - fromB[0]) * e, + fromB[1] + (toB[1] - fromB[1]) * e, + ]; + app.canvas.setDirty(true, true); + if (t < 1) requestAnimationFrame(frame); + } + + requestAnimationFrame(frame); +} + +function executeNodeSwap(canvas, nodeA, nodeB, originalPosA) { + const graph = canvas.graph || app.graph; + + const snapA = snapshotConnections(graph, nodeA); + const snapB = snapshotConnections(graph, nodeB); + + const posA = originalPosA || [nodeA.pos[0], nodeA.pos[1]]; + const posB = [nodeB.pos[0], nodeB.pos[1]]; + + disconnectAll(nodeA); + disconnectAll(nodeB); + + reconnectExternal(graph, snapA, nodeB, nodeA.id); + reconnectExternal(graph, snapB, nodeA, nodeB.id); + + reconnectInternalOneDirection(snapA, nodeB.id, nodeB, nodeA); + reconnectInternalOneDirection(snapB, nodeA.id, nodeA, nodeB); + + if (LiteGraph.vueNodesMode) { + animateSwapVue(nodeA, posB, nodeB, posA, 200); + } else { + animateSwapCanvas(nodeA, [nodeA.pos[0], nodeA.pos[1]], posB, + nodeB, [nodeB.pos[0], nodeB.pos[1]], posA, 200); + } + + graph.change(); + canvas.setDirty(true, true); +} + +app.registerExtension({ + name: "KJNodes.NodeSwap", + + settings: [ + { + id: "KJNodes.nodeSwapEnabled", + name: "Enable node swap on drag", + category: ["KJNodes", "Node Swap", "Enable"], + tooltip: "Hold swap key (default: S, rebindable in Keybindings) and drag a node onto another to swap their positions and reconnect links", + type: "boolean", + defaultValue: true, + }, + ], + commands: [ + { + id: "KJNodes.ToggleNodeSwap", + label: "Toggle node swap on drag", + active: () => app.ui.settings.getSettingValue("KJNodes.nodeSwapEnabled"), + function: () => { + const cur = app.ui.settings.getSettingValue("KJNodes.nodeSwapEnabled"); + app.ui.settings.setSettingValue("KJNodes.nodeSwapEnabled", !cur); + }, + }, + { + id: "KJNodes.NodeSwapMode", + label: "Node swap mode (hold to activate)", + active: () => swapKeyDown, + function: () => { + swapKeyDown = true; + // If already dragging a node, check for overlap immediately + // (handles case where node is already on top of another when key is pressed) + if (swapDraggedNode) { + const graph = app.canvas?.graph || app.graph; + if (graph) { + swapTargetNode = getOverlappingNode(graph, swapDraggedNode); + if (swapTargetNode) startHighlightAnim(app.canvas); + } + } + }, + }, + ], + keybindings: [ + { + commandId: "KJNodes.NodeSwapMode", + combo: { key: "s" }, + targetElementId: "graph-canvas", + }, + ], + + async setup() { + await new Promise((resolve) => { + function check() { + if (app.canvas) return resolve(); + requestAnimationFrame(check); + } + check(); + }); + + const lgCanvas = app.canvas; + const canvasEl = lgCanvas.canvas; + + // The command system handles keydown (sets swapKeyDown = true). + // Any keyup clears it since the user must hold the key. + document.addEventListener("keyup", () => { + if (swapKeyDown) { + swapKeyDown = false; + if (swapTargetNode) { + swapTargetNode = null; + lgCanvas.setDirty(false, true); + } + } + }); + + // Pointermove/pointerup handlers — only attached while pointer is down. + function onPointerMove() { + swapHasMoved = true; + + if (!swapKeyDown) { + if (swapTargetNode) { + swapTargetNode = null; + lgCanvas.setDirty(false, true); + } + return; + } + + if (!swapDraggedNode) return; + + const graph = lgCanvas.graph || app.graph; + if (!graph) return; + + const prev = swapTargetNode; + swapTargetNode = getOverlappingNode(graph, swapDraggedNode); + + if (swapTargetNode && !prev) { + startHighlightAnim(lgCanvas); + } + } + + function onPointerUp(e) { + if (e.button !== 0) return; + + document.removeEventListener("pointermove", onPointerMove); + document.removeEventListener("pointerup", onPointerUp, true); + + const targetNode = swapTargetNode; + const draggedNode = swapDraggedNode; + const startPos = swapDragStartPos; + const hasMoved = swapHasMoved; + clearSwapState(); + + if (!hasMoved || !targetNode || !draggedNode || !swapKeyDown) return; + + if (!targetNode.flags?.pinned && !draggedNode.flags?.pinned) { + executeNodeSwap(lgCanvas, draggedNode, targetNode, startPos); + } + } + + // Record clicked node and its original position. + // Listen on document to catch events from both the canvas and Vue node overlays. + const graphContainer = canvasEl.closest(".graph-canvas-container") || canvasEl.parentElement; + document.addEventListener("pointerdown", (e) => { + if (e.button !== 0) return; + if (!app.ui.settings.getSettingValue("KJNodes.nodeSwapEnabled")) return; + if (!graphContainer?.contains(e.target)) return; + + swapHasMoved = false; + + const graph = lgCanvas.graph || app.graph; + if (!graph) return; + + const [cx, cy] = clientToCanvas(lgCanvas, e.clientX, e.clientY); + const node = getNodeAtPoint(graph, cx, cy); + + if (node && (node.inputs || node.outputs)) { + swapDraggedNode = node; + swapDragStartPos = [node.pos[0], node.pos[1]]; + document.addEventListener("pointermove", onPointerMove); + document.addEventListener("pointerup", onPointerUp, true); + } + }, true); + + // Draw pulsing glow outline on the swap target node + const origOnDrawBg = lgCanvas.onDrawBackground; + lgCanvas.onDrawBackground = function (ctx, ...args) { + origOnDrawBg?.call(this, ctx, ...args); + + const target = swapTargetNode; + if (!target) return; + + const [x, y] = target.pos; + const w = target.size[0] || 100; + const h = target.size[1] || 60; + const titleH = LiteGraph.NODE_TITLE_HEIGHT || 20; + const pulse = 0.5 + 0.3 * Math.sin(Date.now() / 200); + const isPinned = target.flags?.pinned; + const color = isPinned ? "255, 80, 80" : "100, 200, 255"; + const pad = 5; + + ctx.save(); + ctx.strokeStyle = `rgba(${color}, ${pulse})`; + ctx.lineWidth = 3; + ctx.shadowColor = `rgba(${color}, 0.8)`; + ctx.shadowBlur = 15; + ctx.beginPath(); + ctx.roundRect(x - pad, y - titleH - pad, w + pad * 2, h + titleH + pad * 2, 8); + ctx.stroke(); + ctx.fillStyle = `rgba(${color}, ${isPinned ? 0.08 : 0.05})`; + ctx.fill(); + ctx.restore(); + }; + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/performance.js b/custom_nodes/ComfyUI-KJNodes/web/js/performance.js new file mode 100644 index 0000000000000000000000000000000000000000..4420dce237fab545326b298cd175b27bd85e67bd --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/performance.js @@ -0,0 +1,286 @@ +const { app } = window.comfyAPI.app; + +const _savedRoundRadius = typeof LiteGraph !== "undefined" ? LiteGraph.ROUND_RADIUS : 8; + +function getSetting(id, fallback) { + return app.ui.settings.getSettingValue(id) ?? fallback; +} + +function toggleSetting(id) { + const cur = app.ui.settings.getSettingValue(id); + app.ui.settings.setSettingValue(id, !cur); +} + +app.registerExtension({ + name: "KJNodes.Performance", + + commands: [ + { + id: "KJNodes.perf.toggleSingleCanvasPan", + label: "Single-canvas mode during pan", + menubarLabel: "Single-canvas mode during pan", + function: () => toggleSetting("KJNodes.perf.singleCanvasPan"), + active: () => getSetting("KJNodes.perf.singleCanvasPan", false), + }, + { + id: "KJNodes.perf.toggleDisableShadows", + label: "Disable node shadows", + menubarLabel: "Disable node shadows", + function: () => toggleSetting("KJNodes.perf.disableShadows"), + active: () => getSetting("KJNodes.perf.disableShadows", false), + }, + { + id: "KJNodes.perf.toggleDisableConnectionBorders", + label: "Disable connection borders", + menubarLabel: "Disable connection borders", + function: () => toggleSetting("KJNodes.perf.disableConnectionBorders"), + active: () => getSetting("KJNodes.perf.disableConnectionBorders", false), + }, + { + id: "KJNodes.perf.toggleDisableRoundedCorners", + label: "Disable rounded corners", + menubarLabel: "Disable rounded corners", + function: () => toggleSetting("KJNodes.perf.disableRoundedCorners"), + active: () => getSetting("KJNodes.perf.disableRoundedCorners", false), + }, + { + id: "KJNodes.perf.toggleThrottleRenderInfo", + label: "Throttle info overlay", + menubarLabel: "Throttle info overlay", + function: () => toggleSetting("KJNodes.perf.throttleRenderInfo"), + active: () => getSetting("KJNodes.perf.throttleRenderInfo", false), + }, + ], + + menuCommands: [ + { + path: ["KJNodes", "Performance"], + commands: [ + "KJNodes.perf.toggleSingleCanvasPan", + "KJNodes.perf.toggleDisableShadows", + "KJNodes.perf.toggleDisableConnectionBorders", + "KJNodes.perf.toggleDisableRoundedCorners", + "KJNodes.perf.toggleThrottleRenderInfo", + ], + }, + ], + + settings: [ + { + id: "KJNodes.perf.singleCanvasPan", + name: "Single-canvas mode during pan", + category: ["KJNodes", "Performance", "Single-canvas mode during pan"], + tooltip: "Eliminates the expensive bgcanvas-to-canvas copy during panning. Significant improvement without hardware acceleration.", + type: "boolean", + defaultValue: false, + onChange: (value) => { + const canvas = app.canvas; + if (!canvas) return; + if (value) { + installSingleCanvasPan(canvas); + } else { + uninstallSingleCanvasPan(canvas); + } + }, + }, + { + id: "KJNodes.perf.disableShadows", + name: "Disable node shadows", + category: ["KJNodes", "Performance", "Disable node shadows"], + tooltip: "Disable shadow rendering on nodes for better performance without hardware acceleration.", + type: "boolean", + defaultValue: false, + onChange: (value) => { + if (app.canvas) { + app.canvas.render_shadows = !value; + app.canvas.setDirty(true, true); + } + }, + }, + { + id: "KJNodes.perf.disableConnectionBorders", + name: "Disable connection borders", + category: ["KJNodes", "Performance", "Disable connection borders"], + tooltip: "Disable the outer border stroke on connection lines. Each link draws two strokes instead of one when enabled.", + type: "boolean", + defaultValue: false, + onChange: (value) => { + if (app.canvas) { + app.canvas.render_connections_border = !value; + app.canvas.setDirty(true, true); + } + }, + }, + { + id: "KJNodes.perf.disableRoundedCorners", + name: "Disable rounded corners", + category: ["KJNodes", "Performance", "Disable rounded corners"], + tooltip: "Use square corners on nodes instead of rounded. Avoids roundRect calls in software rendering.", + type: "boolean", + defaultValue: false, + onChange: (value) => { + if (typeof LiteGraph !== "undefined") { + LiteGraph.ROUND_RADIUS = value ? 0 : _savedRoundRadius; + app.canvas?.setDirty(true, true); + } + }, + }, + { + id: "KJNodes.perf.throttleRenderInfo", + name: "Throttle info overlay", + category: ["KJNodes", "Performance", "Throttle info overlay"], + tooltip: "Cache the FPS/info text overlay and only re-render it a few times per second. Saves ~20ms per frame in software rendering.", + type: "boolean", + defaultValue: false, + onChange: (value) => { + const canvas = app.canvas; + if (!canvas) return; + if (value) { + installThrottleRenderInfo(canvas); + } else { + uninstallThrottleRenderInfo(canvas); + } + }, + }, + ], + + setup() { + const canvas = app.canvas; + if (!canvas) return; + + if (getSetting("KJNodes.perf.singleCanvasPan", false)) { + installSingleCanvasPan(canvas); + } + if (getSetting("KJNodes.perf.disableShadows", false)) { + canvas.render_shadows = false; + } + if (getSetting("KJNodes.perf.disableConnectionBorders", false)) { + canvas.render_connections_border = false; + } + if (getSetting("KJNodes.perf.disableRoundedCorners", false)) { + LiteGraph.ROUND_RADIUS = 0; + } + if (getSetting("KJNodes.perf.throttleRenderInfo", false)) { + installThrottleRenderInfo(canvas); + } + }, +}); + +// ── Single-canvas mode during pan ─────────────────────────────────────────── + +let _origDraw = null; +let _panInstalledOn = null; + +function installSingleCanvasPan(canvas) { + if (_origDraw && _panInstalledOn === canvas) return; + if (_origDraw && _panInstalledOn !== canvas) { + // Canvas was replaced — discard stale reference + _origDraw = null; + } + _panInstalledOn = canvas; + _origDraw = canvas.draw; + let panning = false; + let savedBgCanvas = null; + let savedBgCtx = null; + + canvas.draw = function (force_canvas, force_bgcanvas) { + if (this.dragging_canvas) { + if (!panning) { + // Flush any pending bg redraw before aliasing to avoid a stale-frame flash + if (this.dirty_bgcanvas) { + _origDraw.call(this, false, true); + } + savedBgCanvas = this.bgcanvas; + savedBgCtx = this.bgctx; + panning = true; + this.bgcanvas = this.canvas; + this.bgctx = this.ctx; + } + } else if (panning) { + panning = false; + this.bgcanvas = savedBgCanvas; + this.bgctx = savedBgCtx; + this.dirty_bgcanvas = true; + this.dirty_canvas = true; + } + _origDraw.call(this, force_canvas, force_bgcanvas); + }; +} + +function uninstallSingleCanvasPan(canvas) { + if (!_origDraw) return; + if (_panInstalledOn === canvas) { + canvas.draw = _origDraw; + } + _origDraw = null; + _panInstalledOn = null; + canvas.dirty_bgcanvas = true; + canvas.dirty_canvas = true; +} + +// ── Throttle renderInfo ───────────────────────────────────────────────────── + +let _origRenderInfo = null; +let _infoCanvas = null; +let _infoInstalledOn = null; + +function installThrottleRenderInfo(canvas) { + if (_origRenderInfo && _infoInstalledOn === canvas) return; + if (_origRenderInfo && _infoInstalledOn !== canvas) { + _origRenderInfo = null; + _infoCanvas = null; + } + _infoInstalledOn = canvas; + _origRenderInfo = canvas.renderInfo; + _infoCanvas = document.createElement("canvas"); + let infoCtx = _infoCanvas.getContext("2d"); + let lastInfoTime = 0; + let cachedDpr = 0; + let _infoDrawX = 10; + let _infoDrawY = 0; + + canvas.renderInfo = function (ctx, x, y) { + const dpr = window.devicePixelRatio || 1; + + // Resize offscreen canvas if DPR changed + if (dpr !== cachedDpr) { + cachedDpr = dpr; + _infoCanvas.width = Math.ceil(200 * dpr); + _infoCanvas.height = Math.ceil(100 * dpr); + infoCtx = _infoCanvas.getContext("2d"); + infoCtx.scale(dpr, dpr); + lastInfoTime = 0; + } + + const now = performance.now(); + if (now - lastInfoTime > 250) { + lastInfoTime = now; + infoCtx.clearRect(0, 0, 200, 100); + // Render with the same x/y the original would use, but offset + // into offscreen canvas coordinates + const origX = x || 10; + const lineHeight = 13; + const lineCount = (this.graph ? 5 : 1) + (this.info_text ? 1 : 0); + const origY = y || (this.canvas.height / dpr - (lineCount + 1) * lineHeight); + // Store the computed position for blitting + _infoDrawX = origX; + _infoDrawY = origY; + // Draw into offscreen at (1, 1), text flows from there + _origRenderInfo.call(this, infoCtx, 1, 1); + } + + // Blit cached text at the original position + ctx.drawImage(_infoCanvas, _infoDrawX - 1, _infoDrawY - 1, 200, 100); + }; +} + +function uninstallThrottleRenderInfo(canvas) { + if (!_origRenderInfo) return; + if (_infoInstalledOn === canvas) { + canvas.renderInfo = _origRenderInfo; + } + _origRenderInfo = null; + _infoCanvas = null; + _infoInstalledOn = null; + canvas.setDirty(true, true); +} diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/play_sound.js b/custom_nodes/ComfyUI-KJNodes/web/js/play_sound.js new file mode 100644 index 0000000000000000000000000000000000000000..cda30caea56623e650836b77a5e532f47c3cae8c --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/play_sound.js @@ -0,0 +1,94 @@ +const { app } = window.comfyAPI.app; +const { api } = window.comfyAPI.api; + +app.registerExtension({ + name: "KJNodes.PlaySound", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== "PlaySoundKJ") return; + + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (output) { + onExecuted?.call(this, output); + + const audios = output?.audio; + if (!audios?.length) return; + + const modeWidget = this.widgets?.find(w => w.name === "mode"); + const volumeWidget = this.widgets?.find(w => w.name === "volume"); + const durationWidget = this.widgets?.find(w => w.name === "duration"); + const mode = modeWidget?.value ?? "always"; + const volume = volumeWidget?.value ?? 0.5; + const duration = durationWidget?.value ?? 0; + + // on_change: skip if audio content hasn't changed + if (mode === "on_change") { + const audioHash = output?.audio_hash?.[0]; + if (audioHash != null && this._kjLastAudioHash === audioHash) return; + this._kjLastAudioHash = audioHash; + } + + // Clean up previous state + if (this._kjStatusListener) { + api.removeEventListener("status", this._kjStatusListener); + this._kjStatusListener = null; + } + clearTimeout(this._kjQueueDebounce); + this._kjPendingAudio = null; + + if (this._kjPlayingAudio) { + this._kjPlayingAudio.pause(); + this._kjPlayingAudio = null; + } + if (this._kjPlayTimer != null) { + clearTimeout(this._kjPlayTimer); + this._kjPlayTimer = null; + } + + const startPlayback = () => { + const { filename, subfolder, type } = audios[0]; + const params = new URLSearchParams({ + filename: filename ?? "", + subfolder: subfolder ?? "", + type: type ?? "temp", + }); + const url = api.apiURL(`/view?${params.toString()}`); + const audio = new Audio(url); + audio.volume = Math.max(0, Math.min(1, volume)); + audio.play().catch(() => {}); + this._kjPlayingAudio = audio; + if (duration > 0) { + this._kjPlayTimer = setTimeout(() => { + audio.pause(); + this._kjPlayingAudio = null; + this._kjPlayTimer = null; + }, duration * 1000); + } + }; + + if (mode === "on_empty_queue") { + this._kjPendingAudio = startPlayback; + this._kjStatusListener = ({ detail }) => { + const remaining = detail?.exec_info?.queue_remaining ?? 0; + if (remaining === 0) { + // Debounce: confirm queue is truly empty + // (status can briefly show 0 between dispatches) + clearTimeout(this._kjQueueDebounce); + this._kjQueueDebounce = setTimeout(() => { + if (this._kjPendingAudio) { + this._kjPendingAudio(); + this._kjPendingAudio = null; + } + api.removeEventListener("status", this._kjStatusListener); + this._kjStatusListener = null; + }, 1000); + } else { + clearTimeout(this._kjQueueDebounce); + } + }; + api.addEventListener("status", this._kjStatusListener); + } else { + startPlayback(); + } + }; + }, +}); diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/preview_override/preview_override.css b/custom_nodes/ComfyUI-KJNodes/web/js/preview_override/preview_override.css new file mode 100644 index 0000000000000000000000000000000000000000..c5197d80df27cbc767d677b78cc260e1522f8e0a --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/preview_override/preview_override.css @@ -0,0 +1,194 @@ +/* Palette matched to ComfyUI-MemoryVisualization. */ + +.kj-pov-root { + width: 100%; + height: 100%; + position: relative; + background: #181818; + color: #b0b0b0; + font: 12px monospace; + display: flex; + flex-direction: column; + overflow: hidden; + user-select: none; +} + +.kj-pov-image-area { + flex: 1 1 auto; + min-height: 80px; + position: relative; + display: flex; + align-items: center; + justify-content: center; + background: #0a0a0a; + overflow: hidden; +} + +/* Absolute-fill so latent-resolution previews scale up rather than staying centred-tiny. */ +.kj-pov-img { + position: absolute; + inset: 0; + width: 100%; + height: 100%; + object-fit: contain; + image-rendering: pixelated; + display: block; +} + +.kj-pov-placeholder { + position: absolute; + inset: 0; + display: flex; + align-items: center; + justify-content: center; + color: #707070; + font: 12px monospace; + pointer-events: none; +} + +.kj-pov-scrub { + position: absolute; + left: 0; + right: 0; + bottom: 0; + height: 6px; + background: rgba(0, 0, 0, 0.5); + cursor: pointer; + z-index: 5; + transition: height 80ms ease; +} +.kj-pov-scrub:hover { height: 10px; } +.kj-pov-scrub-fill { + height: 100%; + width: 0%; + background: #e67e22; + pointer-events: none; +} +.kj-pov-scrub.kj-pov-paused .kj-pov-scrub-fill { + background: #707070; +} + +.kj-pov-panel-grip { + flex: 0 0 5px; + background: #2a2a2a; + cursor: ns-resize; + position: relative; + z-index: 4; + user-select: none; +} +.kj-pov-panel-grip::after { + content: ''; + position: absolute; + left: 50%; + top: 50%; + transform: translate(-50%, -50%); + width: 26px; + height: 2px; + background: #555; + border-radius: 1px; +} +.kj-pov-panel-grip:hover { background: #3a3a3a; } +.kj-pov-panel-grip:hover::after { background: #888; } + +.kj-pov-panel { + flex: 0 0 auto; + height: 140px; + min-height: 60px; + background: #181818; + font: 11px monospace; + color: #b0b0b0; + display: flex; + flex-direction: column; + overflow: hidden; +} + +.kj-pov-panel-header { + display: flex; + justify-content: space-between; + align-items: center; + padding: 4px 10px; + gap: 8px; + background: #202020; + border-bottom: 1px solid #2a2a2a; +} +.kj-pov-header-left { + display: flex; + align-items: center; + gap: 8px; +} +.kj-pov-panel-title { color: #b0b0b0; font-size: 11px; } +.kj-pov-panel-summary { color: #707070; font-size: 10px; } + +.kj-pov-toggle { + background: #2a2a2a; + color: #888; + border: 1px solid #333; + border-radius: 3px; + padding: 1px 7px; + font: 10px monospace; + cursor: pointer; + user-select: none; +} +.kj-pov-toggle:hover { background: #333; color: #b0b0b0; } +.kj-pov-toggle.kj-pov-on { + background: #3a2a1a; + color: #e67e22; + border-color: #c66e1a; +} + +.kj-pov-graphs-grid { + display: flex; + align-items: stretch; + padding: 6px 6px 7px; + gap: 6px; + flex: 1 1 auto; + min-height: 0; + box-sizing: border-box; +} + +.kj-pov-graph-cell { + flex: 1 1 0; + min-width: 0; + display: flex; + flex-direction: column; + background: #0e0e0e; + border: 1px solid #1c1c1c; + border-radius: 3px; + overflow: hidden; +} + +.kj-pov-graph-head { + flex: 0 0 auto; + display: flex; + justify-content: space-between; + align-items: baseline; + padding: 3px 6px 2px; + background: #161616; + border-bottom: 1px solid #1c1c1c; + font-size: 10px; + line-height: 1.1; + gap: 4px; +} + +.kj-pov-graph-label { + color: #707070; + letter-spacing: 0.05em; + flex: 0 0 auto; +} + +.kj-pov-graph-value { + color: #b0b0b0; + font-size: 10px; + text-align: right; + flex: 1 1 auto; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.kj-pov-graph-canvas { + flex: 1 1 auto; + display: block; + width: 100%; + min-height: 0; +} diff --git a/custom_nodes/ComfyUI-KJNodes/web/js/preview_override/preview_override.js b/custom_nodes/ComfyUI-KJNodes/web/js/preview_override/preview_override.js new file mode 100644 index 0000000000000000000000000000000000000000..cc1cf1a8e23e0e38bc0dc2e5dd4d6f64d7ab50cd --- /dev/null +++ b/custom_nodes/ComfyUI-KJNodes/web/js/preview_override/preview_override.js @@ -0,0 +1,948 @@ +import { chainCallback, addMiddleClickPan, addWheelPassthrough } from '../utility.js'; +const { app } = window.comfyAPI.app; +const { api } = window.comfyAPI.api; + +const STYLE_ID = "kj-pov-stylesheet"; +const _cssUrl = new URL("./preview_override.css", import.meta.url).href; +function ensureStyles() { + if (document.getElementById(STYLE_ID)) return; + const link = document.createElement("link"); + link.id = STYLE_ID; + link.rel = "stylesheet"; + link.href = _cssUrl; + document.head.appendChild(link); +} + +// Walks subgraph chain for IDs like "12:7:5". Mirrors getNodeByExecutionId (not exported). +function findNodeByQualifiedId(rootGraph, qid) { + if (!rootGraph || !qid) return null; + const parts = String(qid).split(":"); + let graph = rootGraph; + for (let i = 0; i < parts.length - 1; i++) { + const parentId = parseInt(parts[i], 10); + if (!Number.isFinite(parentId)) return null; + const parentNode = graph?.getNodeById?.(parentId); + if (!parentNode?.subgraph) return null; + graph = parentNode.subgraph; + } + const leafId = parseInt(parts[parts.length - 1], 10); + if (!Number.isFinite(leafId)) return null; + return graph?.getNodeById?.(leafId) || null; +} + +api.addEventListener("kj_preview_override", (e) => { + const data = e.detail; + if (!data || data.node_id == null) return; + const node = findNodeByQualifiedId(app.graph, data.node_id); + if (node?._kjPreviewHandler) node._kjPreviewHandler(data); +}); + +const GRAPH_PAD_X = 4; + +function fmt(n, d) { + return Number.isFinite(n) ? n.toFixed(d) : "—"; +} + +function el(tag, className, parent) { + const e = document.createElement(tag); + if (className) e.className = className; + if (parent) parent.appendChild(e); + return e; +} + +function b64ToBlob(b64, mime) { + const bin = atob(b64); + const arr = new Uint8Array(bin.length); + for (let i = 0; i < bin.length; i++) arr[i] = bin.charCodeAt(i); + return new Blob([arr], { type: mime }); +} + +function syncCanvasDPR(canvas) { + const dpr = window.devicePixelRatio || 1; + const cssW = canvas.clientWidth || canvas.width; + const cssH = canvas.clientHeight || canvas.height; + if (canvas.width !== Math.round(cssW * dpr) || canvas.height !== Math.round(cssH * dpr)) { + canvas.width = Math.round(cssW * dpr); + canvas.height = Math.round(cssH * dpr); + } + const ctx = canvas.getContext("2d"); + ctx.setTransform(dpr, 0, 0, dpr, 0, 0); + return { ctx, W: cssW, H: cssH }; +} + +function drawGridlines(ctx, W, H, padX, padY) { + ctx.strokeStyle = "#1e1e1e"; + ctx.lineWidth = 1; + for (let g = 1; g < 4; g++) { + const y = Math.round(padY + (g / 4) * (H - 2 * padY)) + 0.5; + ctx.beginPath(); + ctx.moveTo(padX, y); + ctx.lineTo(W - padX, y); + ctx.stroke(); + } +} + +// σ + Δ overlaid; each series normalised to its own max for comparable shapes. +function drawSigmaDeltaGraph(canvas, sigmas, deltas, step, totalSteps, hoverStep, dbCurve, lockedStep) { + const { ctx, W, H } = syncCanvasDPR(canvas); + const padX = GRAPH_PAD_X, padY = 3; + const iW = W - 2 * padX, iH = H - 2 * padY; + ctx.clearRect(0, 0, W, H); + drawGridlines(ctx, W, H, padX, padY); + + const n = sigmas?.length || 0; + const xSteps = Math.max(totalSteps || n, n, deltas?.length || 0); + const xAt = i => padX + (i / Math.max(1, xSteps - 1)) * iW; + + let sYAt = null; + if (n > 1) { + let sMax = -Infinity, sMin = Infinity; + for (const s of sigmas) { if (s > sMax) sMax = s; if (s < sMin) sMin = s; } + if (sMin > 0) sMin = 0; + const sRange = Math.max(sMax - sMin, 1e-6); + sYAt = v => padY + (1 - (v - sMin) / sRange) * iH; + + ctx.strokeStyle = "rgba(208, 208, 208, 0.55)"; + ctx.lineWidth = 1; + ctx.setLineDash([3, 3]); + ctx.beginPath(); + for (let i = 0; i < n; i++) { + const px = padX + (i / (n - 1)) * iW; + const py = sYAt(sigmas[i]); + if (i === 0) ctx.moveTo(px, py); else ctx.lineTo(px, py); + } + ctx.stroke(); + ctx.setLineDash([]); + + const i = Math.max(0, Math.min(n - 1, step)); + const mx = padX + (i / Math.max(1, n - 1)) * iW; + const my = sYAt(sigmas[i]); + ctx.fillStyle = "#d0d0d0"; + ctx.beginPath(); + ctx.arc(mx, my, 2.5, 0, Math.PI * 2); + ctx.fill(); + } + + let dYAt = null; + if (deltas && deltas.length >= 1) { + let dMax = -Infinity; + for (const v of deltas) if (Number.isFinite(v) && v > dMax) dMax = v; + const dRange = Math.max(dMax, 1e-6); + dYAt = v => padY + (1 - v / dRange) * iH; + + // delta[i] is plotted at boundary (i+1); flat-extend delta[0] back to boundary 0. + const lastB = deltas.length; + ctx.beginPath(); + ctx.moveTo(xAt(0), H - padY); + ctx.lineTo(xAt(0), dYAt(deltas[0])); + for (let i = 0; i < deltas.length; i++) ctx.lineTo(xAt(i + 1), dYAt(deltas[i])); + ctx.lineTo(xAt(lastB), H - padY); + ctx.closePath(); + ctx.fillStyle = "rgba(230, 126, 34, 0.15)"; + ctx.fill(); + + ctx.strokeStyle = "#e67e22"; + ctx.lineWidth = 1.3; + if (deltas.length === 1) { + ctx.fillStyle = "#e67e22"; + ctx.beginPath(); + ctx.arc(xAt(1), dYAt(deltas[0]), 2, 0, Math.PI * 2); + ctx.fill(); + } else { + ctx.beginPath(); + for (let i = 0; i < deltas.length; i++) { + const px = xAt(i + 1), py = dYAt(deltas[i]); + if (i === 0) ctx.moveTo(px, py); else ctx.lineTo(px, py); + } + ctx.stroke(); + } + } + + // SamplerDetailBoost curve. `null` entries = outside the gate (drawn as faint baseline), + // finite entries = active (drawn as the bright cyan line, auto-scaled to its own peak). + if (Array.isArray(dbCurve) && dbCurve.length > 1) { + let dbMaxAbs = 0; + for (const v of dbCurve) if (Number.isFinite(v) && Math.abs(v) > dbMaxAbs) dbMaxAbs = Math.abs(v); + const dbXAt = i => padX + (i / Math.max(1, dbCurve.length - 1)) * iW; + const baselineY = H - padY - 0.5; + const isActive = v => Number.isFinite(v); + + ctx.strokeStyle = "rgba(120, 200, 220, 0.25)"; + ctx.lineWidth = 1; + ctx.beginPath(); + let inactiveOpen = false; + for (let i = 0; i < dbCurve.length; i++) { + if (!isActive(dbCurve[i])) { + if (!inactiveOpen) { ctx.moveTo(dbXAt(i), baselineY); inactiveOpen = true; } + else ctx.lineTo(dbXAt(i), baselineY); + } else { + inactiveOpen = false; + } + } + ctx.stroke(); + + if (dbMaxAbs > 1e-9) { + const dbYAt = v => padY + (1 - Math.abs(v) / dbMaxAbs) * iH; + ctx.strokeStyle = "rgba(120, 200, 220, 0.85)"; + ctx.fillStyle = "rgba(120, 200, 220, 0.85)"; + ctx.lineWidth = 1.2; + ctx.beginPath(); + let segStart = -1; + for (let i = 0; i <= dbCurve.length; i++) { + const v = i < dbCurve.length ? dbCurve[i] : null; + if (isActive(v)) { + if (segStart < 0) { segStart = i; ctx.moveTo(dbXAt(i), dbYAt(v)); } + else ctx.lineTo(dbXAt(i), dbYAt(v)); + } else if (segStart >= 0) { + if (i - segStart === 1) { + const v0 = dbCurve[segStart]; + ctx.fillRect(dbXAt(segStart) - 1, dbYAt(v0) - 1, 2, 2); + } + segStart = -1; + } + } + ctx.stroke(); + } + } + + if (lockedStep != null && lockedStep >= 0 && lockedStep < xSteps) { + const lx = xAt(lockedStep) + 0.5; + ctx.strokeStyle = "rgba(245, 200, 60, 0.9)"; + ctx.lineWidth = 1.2; + ctx.setLineDash([4, 2]); + ctx.beginPath(); + ctx.moveTo(lx, padY); + ctx.lineTo(lx, H - padY); + ctx.stroke(); + ctx.setLineDash([]); + } + + if (hoverStep != null && hoverStep >= 0 && hoverStep < xSteps) { + const hx = xAt(hoverStep) + 0.5; + ctx.strokeStyle = "rgba(208, 208, 208, 0.5)"; + ctx.lineWidth = 1; + ctx.beginPath(); + ctx.moveTo(hx, padY); + ctx.lineTo(hx, H - padY); + ctx.stroke(); + if (sYAt && hoverStep < n) { + ctx.fillStyle = "#d0d0d0"; + ctx.beginPath(); + ctx.arc(hx - 0.5, sYAt(sigmas[hoverStep]), 2.5, 0, Math.PI * 2); + ctx.fill(); + } + // delta[k-1] is plotted at boundary k. + if (dYAt && hoverStep >= 1 && (hoverStep - 1) < deltas.length) { + ctx.fillStyle = "#e67e22"; + ctx.beginPath(); + ctx.arc(hx - 0.5, dYAt(deltas[hoverStep - 1]), 2.5, 0, Math.PI * 2); + ctx.fill(); + } + } +} + +// totalSteps fixes the x-axis so the line grows left-to-right, not stretching to fill. +function drawLineGraph(canvas, values, totalSteps) { + const { ctx, W, H } = syncCanvasDPR(canvas); + const padX = GRAPH_PAD_X, padY = 3; + const iW = W - 2 * padX, iH = H - 2 * padY; + ctx.clearRect(0, 0, W, H); + drawGridlines(ctx, W, H, padX, padY); + if (!values || values.length < 1) return; + + let vMax = -Infinity, vMin = Infinity; + for (const v of values) { if (v > vMax) vMax = v; if (v < vMin) vMin = v; } + if (vMin > 0) vMin = 0; + const vRange = Math.max(vMax - vMin, 1e-6); + + const xSteps = Math.max(totalSteps || values.length, values.length); + + ctx.beginPath(); + ctx.moveTo(padX, H - padY); + for (let i = 0; i < values.length; i++) { + const px = padX + (i / Math.max(1, xSteps - 1)) * iW; + const py = padY + (1 - (values[i] - vMin) / vRange) * iH; + ctx.lineTo(px, py); + } + ctx.lineTo(padX + ((values.length - 1) / Math.max(1, xSteps - 1)) * iW, H - padY); + ctx.closePath(); + ctx.fillStyle = "rgba(230, 126, 34, 0.15)"; + ctx.fill(); + + ctx.strokeStyle = "#e67e22"; + ctx.lineWidth = 1.3; + ctx.beginPath(); + for (let i = 0; i < values.length; i++) { + const px = padX + (i / Math.max(1, xSteps - 1)) * iW; + const py = padY + (1 - (values[i] - vMin) / vRange) * iH; + if (i === 0) ctx.moveTo(px, py); else ctx.lineTo(px, py); + } + ctx.stroke(); +} + +app.registerExtension({ + name: "KJNodes.ModelPreviewOverride", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== "ModelPreviewOverrideKJ") return; + + chainCallback(nodeType.prototype, "onNodeCreated", function () { + ensureStyles(); + const node = this; + + const root = el("div", "kj-pov-root"); + + const imageArea = el("div", "kj-pov-image-area", root); + // Double-buffered: decode() on the visible-to-be element so the bitmap is reused. + const imgA = el("img", "kj-pov-img", imageArea); + const imgB = el("img", "kj-pov-img", imageArea); + imgA.draggable = false; + imgB.draggable = false; + imgB.style.opacity = "0"; + let visibleImg = imgA; + let pendingImg = imgB; + const img = imgA; // alias used by class-toggle code; applied to both buffers + // WebP path: ImageDecoder → VideoFrame[] → canvas, driven by a global timer. + const videoCanvas = el("canvas", "kj-pov-img", imageArea); + videoCanvas.style.opacity = "0"; + const videoCtx = videoCanvas.getContext("2d"); + // MP4 path: double-buffered