From 5eda61901a42d210bf4c3fc71e6caef4fcedb072 Mon Sep 17 00:00:00 2001 From: Satadal Dhara Date: Mon, 20 Jul 2026 16:03:08 +0530 Subject: [PATCH] v2.0: Architecture-aware upscaling, dynamic precision/tiling, GPU-bound frequency splitting, and Triton safety checks --- example_workflows/Sample_Workflow.json | 975 ++++++++++++++++++++++++- nodes/upscale_machine.py | 377 +++++----- 2 files changed, 1183 insertions(+), 169 deletions(-) diff --git a/example_workflows/Sample_Workflow.json b/example_workflows/Sample_Workflow.json index 7c51acc..5d2391d 100644 --- a/example_workflows/Sample_Workflow.json +++ b/example_workflows/Sample_Workflow.json @@ -1 +1,974 @@ -{"id":"84f8fb1b-dc2e-45b6-8c1e-0645a7cfc0d7","revision":0,"last_node_id":13,"last_link_id":17,"nodes":[{"id":7,"type":"VAEDecode","pos":[4687.470479278941,-139.45858388571898],"size":[140,46],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"samples","name":"samples","type":"LATENT","link":3},{"localized_name":"vae","name":"vae","type":"VAE","link":16}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[2]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"VAEDecode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":[]},{"id":10,"type":"EmptyLatentImage","pos":[3564.7126840943456,219.61366552815454],"size":[270,106],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":6},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":7},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[8]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"EmptyLatentImage","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":[512,512,1]},{"id":8,"type":"CLIPTextEncode","pos":[4146.234557539965,-115.43694680669059],"size":[400,200],"flags":{"collapsed":true},"order":7,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":14},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":9}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","links":[4]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"CLIPTextEncode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":[""],"color":"#232","bgcolor":"#353"},{"id":9,"type":"CLIPTextEncode","pos":[4146.230141291092,-69.40016084706781],"size":[400,200],"flags":{"collapsed":true},"order":8,"mode":0,"inputs":[{"localized_name":"clip","name":"clip","type":"CLIP","link":15},{"localized_name":"text","name":"text","type":"STRING","widget":{"name":"text"},"link":10}],"outputs":[{"localized_name":"CONDITIONING","name":"CONDITIONING","type":"CONDITIONING","links":[5]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"CLIPTextEncode","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":[""],"color":"#322","bgcolor":"#533"},{"id":5,"type":"CheckpointLoaderSimple","pos":[3506.160059451129,-198.51163106320556],"size":[270,98],"flags":{},"order":0,"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":[13]},{"localized_name":"CLIP","name":"CLIP","type":"CLIP","links":[14,15]},{"localized_name":"VAE","name":"VAE","type":"VAE","links":[16]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"CheckpointLoaderSimple","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["cyberrealistic_v70DMD2.safetensors"]},{"id":12,"type":"StringConcatenate","pos":[3964.8174368533214,-3.340835468014089],"size":[400,200],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"string_a","name":"string_a","type":"STRING","widget":{"name":"string_a"},"link":12},{"localized_name":"string_b","name":"string_b","type":"STRING","widget":{"name":"string_b"},"link":null},{"localized_name":"delimiter","name":"delimiter","type":"STRING","widget":{"name":"delimiter"},"link":null}],"outputs":[{"localized_name":"STRING","name":"STRING","type":"STRING","links":[10]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"StringConcatenate","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["","NSFW",""],"color":"#322","bgcolor":"#533"},{"id":2,"type":"Resolution_Machine","pos":[3255.2472676526163,244.1335282479651],"size":[270,150],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"resolution","name":"resolution","type":"COMBO","widget":{"name":"resolution"},"link":null},{"localized_name":"custom_width","name":"custom_width","type":"INT","widget":{"name":"custom_width"},"link":null},{"localized_name":"custom_height","name":"custom_height","type":"INT","widget":{"name":"custom_height"},"link":null}],"outputs":[{"localized_name":"width","name":"width","type":"INT","links":[6]},{"localized_name":"height","name":"height","type":"INT","links":[7]}],"properties":{"aux_id":"SatadalAI/SATA_UtilityNode","ver":"19e94ff1e0a534eeed7a8da88fd5c4fdbadb126b","Node name for S&R":"Resolution_Machine","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["SDXL_Illustrious_Pony","Best Portrait (13:19) 832x1216",832,1216]},{"id":4,"type":"Upscale_Machine","pos":[4854.531757397435,-141.42108428884757],"size":[270,154],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":2},{"localized_name":"upscale_model","name":"upscale_model","type":"COMBO","widget":{"name":"upscale_model"},"link":null},{"localized_name":"chained_model","name":"chained_model","type":"COMBO","widget":{"name":"chained_model"},"link":null},{"localized_name":"rescale_factor","name":"rescale_factor","type":"FLOAT","widget":{"name":"rescale_factor"},"link":null},{"localized_name":"supersample","name":"supersample","type":"COMBO","widget":{"name":"supersample"},"link":null},{"localized_name":"rounding_modulus","name":"rounding_modulus","type":"INT","widget":{"name":"rounding_modulus"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[1]}],"properties":{"aux_id":"SatadalAI/SATA_UtilityNode","ver":"19e94ff1e0a534eeed7a8da88fd5c4fdbadb126b","Node name for S&R":"Upscale_Machine","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["4x-UltraSharpV2.safetensors","None",1.5,"true",8]},{"id":6,"type":"KSampler","pos":[4389.953637891998,-139.45860297634326],"size":[270,262],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"model","name":"model","type":"MODEL","link":13},{"localized_name":"positive","name":"positive","type":"CONDITIONING","link":4},{"localized_name":"negative","name":"negative","type":"CONDITIONING","link":5},{"localized_name":"latent_image","name":"latent_image","type":"LATENT","link":8},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"localized_name":"steps","name":"steps","type":"INT","widget":{"name":"steps"},"link":null},{"localized_name":"cfg","name":"cfg","type":"FLOAT","widget":{"name":"cfg"},"link":null},{"localized_name":"sampler_name","name":"sampler_name","type":"COMBO","widget":{"name":"sampler_name"},"link":null},{"localized_name":"scheduler","name":"scheduler","type":"COMBO","widget":{"name":"scheduler"},"link":null},{"localized_name":"denoise","name":"denoise","type":"FLOAT","widget":{"name":"denoise"},"link":null}],"outputs":[{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":[3]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"KSampler","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":[459887911092090,"randomize",8,1.5,"lcm","exponential",1]},{"id":3,"type":"ImageSaveWithMetadata","pos":[5163.316100152279,-140.23340376034693],"size":[416.464832943705,471.03115416393564],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":1},{"localized_name":"path_and_filename","name":"path_and_filename","type":"STRING","widget":{"name":"path_and_filename"},"link":null},{"localized_name":"extension","name":"extension","type":"COMBO","widget":{"name":"extension"},"link":null},{"localized_name":"custom_string","name":"custom_string","shape":7,"type":"STRING","widget":{"name":"custom_string"},"link":null},{"localized_name":"lossless_webp","name":"lossless_webp","shape":7,"type":"BOOLEAN","widget":{"name":"lossless_webp"},"link":null},{"localized_name":"quality_jpeg_or_webp","name":"quality_jpeg_or_webp","shape":7,"type":"INT","widget":{"name":"quality_jpeg_or_webp"},"link":null}],"outputs":[],"properties":{"aux_id":"SatadalAI/SATA_UtilityNode","ver":"19e94ff1e0a534eeed7a8da88fd5c4fdbadb126b","Node name for S&R":"ImageSaveWithMetadata","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["%time","webp","",true,100]},{"id":11,"type":"StringConcatenate","pos":[3531.432838773623,-27.357903573158715],"size":[400,200],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"string_a","name":"string_a","type":"STRING","widget":{"name":"string_a"},"link":11},{"localized_name":"string_b","name":"string_b","type":"STRING","widget":{"name":"string_b"},"link":null},{"localized_name":"delimiter","name":"delimiter","type":"STRING","widget":{"name":"delimiter"},"link":null}],"outputs":[{"localized_name":"STRING","name":"STRING","type":"STRING","links":[9]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"StringConcatenate","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["","rapunzel_\\(disney\\) cosplay, done by Indian girl,",""],"color":"#232","bgcolor":"#353"},{"id":13,"type":"ShowText|pysssss","pos":[3209.8748108791883,-205.30488323760082],"size":[269.0952238525533,122.84906196257143],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","link":17}],"outputs":[{"localized_name":"STRING","name":"STRING","shape":6,"type":"STRING","links":null}],"properties":{"cnr_id":"comfyui-custom-scripts","ver":"f2838ed5e59de4d73cde5c98354b87a8d3200190","Node name for S&R":"ShowText|pysssss","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["sd1.5"]},{"id":1,"type":"Prompt_Machine","pos":[3205.7149982758856,5.5565136185743595],"size":[270,122],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"csv_file","name":"csv_file","type":"COMBO","widget":{"name":"csv_file"},"link":null},{"localized_name":"name","name":"name","type":"STRING","widget":{"name":"name"},"link":null}],"outputs":[{"localized_name":"positive","name":"positive","type":"STRING","links":[11]},{"localized_name":"negative","name":"negative","type":"STRING","links":[12]},{"localized_name":"note","name":"note","type":"STRING","links":[17]}],"properties":{"aux_id":"SatadalAI/SATA_UtilityNode","ver":"19e94ff1e0a534eeed7a8da88fd5c4fdbadb126b","Node name for S&R":"Prompt_Machine","ue_properties":{"widget_ue_connectable":{},"input_ue_unconnectable":{},"version":"7.4.1"}},"widgets_values":["models.csv","RealVis"]}],"links":[[1,4,0,3,0,"IMAGE"],[2,7,0,4,0,"IMAGE"],[3,6,0,7,0,"LATENT"],[4,8,0,6,1,"CONDITIONING"],[5,9,0,6,2,"CONDITIONING"],[6,2,0,10,0,"INT"],[7,2,1,10,1,"INT"],[8,10,0,6,3,"LATENT"],[9,11,0,8,1,"STRING"],[10,12,0,9,1,"STRING"],[11,1,0,11,0,"STRING"],[12,1,1,12,0,"STRING"],[13,5,0,6,0,"MODEL"],[14,5,1,8,0,"CLIP"],[15,5,1,9,0,"CLIP"],[16,5,2,7,1,"VAE"],[17,1,2,13,0,"STRING"]],"groups":[],"config":{},"extra":{"ue_links":[],"links_added_by_ue":[],"ds":{"scale":0.7684248783886432,"offset":[-2849.709530002181,361.5277939468702]}},"version":0.4} \ No newline at end of file +{ + "id": "4efb5c6c-dfe3-4cf4-ac1b-b0004714bfd4", + "revision": 0, + "last_node_id": 68, + "last_link_id": 117, + "nodes": [ + { + "id": 14, + "type": "Save_Machine", + "pos": [ + 5152.081565225479, + -142.37814847211683 + ], + "size": [ + 417.31616451357604, + 503.8613833815003 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 18 + } + ], + "outputs": [ + { + "name": "file_path", + "type": "STRING", + "links": null + } + ], + "properties": { + "Node name for S&R": "Save_Machine", + "aux_id": "SatadalAI/SATA_UtilityNode", + "ver": "b4cd890c9e599b6f01361acd582662483a9f7189", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.5.2" + } + }, + "widgets_values": [ + "wallpaper/%time", + "webp" + ] + }, + { + "id": 63, + "type": "UNETLoader", + "pos": [ + 3261.5444578237834, + -108.89048389712944 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 111 + ] + } + ], + "properties": { + "Node name for S&R": "UNETLoader", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.8" + } + }, + "widgets_values": [ + "z_image_turbo_int8_convrot.safetensors", + "default" + ] + }, + { + "id": 62, + "type": "CLIPLoader", + "pos": [ + 3268.220695664241, + 18.860210383864736 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 110 + ] + } + ], + "properties": { + "Node name for S&R": "CLIPLoader", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.8" + } + }, + "widgets_values": [ + "qwen_3_4b_int8_convrot.safetensors", + "lumina2", + "default" + ] + }, + { + "id": 21, + "type": "VAELoader", + "pos": [ + 3265.231337729953, + 178.9079538641949 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 38 + ] + } + ], + "properties": { + "Node name for S&R": "VAELoader", + "cnr_id": "comfy-core", + "ver": "0.3.73", + "ue_properties": { + "widget_ue_connectable": {}, + "version": "7.5.2", + "input_ue_unconnectable": {} + }, + "models": [ + { + "name": "ae.safetensors", + "url": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/vae/ae.safetensors", + "directory": "vae" + } + ], + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "Zae.safetensors" + ] + }, + { + "id": 26, + "type": "ModelSamplingAuraFlow", + "pos": [ + 3594.6013142917627, + -107.54723647282863 + ], + "size": [ + 210, + 79.62649507714298 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 111 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 33 + ] + } + ], + "properties": { + "Node name for S&R": "ModelSamplingAuraFlow", + "cnr_id": "comfy-core", + "ver": "0.3.64", + "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": [ + 3 + ] + }, + { + "id": 11, + "type": "StringConcatenate", + "pos": [ + 3610.855545233338, + 20.721433684519976 + ], + "size": [ + 407.76454013176954, + 166 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "string_a", + "type": "STRING", + "widget": { + "name": "string_a" + }, + "link": 116 + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 51 + ] + } + ], + "properties": { + "Node name for S&R": "StringConcatenate", + "cnr_id": "comfy-core", + "ver": "0.3.68", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.4.1" + } + }, + "widgets_values": [ + "", + "full body, high quality, 8k, hyper detailed, award winning, masterpiece, highest quality, facing camera, very detailed skin, skin texture,", + "," + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 24, + "type": "CLIPTextEncode", + "pos": [ + 3817.204506955194, + -109.01087387481428 + ], + "size": [ + 218.659765625, + 88 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 110 + }, + { + "name": "text", + "type": "STRING", + "widget": { + "name": "text" + }, + "link": 51 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 35, + 36 + ] + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode", + "cnr_id": "comfy-core", + "ver": "0.3.73", + "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": 17, + "type": "ConditioningZeroOut", + "pos": [ + 4091.6785047538547, + -146.51380029246832 + ], + "size": [ + 204.134765625, + 26 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 36 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 37 + ] + } + ], + "properties": { + "Node name for S&R": "ConditioningZeroOut", + "cnr_id": "comfy-core", + "ver": "0.3.73", + "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": 65, + "type": "Latent_Machine", + "pos": [ + 4096.016727994729, + -55.31161363152323 + ], + "size": [ + 210, + 342 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": 113 + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": 114 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 115 + ] + }, + { + "name": "IMAGE", + "type": "IMAGE", + "links": null + } + ], + "properties": { + "Node name for S&R": "Latent_Machine", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.8" + } + }, + "widgets_values": [ + "Z-Image Family", + "Landscape", + "Custom", + "SD3", + "Gaussian (White): Sharp Architecture, Text, intricate mechanics", + 1, + "Pink (1/f): Photorealism, Portraits, Nature", + 1, + 1085431197464353, + "randomize", + true, + null + ] + }, + { + "id": 2, + "type": "Resolution_Machine", + "pos": [ + 3821.546799515842, + 301.92650724760136 + ], + "size": [ + 210, + 222 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "width", + "type": "INT", + "links": [ + 113 + ] + }, + { + "name": "height", + "type": "INT", + "links": [ + 114 + ] + } + ], + "properties": { + "Node name for S&R": "Resolution_Machine", + "aux_id": "SatadalAI/SATA_UtilityNode", + "ver": "19e94ff1e0a534eeed7a8da88fd5c4fdbadb126b", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.4.1" + } + }, + "widgets_values": [ + "Z-Image Family", + "Landscape", + "Custom", + 1344, + 768, + "None", + null + ] + }, + { + "id": 6, + "type": "KSampler", + "pos": [ + 4340.334703541104, + -149.75690572490768 + ], + "size": [ + 270, + 262 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 33 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 35 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 37 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 115 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 3 + ] + } + ], + "properties": { + "Node name for S&R": "KSampler", + "cnr_id": "comfy-core", + "ver": "0.3.68", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.4.1" + } + }, + "widgets_values": [ + 128626698311720, + "randomize", + 9, + 1, + "res_multistep", + "simple", + 1 + ] + }, + { + "id": 7, + "type": "VAEDecode", + "pos": [ + 4644.405099597037, + -149.75683306415954 + ], + "size": [ + 140, + 46 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 3 + }, + { + "name": "vae", + "type": "VAE", + "link": 38 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 2, + 112 + ] + } + ], + "properties": { + "Node name for S&R": "VAEDecode", + "cnr_id": "comfy-core", + "ver": "0.3.68", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.4.1" + } + }, + "widgets_values": [] + }, + { + "id": 64, + "type": "Preview_Machine", + "pos": [ + 4798.8673871537585, + 72.24708306565708 + ], + "size": [ + 322.0207042694277, + 268.30575207627953 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 112 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "Preview_Machine", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.8" + } + }, + "widgets_values": [] + }, + { + "id": 67, + "type": "Prompt_Machine", + "pos": [ + 2821.958194903464, + -140.59048020621123 + ], + "size": [ + 400, + 382 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "technical", + "type": "STRING", + "widget": { + "name": "technical" + }, + "link": 117 + } + ], + "outputs": [ + { + "name": "prompt", + "type": "STRING", + "links": [ + 116 + ] + } + ], + "properties": { + "Node name for S&R": "Prompt_Machine", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.8" + } + }, + "widgets_values": [ + "a lady walking down the street with her little daughter in a rainy day. both are wearing multicolored dress.", + "a modern photograph of a rainy day.", + "gloomy daylight.", + "the little girl has jumped to a stagnent water. mother is holding one hand of her daughter. mother has a side bag on her shoulder and a closed umbrella holding in her hand. ", + "mother has smiling expression on her face and daughter has sprakling eyes with a cute smile.", + "", + false + ] + }, + { + "id": 68, + "type": "Prompt_Style_Machine", + "pos": [ + 2565.999901970176, + -69.93013216192438 + ], + "size": [ + 210, + 136 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "prompt", + "type": "STRING", + "links": [ + 117 + ] + } + ], + "properties": { + "Node name for S&R": "Prompt_Style_Machine", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.8" + } + }, + "widgets_values": [ + "styles.csv", + "Chibi", + "No preview available" + ] + }, + { + "id": 4, + "type": "Upscale_Machine", + "pos": [ + 4829.254230193369, + -149.8469326424389 + ], + "size": [ + 270, + 154 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 2 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 18 + ] + } + ], + "properties": { + "Node name for S&R": "Upscale_Machine", + "aux_id": "SatadalAI/SATA_UtilityNode", + "ver": "19e94ff1e0a534eeed7a8da88fd5c4fdbadb126b", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.4.1" + } + }, + "widgets_values": [ + "4xNomos2_realplksr_dysample.pth", + "None", + 1.5, + "true" + ] + } + ], + "links": [ + [ + 2, + 7, + 0, + 4, + 0, + "IMAGE" + ], + [ + 3, + 6, + 0, + 7, + 0, + "LATENT" + ], + [ + 18, + 4, + 0, + 14, + 0, + "IMAGE" + ], + [ + 33, + 26, + 0, + 6, + 0, + "MODEL" + ], + [ + 35, + 24, + 0, + 6, + 1, + "CONDITIONING" + ], + [ + 36, + 24, + 0, + 17, + 0, + "CONDITIONING" + ], + [ + 37, + 17, + 0, + 6, + 2, + "CONDITIONING" + ], + [ + 38, + 21, + 0, + 7, + 1, + "VAE" + ], + [ + 51, + 11, + 0, + 24, + 1, + "STRING" + ], + [ + 110, + 62, + 0, + 24, + 0, + "CLIP" + ], + [ + 111, + 63, + 0, + 26, + 0, + "MODEL" + ], + [ + 112, + 7, + 0, + 64, + 0, + "IMAGE" + ], + [ + 113, + 2, + 0, + 65, + 0, + "INT" + ], + [ + 114, + 2, + 1, + 65, + 1, + "INT" + ], + [ + 115, + 65, + 0, + 6, + 3, + "LATENT" + ], + [ + 116, + 67, + 0, + 11, + 0, + "STRING" + ], + [ + 117, + 68, + 0, + 67, + 0, + "STRING" + ] + ], + "groups": [ + { + "id": 2, + "title": "Step3 - Prompt", + "bounding": [ + 3580.3970390097234, + -182.48227215100252, + 452.36970022695004, + 424.24216158491276 + ], + "color": "#3f789e", + "flags": {} + }, + { + "id": 3, + "title": "Step1 - Load models", + "bounding": [ + 3244.0164909689006, + -181.3559554697695, + 311.11840937274474, + 431.3015218159521 + ], + "color": "#3f789e", + "flags": {} + } + ], + "config": {}, + "extra": { + "ue_links": [], + "links_added_by_ue": [], + "workflowRendererVersion": "LG", + "ds": { + "scale": 0.7199453959465028, + "offset": [ + -2969.6857606137182, + 381.63065761713165 + ] + }, + "frontendVersion": "1.45.21", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes/upscale_machine.py b/nodes/upscale_machine.py index 4d9d71a..cecb8df 100644 --- a/nodes/upscale_machine.py +++ b/nodes/upscale_machine.py @@ -1,7 +1,5 @@ -import os import torch -import cv2 -import numpy as np +import torch.nn.functional as F import folder_paths import comfy.utils from comfy import model_management @@ -42,45 +40,90 @@ def generate_blue_noise(batch_size, c, h, w, device, beta=1.5): return blue_noise -def resize_tensor_opencv(tensor_chw, target_width, target_height, supersample='true', factor=2.0): +def _make_gaussian_kernel(kernel_size, sigma, device): + """Build a separable 2D Gaussian kernel tensor (1,1,k,k) on device.""" + k = kernel_size + coords = torch.arange(k, dtype=torch.float32, device=device) - k // 2 + gauss_1d = torch.exp(-0.5 * (coords / sigma) ** 2) + gauss_1d = gauss_1d / gauss_1d.sum() + kernel_2d = gauss_1d[:, None] * gauss_1d[None, :] + return kernel_2d.view(1, 1, k, k) + + +def fast_gaussian_blur_bchw(tensor_bchw, kernel_size=9, sigma=3.0): """ - CPU-based resizing using OpenCV for 'rescale' mode. - Input: tensor_chw -> torch.FloatTensor (C,H,W) in [0,1] - Output: torch.FloatTensor (C,H,W) in [0,1], guaranteed 3-channel RGB + Fast per-channel GPU Gaussian blur via depthwise convolution. + Input/Output: BCHW float tensor. Kernel is always created on the same device. """ - new_width = max(1, int(target_width)) - new_height = max(1, int(target_height)) + device = tensor_bchw.device + k = kernel_size + kernel = _make_gaussian_kernel(k, sigma, device) # on same device + C = tensor_bchw.shape[1] + kernel_c = kernel.expand(C, 1, k, k).to(device) # guarantee same device + pad = k // 2 + return F.conv2d(tensor_bchw.contiguous(), kernel_c, padding=pad, groups=C) - interp = cv2.INTER_LANCZOS4 if factor > 1.0 else cv2.INTER_AREA - np_img = tensor_chw.permute(1, 2, 0).detach().cpu().numpy() - np_img = (np.clip(np_img, 0.0, 1.0) * 255.0).astype(np.uint8) +def resize_bchw(tensor_bchw, target_h, target_w): + """ + Fast GPU resize — stays on device, no antialias overhead. + Bilinear for downscale (smoother), bicubic for upscale (sharper). + """ + in_h, in_w = tensor_bchw.shape[2], tensor_bchw.shape[3] + if in_h == target_h and in_w == target_w: + return tensor_bchw + downscaling = (target_h < in_h) or (target_w < in_w) + mode = 'bilinear' if downscaling else 'bicubic' + return F.interpolate(tensor_bchw, size=(target_h, target_w), mode=mode, align_corners=False) - if np_img.ndim == 2: - np_img = cv2.cvtColor(np_img, cv2.COLOR_GRAY2RGB) - elif np_img.shape[2] == 1: - np_img = cv2.cvtColor(np_img, cv2.COLOR_GRAY2RGB) - elif np_img.shape[2] == 4: - np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB) - elif np_img.shape[2] > 3: - np_img = np_img[:, :, :3] - if np_img.shape[0] == 0 or np_img.shape[1] == 0: - raise ValueError("Input image for OpenCV resize has zero width or height.") +# ───────────────────────────────────────────────────────────────────────────── +# Architecture optimization profiles for spandrel models. +# Keys match the inner nn.Module class name (type(model.model).__name__). +# +# tile : starting tile size for tiled inference (OOM fallback halves it) +# fp16_model : cast model weights to FP16 before inference (Tensor Core boost) +# compile : torch.compile mode — 'reduce-overhead' for fixed-shape tiling +# ───────────────────────────────────────────────────────────────────────────── +ARCH_PROFILES = { + # ── Fast CNN / RRDB models — full FP16 support, large tiles ────────────── + "RRDBNet": {"tile": 768, "fp16_model": True, "compile": "reduce-overhead"}, # ESRGAN, Real-ESRGAN + "SRVGGNetCompact": {"tile": 896, "fp16_model": True, "compile": "reduce-overhead"}, # Real-ESRGAN Compact + "SPAN": {"tile": 896, "fp16_model": True, "compile": "reduce-overhead"}, # Swift, very fast + "PLKSR": {"tile": 768, "fp16_model": False, "compile": "reduce-overhead"}, # Lightweight + "RealPLKSR": {"tile": 768, "fp16_model": False, "compile": "reduce-overhead"}, # Lightweight photo + "RCAN": {"tile": 768, "fp16_model": True, "compile": "reduce-overhead"}, # Residual channel attn + "SAFMN": {"tile": 896, "fp16_model": True, "compile": "reduce-overhead"}, # Very fast + "DITN": {"tile": 768, "fp16_model": True, "compile": "reduce-overhead"}, + "CRAFT": {"tile": 768, "fp16_model": True, "compile": "reduce-overhead"}, + "NAFNet": {"tile": 768, "fp16_model": True, "compile": "reduce-overhead"}, + "OmniSR": {"tile": 640, "fp16_model": True, "compile": "reduce-overhead"}, + "RGT": {"tile": 640, "fp16_model": True, "compile": "reduce-overhead"}, + # ── Transformer / attention models — autocast only, conservative tiles ──── + "DAT": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, # Dual Attention Transformer + "DAT_S": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, + "HAT": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, # Hybrid Attention + "HAT_L": {"tile": 480, "fp16_model": False, "compile": "reduce-overhead"}, # HAT Large — very heavy + "SwinIR": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, # Swin Transformer IR + "Swin2SR": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, + "SAN": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, # Second-order Attn + "DRCT": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, + "ATD": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, + "GRL": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, + "FDAT": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, + "AuraSR": {"tile": 512, "fp16_model": False, "compile": "reduce-overhead"}, +} - if supersample == 'true': - ss_width = max(1, new_width * 8) - ss_height = max(1, new_height * 8) - np_img = cv2.resize(np_img, (ss_width, ss_height), interpolation=interp) - if np_img.shape[0] == 0 or np_img.shape[1] == 0: - raise ValueError("Image became empty after supersample resize.") - - resized = cv2.resize(np_img, (new_width, new_height), interpolation=interp) - tensor_out = torch.from_numpy(resized).float().div(255.0).permute(2, 0, 1) - return tensor_out +# Fallback profile for any architecture not listed above. +# Uses supports_half from spandrel for fp16_model, 512 tile as safe default. +_DEFAULT_PROFILE = {"tile": 512, "fp16_model": None, "compile": "reduce-overhead"} class Upscale_Machine: + # Session-level cache: model object id -> torch.compile'd inner module + # Avoids recompiling the same model on repeated runs in the same session. + _compiled_cache: dict = {} + @classmethod def INPUT_TYPES(cls): return { @@ -99,10 +142,7 @@ class Upscale_Machine: CATEGORY = "SATA_UtilityNode" def load_model(self, model_name): - """Load ESRGAN/RealESRGAN/AESRGAN from the project's upscale_models folder. - - Only legacy PyTorch/spandrel models are supported by this simplified node. - """ + """Load ESRGAN/RealESRGAN/AESRGAN from the project's upscale_models folder.""" if not model_name: raise ValueError("No upscale model selected or provided.") @@ -116,40 +156,109 @@ class Upscale_Machine: model.eval() return model - def upscale_with_model(self, upscale_model, image_bhwc, pbar=None): + def _get_arch_profile(self, upscale_model): """ - Run tiled upscale with spandrel model. - Input: image_bhwc -> (B,H,W,C) float [0,1] - Output: (B,C,H',W') float [0,1] + Detects the model architecture and returns its optimisation profile. + Detection priority: + 1. Inner nn.Module class name (most specific) + 2. spandrel supports_half flag (fallback for unknown architectures) """ - device = model_management.get_torch_device() + inner_cls = type(upscale_model.model).__name__ + profile = ARCH_PROFILES.get(inner_cls) + if profile is None: + # Unknown architecture — use spandrel flag for fp16 decision + fp16_ok = getattr(upscale_model, "supports_half", False) + profile = {**_DEFAULT_PROFILE, "fp16_model": fp16_ok} + return inner_cls, profile + + def _get_compiled_model(self, upscale_model, compile_mode): + """ + Returns a torch.compile()'d version of the model's inner nn.Module. + Cached by object id — compiles once per session per loaded model. + Falls back silently to eager mode if compile fails or Triton is missing. + """ + inner = upscale_model.model + cache_key = id(inner) + if cache_key not in Upscale_Machine._compiled_cache: + # Check for triton before attempting compile, as torch.compile is lazy + # and will otherwise crash on the first tile execution on Windows. + try: + import triton + has_triton = True + except ImportError: + has_triton = False + + if has_triton: + try: + compiled = torch.compile(inner, mode=compile_mode, dynamic=False) + Upscale_Machine._compiled_cache[cache_key] = compiled + except Exception as ex: + print(f"[Upscale_Machine] torch.compile unavailable ({ex}), using eager mode.") + Upscale_Machine._compiled_cache[cache_key] = inner + else: + print("[Upscale_Machine] Triton not installed (Windows), skipping torch.compile.") + Upscale_Machine._compiled_cache[cache_key] = inner + + return Upscale_Machine._compiled_cache[cache_key] + + def upscale_with_model(self, upscale_model, image_bchw, device, pbar=None): + """ + Architecture-aware tiled upscale: + - Detects model type and loads optimisation profile + - FP16 weights (Tensor Cores) for CNN models that support it + - torch.autocast for mixed-precision on every tile (all architectures) + - torch.compile fused graph (cached per session, falls back safely) + Input: image_bchw -> (B,C,H,W) float32 [0,1] + Output: (B,C,H',W') float32 [0,1] + """ + arch_name, profile = self._get_arch_profile(upscale_model) + spandrel_supports_half = getattr(upscale_model, "supports_half", False) + use_fp16 = profile["fp16_model"] and spandrel_supports_half + tile = profile["tile"] + comp_mode = profile["compile"] + device_type = device.type if hasattr(device, "type") else str(device).split(":")[0] + + print(f"[Upscale_Machine] {arch_name} | " + f"tile={tile} | fp16={'yes' if use_fp16 else 'no'} | compile={comp_mode}") + upscale_model.to(device) + if use_fp16: + upscale_model.half() + in_tensor = image_bchw.to(device).half() + else: + in_tensor = image_bchw.to(device).float() + + # Compile inner nn.Module once per session + compiled_fn = self._get_compiled_model(upscale_model, comp_mode) try: - in_img = image_bhwc.movedim(-1, -3).to(device) - - tile = 512 overlap = 32 oom = True while oom: try: if not pbar: - steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps( - in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap + steps = in_tensor.shape[0] * comfy.utils.get_tiled_scale_steps( + in_tensor.shape[3], in_tensor.shape[2], tile_x=tile, tile_y=tile, overlap=overlap ) local_pbar = comfy.utils.ProgressBar(steps) else: local_pbar = pbar - s = comfy.utils.tiled_scale( - in_img, - lambda a: upscale_model(a), - tile_x=tile, - tile_y=tile, - overlap=overlap, - upscale_amount=getattr(upscale_model, "scale", 4), - pbar=local_pbar - ) + import contextlib + # If model doesn't support fp16 (like DAT/HAT), autocast to fp16 will also cause NaNs. + # We only autocast if use_fp16 is True to be perfectly safe. + ctx = torch.autocast(device_type=device_type, dtype=torch.float16) if use_fp16 else contextlib.nullcontext() + + with ctx: + s = comfy.utils.tiled_scale( + in_tensor, + lambda a: compiled_fn(a), + tile_x=tile, + tile_y=tile, + overlap=overlap, + upscale_amount=getattr(upscale_model, "scale", 4), + pbar=local_pbar + ) oom = False except model_management.OOM_EXCEPTION as e: model_management.soft_empty_cache() @@ -157,10 +266,9 @@ class Upscale_Machine: if tile < 128: raise e finally: - upscale_model.cpu() + upscale_model.cpu().float() # Always restore to FP32 on CPU - s = torch.clamp(s, min=0.0, max=1.0) - return s + return torch.clamp(s.float(), min=0.0, max=1.0) def _round_to_modulus(self, value, modulus): if modulus is None or modulus <= 1: @@ -173,127 +281,60 @@ class Upscale_Machine: if image.ndim != 4: raise ValueError("Expected IMAGE tensor with 4 dims (B,H,W,C).") + device = model_management.get_torch_device() + original_height = int(image.shape[1]) original_width = int(image.shape[2]) - target_width = self._round_to_modulus(original_width * rescale_factor, rounding_modulus) - target_height = self._round_to_modulus(original_height * rescale_factor, rounding_modulus) + target_w = self._round_to_modulus(original_width * rescale_factor, rounding_modulus) + target_h = self._round_to_modulus(original_height * rescale_factor, rounding_modulus) - current_bhwc = image - - # Calculate total steps for progress bar - total_steps = 0 - if upscale_model: - total_steps += image.shape[0] * comfy.utils.get_tiled_scale_steps( - image.shape[2], image.shape[1], tile_x=512, tile_y=512, overlap=32 - ) - if chained_model and chained_model != "None": - # Rough estimate: if chained model runs on target_width/height after intermediate downscale - total_steps += image.shape[0] * comfy.utils.get_tiled_scale_steps( - target_width, target_height, tile_x=512, tile_y=512, overlap=32 - ) + # Move to GPU once as BCHW — stay there the entire pipeline + current_bchw = image.movedim(-1, 1).to(device) + if current_bchw.shape[1] == 1: + current_bchw = current_bchw.repeat(1, 3, 1, 1) + elif current_bchw.shape[1] > 3: + current_bchw = current_bchw[:, :3, :, :] - pbar = comfy.utils.ProgressBar(total_steps) if total_steps > 0 else None + # Keep a clean GPU copy of original for frequency split baseline + original_bchw = current_bchw.clone() - # First upscale + # ── First upscale model ─────────────────────────────────────────────── if upscale_model: up_model = self.load_model(upscale_model) - up_bchw = self.upscale_with_model(up_model, current_bhwc, pbar) - current_bhwc = up_bchw.movedim(1, -1).contiguous() + current_bchw = self.upscale_with_model(up_model, current_bchw, device) + # tiled_scale may return CPU tensor — pin back to GPU + current_bchw = current_bchw.to(device) + current_bchw = resize_bchw(current_bchw, target_h, target_w) - # Downscale to target size before chaining - batch_downscaled = [] - for i in range(current_bhwc.shape[0]): - chw = current_bhwc[i].movedim(-1, 0) # (H,W,C) -> (C,H,W) - chw = torch.clamp(chw, 0.0, 1.0) - resized_chw = resize_tensor_opencv( - chw, target_width=target_width, target_height=target_height, - supersample=supersample, factor=rescale_factor - ) - if resized_chw.shape[0] == 1: - resized_chw = resized_chw.repeat(3, 1, 1) - elif resized_chw.shape[0] > 3: - resized_chw = resized_chw[:3, :, :] - batch_downscaled.append(resized_chw.unsqueeze(0)) - current_bhwc = torch.cat(batch_downscaled, dim=0) - current_bhwc = current_bhwc.movedim(1, -1).contiguous() - - # Second upscale (chained model) if selected - if chained_model and chained_model != "None": - # chained model is a model name + # ── Chained model ───────────────────────────────────────────────────── + has_chain = chained_model and chained_model != "None" + if has_chain: chain_model = self.load_model(chained_model) - chain_bchw = self.upscale_with_model(chain_model, current_bhwc, pbar) - current_bhwc = chain_bchw.movedim(1, -1).contiguous() + current_bchw = self.upscale_with_model(chain_model, current_bchw, device) + # tiled_scale may return CPU tensor — pin back to GPU + current_bchw = current_bchw.to(device) + current_bchw = resize_bchw(current_bchw, target_h, target_w) - # If no chained model, ensure output is at target size - if not (chained_model and chained_model != "None"): - # Already downscaled after first upscale, so just output - images_out = current_bhwc - else: - # After chained model, do a final resize to target size - batch_out = [] - for i in range(current_bhwc.shape[0]): - chw = current_bhwc[i].movedim(-1, 0) - chw = torch.clamp(chw, 0.0, 1.0) - resized_chw = resize_tensor_opencv( - chw, target_width=target_width, target_height=target_height, - supersample=supersample, factor=rescale_factor - ) - if resized_chw.shape[0] == 1: - resized_chw = resized_chw.repeat(3, 1, 1) - elif resized_chw.shape[0] > 3: - resized_chw = resized_chw[:3, :, :] - batch_out.append(resized_chw.unsqueeze(0)) - out_bchw = torch.cat(batch_out, dim=0) - images_out = out_bchw.movedim(1, -1).contiguous() + images_out = current_bchw - # Apply Frequency-Split SR if enabled + # ── Frequency-Split SR ──────────────────────────────────────────────── if frequency_split and upscale_model: - import torchvision.transforms.functional as TF - - # Upscale original image to target resolution using basic bicubic - batch_base = [] - for i in range(image.shape[0]): - chw = image[i].movedim(-1, 0) - chw = torch.clamp(chw, 0.0, 1.0) - resized_chw = resize_tensor_opencv( - chw, target_width=target_width, target_height=target_height, - supersample='false', factor=rescale_factor - ) - if resized_chw.shape[0] == 1: - resized_chw = resized_chw.repeat(3, 1, 1) - elif resized_chw.shape[0] > 3: - resized_chw = resized_chw[:3, :, :] - batch_base.append(resized_chw.unsqueeze(0)) - - bicubic_out = torch.cat(batch_base, dim=0).movedim(1, -1).contiguous() - - def blur_bhwc(tensor_bhwc): - tensor_bchw = tensor_bhwc.movedim(-1, 1) - h, w = tensor_bchw.shape[-2:] - k_y = min(15, h if h % 2 != 0 else h - 1) - k_x = min(15, w if w % 2 != 0 else w - 1) - k_y = max(1, k_y) - k_x = max(1, k_x) - if k_y >= 3 and k_x >= 3: - blurred = TF.gaussian_blur(tensor_bchw, [k_y, k_x], [3.0, 3.0]) - else: - blurred = tensor_bchw - return blurred.movedim(1, -1) - - sr_blurred = blur_bhwc(images_out) - sr_high_freq = images_out - sr_blurred - - bicubic_blurred = blur_bhwc(bicubic_out) - images_out = bicubic_blurred + sr_high_freq - images_out = torch.clamp(images_out, 0.0, 1.0) + # Ensure both tensors are on the same device before any arithmetic + bicubic_bchw = resize_bchw(original_bchw.to(device), target_h, target_w) + images_out = images_out.to(device) - if chained_model and chained_model != "None": - inject_noise_for_realism = 0.15 - B, H, W, C = images_out.shape - device = images_out.device - noise = generate_blue_noise(B, C, H, W, device) # shape (B, C, H, W) - noise = noise.movedim(1, -1) # shape (B, H, W, C) - images_out = images_out + noise * inject_noise_for_realism - images_out = torch.clamp(images_out, 0.0, 1.0) + sr_low = fast_gaussian_blur_bchw(images_out) + sr_high = images_out - sr_low - return (images_out,) \ No newline at end of file + bic_low = fast_gaussian_blur_bchw(bicubic_bchw) + images_out = torch.clamp(bic_low + sr_high, 0.0, 1.0) + + # ── Blue Noise realism (chained only) ───────────────────────────────── + if has_chain: + images_out = images_out.to(device) + B, C, H, W = images_out.shape + noise = generate_blue_noise(B, C, H, W, device) + images_out = torch.clamp(images_out + noise * 0.15, 0.0, 1.0) + + # Convert back to BHWC for ComfyUI + return (images_out.movedim(1, -1).contiguous(),) \ No newline at end of file