diff --git a/README.md b/README.md index 950635e..9a8ae93 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ # 🔥 Available now: Super Pop Color Adjustment Model & Nodes

- SuperBeasts.AI Super Pop Color Adjustment + SuperBeasts.AI Super Pop Color Adjustment

**Super Pop Color Adjustment** · *A fast residual color adjustment model for vibrant palettes, deep blacks, crisp highlights, HDR bite – in one shot.* diff --git a/docs/SuperBeasts.AI-SuperPopColorAdjustment.json b/docs/SuperBeasts.AI-SuperPopColorAdjustment.json new file mode 100644 index 0000000..44ea8f2 --- /dev/null +++ b/docs/SuperBeasts.AI-SuperPopColorAdjustment.json @@ -0,0 +1,792 @@ +{ + "id": "f616e3a0-cf66-4f06-9787-fd37ac226747", + "revision": 0, + "last_node_id": 133, + "last_link_id": 306, + "nodes": [ + { + "id": 91, + "type": "HDR Effects (SuperBeasts.AI)", + "pos": [ + -660, + -20 + ], + "size": [ + 390, + 180 + ], + "flags": {}, + "order": 5, + "mode": 4, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 275 + } + ], + "outputs": [ + { + "name": "result_img", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 238 + ] + } + ], + "properties": { + "Node name for S&R": "HDR Effects (SuperBeasts.AI)", + "cnr_id": "comfyui-superbeasts", + "ver": "f684c86d0031b82123dc1f7931a8de43c9946fee", + "widget_ue_connectable": {} + }, + "widgets_values": [ + 0.76, + 0.25, + 0.9, + 0, + 0.1, + 0.6400000000000001 + ], + "color": "#14145c", + "bgcolor": "#000048" + }, + { + "id": 113, + "type": "Note", + "pos": [ + -660, + -800 + ], + "size": [ + 390, + 150 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Super Pop Color Adjustment Recommendations", + "properties": { + "widget_ue_connectable": {} + }, + "widgets_values": [ + "- Image(s) 512 - 2048px\n- Context Image - Empty uses it's self for correction context\n- Max strength 1.0\n- Count 1\n- Overlap at least 0.3 (higher is better up to ~0.9 if GPU allows for final images)\n\nExamples and instructions: https://github.com/SuperBeastsAI/ComfyUI-SuperBeasts/blob/master/docs/super-pop-color-adjustment.md" + ], + "color": "#d67023", + "bgcolor": "#c25c0f" + }, + { + "id": 121, + "type": "Super Pop Residual Blend (SuperBeasts.AI)", + "pos": [ + 920, + 160 + ], + "size": [ + 374.6089782714844, + 78 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 282 + }, + { + "name": "residual", + "type": "SPCA_RESIDUAL", + "link": 288 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 297 + ] + }, + { + "name": "filename_prefix", + "type": "STRING", + "links": null + } + ], + "properties": { + "Node name for S&R": "Super Pop Residual Blend (SuperBeasts.AI)", + "cnr_id": "comfyui-superbeasts", + "ver": "117a4529ab5f0ec0145da3e50d61ef73046fa821", + "widget_ue_connectable": {} + }, + "widgets_values": [ + 0.5000000000000001 + ], + "color": "#14145c", + "bgcolor": "#000048" + }, + { + "id": 126, + "type": "Image Comparer (rgthree)", + "pos": [ + 1630, + -800 + ], + "size": [ + 670, + 900 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "dir": 3, + "name": "image_a", + "type": "IMAGE", + "link": 296 + }, + { + "dir": 3, + "name": "image_b", + "type": "IMAGE", + "link": 299 + } + ], + "outputs": [], + "title": "Comparison - Original V.S. Super Pop Colour Adjustment (SuberBeasts.AI)", + "properties": { + "cnr_id": "rgthree-comfy", + "ver": "944d5353a1b0a668f40844018c3dc956b95a67d7", + "comparer_mode": "Slide", + "widget_ue_connectable": {} + }, + "widgets_values": [ + [ + { + "name": "A", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_zisus_00025_.png&type=temp&subfolder=&preview=webp&rand=0.12061521654902096" + }, + { + "name": "B", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_zisus_00026_.png&type=temp&subfolder=&preview=webp&rand=0.3100928355954683" + } + ] + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 125, + "type": "Image Comparer (rgthree)", + "pos": [ + 920, + -800 + ], + "size": [ + 670, + 900 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "dir": 3, + "name": "image_a", + "type": "IMAGE", + "link": 297 + }, + { + "dir": 3, + "name": "image_b", + "type": "IMAGE", + "link": 298 + } + ], + "outputs": [], + "title": "Comparison - Original V.S. Super Pop Colour Adjustment (SuberBeasts.AI)", + "properties": { + "cnr_id": "rgthree-comfy", + "ver": "944d5353a1b0a668f40844018c3dc956b95a67d7", + "comparer_mode": "Slide", + "widget_ue_connectable": {} + }, + "widgets_values": [ + [ + { + "name": "A", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_difuo_00025_.png&type=temp&subfolder=&preview=webp&rand=0.2381890175765341" + }, + { + "name": "B", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_difuo_00026_.png&type=temp&subfolder=&preview=webp&rand=0.1267883942778919" + } + ] + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 114, + "type": "SaveImage", + "pos": [ + -130, + -800 + ], + "size": [ + 210, + 282 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 284 + }, + { + "name": "filename_prefix", + "type": "STRING", + "widget": { + "name": "filename_prefix" + }, + "link": 285 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.44", + "widget_ue_connectable": {} + }, + "widgets_values": [ + "SuperBeasts.AI_SPCA_" + ] + }, + { + "id": 69, + "type": "Image Comparer (rgthree)", + "pos": [ + 110, + -800 + ], + "size": [ + 670, + 900 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "dir": 3, + "name": "image_a", + "type": "IMAGE", + "link": 238 + }, + { + "dir": 3, + "name": "image_b", + "type": "IMAGE", + "link": 189 + } + ], + "outputs": [], + "title": "Comparison - Original V.S. Super Pop Colour Adjustment (SuberBeasts.AI)", + "properties": { + "cnr_id": "rgthree-comfy", + "ver": "944d5353a1b0a668f40844018c3dc956b95a67d7", + "comparer_mode": "Slide", + "widget_ue_connectable": {} + }, + "widgets_values": [ + [ + { + "name": "A", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_ycmde_00025_.png&type=temp&subfolder=&preview=webp&rand=0.7170845396129458" + }, + { + "name": "B", + "selected": true, + "url": "/api/view?filename=rgthree.compare._temp_ycmde_00026_.png&type=temp&subfolder=&preview=webp&rand=0.42308399021076515" + } + ] + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 102, + "type": "SB Load Model (SuperBeasts.AI)", + "pos": [ + -660, + -420 + ], + "size": [ + 390, + 82 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "model", + "type": "SBMODEL", + "links": [ + 273 + ] + } + ], + "properties": { + "Node name for S&R": "SB Load Model (SuperBeasts.AI)", + "cnr_id": "comfyui-superbeasts", + "ver": "73a6a22138b32cd76318a48c594050b91906901d", + "widget_ue_connectable": {} + }, + "widgets_values": [ + "SuperBeasts_ColorAdjustment_512px_V1.safetensors", + "GPU" + ], + "color": "#14145c", + "bgcolor": "#000048" + }, + { + "id": 93, + "type": "LoadImage", + "pos": [ + -1410, + -800 + ], + "size": [ + 630, + 1000 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 306 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.44", + "Node name for S&R": "LoadImage", + "widget_ue_connectable": {} + }, + "widgets_values": [ + "39a.png", + "image" + ] + }, + { + "id": 81, + "type": "JWImageResizeByLongerSide", + "pos": [ + -660, + -580 + ], + "size": [ + 390, + 82 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 306 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 189, + 274, + 282, + 283, + 298, + 299 + ] + } + ], + "properties": { + "cnr_id": "comfyui-various", + "ver": "5bd85aaf7616878471469c4ec7e11bbd0cef3bf2", + "Node name for S&R": "JWImageResizeByLongerSide", + "widget_ue_connectable": {} + }, + "widgets_values": [ + 1024, + "bicubic" + ] + }, + { + "id": 120, + "type": "Super Pop Color Adjustment (SuperBeasts.AI)", + "pos": [ + -660, + -270 + ], + "size": [ + 387.8160095214844, + 170 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "SBMODEL", + "link": 273 + }, + { + "name": "image", + "type": "IMAGE", + "link": 274 + }, + { + "name": "context", + "shape": 7, + "type": "IMAGE", + "link": null + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 275, + 284 + ] + }, + { + "name": "filename_prefix", + "type": "STRING", + "links": [ + 285 + ] + }, + { + "name": "residuals", + "type": "SPCA_RESIDUAL", + "links": [ + 288, + 289 + ] + } + ], + "properties": { + "Node name for S&R": "Super Pop Color Adjustment (SuperBeasts.AI)", + "cnr_id": "comfyui-superbeasts", + "ver": "117a4529ab5f0ec0145da3e50d61ef73046fa821", + "widget_ue_connectable": {} + }, + "widgets_values": [ + 1, + 1, + 0.7500000000000001, + false + ], + "color": "#14145c", + "bgcolor": "#000048" + }, + { + "id": 118, + "type": "Super Pop Residual Blend (SuperBeasts.AI)", + "pos": [ + 1630, + 160 + ], + "size": [ + 374.6089782714844, + 78 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 283 + }, + { + "name": "residual", + "type": "SPCA_RESIDUAL", + "link": 289 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 296 + ] + }, + { + "name": "filename_prefix", + "type": "STRING", + "links": [] + } + ], + "properties": { + "Node name for S&R": "Super Pop Residual Blend (SuperBeasts.AI)", + "cnr_id": "comfyui-superbeasts", + "ver": "117a4529ab5f0ec0145da3e50d61ef73046fa821", + "widget_ue_connectable": {} + }, + "widgets_values": [ + 1.5000000000000002 + ], + "color": "#14145c", + "bgcolor": "#000048" + } + ], + "links": [ + [ + 189, + 81, + 0, + 69, + 1, + "IMAGE" + ], + [ + 238, + 91, + 0, + 69, + 0, + "IMAGE" + ], + [ + 273, + 102, + 0, + 120, + 0, + "SBMODEL" + ], + [ + 274, + 81, + 0, + 120, + 1, + "IMAGE" + ], + [ + 275, + 120, + 0, + 91, + 0, + "IMAGE" + ], + [ + 282, + 81, + 0, + 121, + 0, + "IMAGE" + ], + [ + 283, + 81, + 0, + 118, + 0, + "IMAGE" + ], + [ + 284, + 120, + 0, + 114, + 0, + "IMAGE" + ], + [ + 285, + 120, + 1, + 114, + 1, + "STRING" + ], + [ + 288, + 120, + 2, + 121, + 1, + "SPCA_RESIDUAL" + ], + [ + 289, + 120, + 2, + 118, + 1, + "SPCA_RESIDUAL" + ], + [ + 296, + 118, + 0, + 126, + 0, + "IMAGE" + ], + [ + 297, + 121, + 0, + 125, + 0, + "IMAGE" + ], + [ + 298, + 81, + 0, + 125, + 1, + "IMAGE" + ], + [ + 299, + 81, + 0, + 126, + 1, + "IMAGE" + ], + [ + 306, + 93, + 0, + 81, + 0, + "IMAGE" + ] + ], + "groups": [ + { + "id": 1, + "title": "SuperBeasts.AI Adjustments", + "bounding": [ + -700, + -900, + 480, + 1110 + ], + "color": "#c25c0f", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "Image results", + "bounding": [ + -150, + -900, + 990, + 1040 + ], + "color": "#a1309b", + "font_size": 24, + "flags": {} + }, + { + "id": 4, + "title": "Image(s) Input", + "bounding": [ + -1440, + -900, + 690, + 1120 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 6, + "title": "Alternative SPCA Strength Tests", + "bounding": [ + 880, + -890, + 1460, + 1170 + ], + "color": "#a1309b", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 0.7513148009015783, + "offset": [ + 1300.3125958595183, + 919.6808269326773 + ] + }, + "frontendVersion": "1.23.4", + "groupNodes": {}, + "ue_links": [], + "links_added_by_ue": [], + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true, + "node_versions": { + "comfy-core": "0.3.14", + "ComfyUI-KJNodes": "80977db1ea8472a7934de50f3c878aed405e8d73", + "ComfyUI-SuperBeasts": "73a6a22138b32cd76318a48c594050b91906901d", + "ComfyUI-VideoHelperSuite": "8629188458dc6cb832f871ece3bd273507e8a766", + "comfyui-various": "36454f91606bbff4fc36d90234981ca4a47e2695", + "rgthree-comfy": "5d771b8b56a343c24a26e8cea1f0c87c3d58102f" + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/docs/assets/examples/SPCA/SuperBeasts.AI-SuperPopColorAdjustment.json b/docs/Workflows/SPCA/SuperBeasts.AI-SuperPopColorAdjustment.json similarity index 100% rename from docs/assets/examples/SPCA/SuperBeasts.AI-SuperPopColorAdjustment.json rename to docs/Workflows/SPCA/SuperBeasts.AI-SuperPopColorAdjustment.json diff --git a/docs/assets/examples/SPCA/workflow.jpg b/docs/Workflows/SPCA/workflow.jpg similarity index 100% rename from docs/assets/examples/SPCA/workflow.jpg rename to docs/Workflows/SPCA/workflow.jpg diff --git a/docs/assets/examples/BatchManagers.jpg b/docs/assets/examples/BatchManagers.jpg deleted file mode 100644 index f373996..0000000 Binary files a/docs/assets/examples/BatchManagers.jpg and /dev/null differ diff --git a/docs/assets/examples/SPCA-Header.webp b/docs/assets/examples/SPCA-Header.webp deleted file mode 100644 index b8d54b0..0000000 Binary files a/docs/assets/examples/SPCA-Header.webp and /dev/null differ diff --git a/docs/assets/examples/SPCA/10a.png b/docs/assets/examples/SPCA/10a.png deleted file mode 100644 index 3a0d45e..0000000 Binary files a/docs/assets/examples/SPCA/10a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/10b.png b/docs/assets/examples/SPCA/10b.png deleted file mode 100644 index 226e446..0000000 Binary files a/docs/assets/examples/SPCA/10b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/11a.png b/docs/assets/examples/SPCA/11a.png deleted file mode 100644 index ccfb3ce..0000000 Binary files a/docs/assets/examples/SPCA/11a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/11b.png b/docs/assets/examples/SPCA/11b.png deleted file mode 100644 index 7aab939..0000000 Binary files a/docs/assets/examples/SPCA/11b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/12a.png b/docs/assets/examples/SPCA/12a.png deleted file mode 100644 index cf1f7de..0000000 Binary files a/docs/assets/examples/SPCA/12a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/12b.png b/docs/assets/examples/SPCA/12b.png deleted file mode 100644 index 0ebcbb7..0000000 Binary files a/docs/assets/examples/SPCA/12b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/13a.png b/docs/assets/examples/SPCA/13a.png deleted file mode 100644 index 3b1d9d1..0000000 Binary files a/docs/assets/examples/SPCA/13a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/13b.png b/docs/assets/examples/SPCA/13b.png deleted file mode 100644 index e9fab6c..0000000 Binary files a/docs/assets/examples/SPCA/13b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/14a.png b/docs/assets/examples/SPCA/14a.png deleted file mode 100644 index 94f10a1..0000000 Binary files a/docs/assets/examples/SPCA/14a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/14b.png b/docs/assets/examples/SPCA/14b.png deleted file mode 100644 index 5f7543a..0000000 Binary files a/docs/assets/examples/SPCA/14b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/15a.png b/docs/assets/examples/SPCA/15a.png deleted file mode 100644 index f276b5c..0000000 Binary files a/docs/assets/examples/SPCA/15a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/15b.png b/docs/assets/examples/SPCA/15b.png deleted file mode 100644 index 5e159a3..0000000 Binary files a/docs/assets/examples/SPCA/15b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/16a.png b/docs/assets/examples/SPCA/16a.png deleted file mode 100644 index 2f76720..0000000 Binary files a/docs/assets/examples/SPCA/16a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/16b.png b/docs/assets/examples/SPCA/16b.png deleted file mode 100644 index 42664fc..0000000 Binary files a/docs/assets/examples/SPCA/16b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/17a.png b/docs/assets/examples/SPCA/17a.png deleted file mode 100644 index 5193447..0000000 Binary files a/docs/assets/examples/SPCA/17a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/17b.png b/docs/assets/examples/SPCA/17b.png deleted file mode 100644 index de892da..0000000 Binary files a/docs/assets/examples/SPCA/17b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/18a.png b/docs/assets/examples/SPCA/18a.png deleted file mode 100644 index b5a818a..0000000 Binary files a/docs/assets/examples/SPCA/18a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/18b.png b/docs/assets/examples/SPCA/18b.png deleted file mode 100644 index d3e62e2..0000000 Binary files a/docs/assets/examples/SPCA/18b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/19a.png b/docs/assets/examples/SPCA/19a.png deleted file mode 100644 index cce22d7..0000000 Binary files a/docs/assets/examples/SPCA/19a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/19b.png b/docs/assets/examples/SPCA/19b.png deleted file mode 100644 index adf394c..0000000 Binary files a/docs/assets/examples/SPCA/19b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/1a.png b/docs/assets/examples/SPCA/1a.png deleted file mode 100644 index 1f7ae4e..0000000 Binary files a/docs/assets/examples/SPCA/1a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/1b.png b/docs/assets/examples/SPCA/1b.png deleted file mode 100644 index 7b0eb83..0000000 Binary files a/docs/assets/examples/SPCA/1b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/20a.png b/docs/assets/examples/SPCA/20a.png deleted file mode 100644 index f4d1a37..0000000 Binary files a/docs/assets/examples/SPCA/20a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/20b.png b/docs/assets/examples/SPCA/20b.png deleted file mode 100644 index 3081f07..0000000 Binary files a/docs/assets/examples/SPCA/20b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/21a.png b/docs/assets/examples/SPCA/21a.png deleted file mode 100644 index 101ab69..0000000 Binary files a/docs/assets/examples/SPCA/21a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/21b.png b/docs/assets/examples/SPCA/21b.png deleted file mode 100644 index 0d65b00..0000000 Binary files a/docs/assets/examples/SPCA/21b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/22a.png b/docs/assets/examples/SPCA/22a.png deleted file mode 100644 index 3852af4..0000000 Binary files a/docs/assets/examples/SPCA/22a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/22b.png b/docs/assets/examples/SPCA/22b.png deleted file mode 100644 index 907342d..0000000 Binary files a/docs/assets/examples/SPCA/22b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/23a.png b/docs/assets/examples/SPCA/23a.png deleted file mode 100644 index 5a73068..0000000 Binary files a/docs/assets/examples/SPCA/23a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/23b.png b/docs/assets/examples/SPCA/23b.png deleted file mode 100644 index 09af354..0000000 Binary files a/docs/assets/examples/SPCA/23b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/24a.png b/docs/assets/examples/SPCA/24a.png deleted file mode 100644 index d0c865c..0000000 Binary files a/docs/assets/examples/SPCA/24a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/24b.png b/docs/assets/examples/SPCA/24b.png deleted file mode 100644 index 4c3d346..0000000 Binary files a/docs/assets/examples/SPCA/24b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/25a.png b/docs/assets/examples/SPCA/25a.png deleted file mode 100644 index aedd6a4..0000000 Binary files a/docs/assets/examples/SPCA/25a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/25b.png b/docs/assets/examples/SPCA/25b.png deleted file mode 100644 index 842bef2..0000000 Binary files a/docs/assets/examples/SPCA/25b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/26a.png b/docs/assets/examples/SPCA/26a.png deleted file mode 100644 index 0c7ea69..0000000 Binary files a/docs/assets/examples/SPCA/26a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/26b.png b/docs/assets/examples/SPCA/26b.png deleted file mode 100644 index cdb77ae..0000000 Binary files a/docs/assets/examples/SPCA/26b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/27a.png b/docs/assets/examples/SPCA/27a.png deleted file mode 100644 index 9be1375..0000000 Binary files a/docs/assets/examples/SPCA/27a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/27b.png b/docs/assets/examples/SPCA/27b.png deleted file mode 100644 index 76019c3..0000000 Binary files a/docs/assets/examples/SPCA/27b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/28a.png b/docs/assets/examples/SPCA/28a.png deleted file mode 100644 index 920d727..0000000 Binary files a/docs/assets/examples/SPCA/28a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/28b.png b/docs/assets/examples/SPCA/28b.png deleted file mode 100644 index 03d36e6..0000000 Binary files a/docs/assets/examples/SPCA/28b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/29a.png b/docs/assets/examples/SPCA/29a.png deleted file mode 100644 index 2719c70..0000000 Binary files a/docs/assets/examples/SPCA/29a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/29b.png b/docs/assets/examples/SPCA/29b.png deleted file mode 100644 index 43c083d..0000000 Binary files a/docs/assets/examples/SPCA/29b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/2a.png b/docs/assets/examples/SPCA/2a.png deleted file mode 100644 index 2b63789..0000000 Binary files a/docs/assets/examples/SPCA/2a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/2b.png b/docs/assets/examples/SPCA/2b.png deleted file mode 100644 index e1185e1..0000000 Binary files a/docs/assets/examples/SPCA/2b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/30a.png b/docs/assets/examples/SPCA/30a.png deleted file mode 100644 index 0881663..0000000 Binary files a/docs/assets/examples/SPCA/30a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/30b.png b/docs/assets/examples/SPCA/30b.png deleted file mode 100644 index 11fed4c..0000000 Binary files a/docs/assets/examples/SPCA/30b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/31a.png b/docs/assets/examples/SPCA/31a.png deleted file mode 100644 index 4236b04..0000000 Binary files a/docs/assets/examples/SPCA/31a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/31b.png b/docs/assets/examples/SPCA/31b.png deleted file mode 100644 index efb04f0..0000000 Binary files a/docs/assets/examples/SPCA/31b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/32a.png b/docs/assets/examples/SPCA/32a.png deleted file mode 100644 index 11214cc..0000000 Binary files a/docs/assets/examples/SPCA/32a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/32b.png b/docs/assets/examples/SPCA/32b.png deleted file mode 100644 index f921541..0000000 Binary files a/docs/assets/examples/SPCA/32b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/33a.png b/docs/assets/examples/SPCA/33a.png deleted file mode 100644 index 23a5b36..0000000 Binary files a/docs/assets/examples/SPCA/33a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/33b.png b/docs/assets/examples/SPCA/33b.png deleted file mode 100644 index cf53582..0000000 Binary files a/docs/assets/examples/SPCA/33b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/34a.png b/docs/assets/examples/SPCA/34a.png deleted file mode 100644 index 98085fd..0000000 Binary files a/docs/assets/examples/SPCA/34a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/34b.png b/docs/assets/examples/SPCA/34b.png deleted file mode 100644 index 717dc6e..0000000 Binary files a/docs/assets/examples/SPCA/34b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/35a.png b/docs/assets/examples/SPCA/35a.png deleted file mode 100644 index 016fd1a..0000000 Binary files a/docs/assets/examples/SPCA/35a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/35b.png b/docs/assets/examples/SPCA/35b.png deleted file mode 100644 index e9c9339..0000000 Binary files a/docs/assets/examples/SPCA/35b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/36a.png b/docs/assets/examples/SPCA/36a.png deleted file mode 100644 index c1c7dfb..0000000 Binary files a/docs/assets/examples/SPCA/36a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/36b.png b/docs/assets/examples/SPCA/36b.png deleted file mode 100644 index d968d90..0000000 Binary files a/docs/assets/examples/SPCA/36b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/37a.png b/docs/assets/examples/SPCA/37a.png deleted file mode 100644 index 97e11ad..0000000 Binary files a/docs/assets/examples/SPCA/37a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/37b.png b/docs/assets/examples/SPCA/37b.png deleted file mode 100644 index 0007d99..0000000 Binary files a/docs/assets/examples/SPCA/37b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/38a.png b/docs/assets/examples/SPCA/38a.png deleted file mode 100644 index 21e833f..0000000 Binary files a/docs/assets/examples/SPCA/38a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/38b.png b/docs/assets/examples/SPCA/38b.png deleted file mode 100644 index 8a6f11c..0000000 Binary files a/docs/assets/examples/SPCA/38b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/39a.png b/docs/assets/examples/SPCA/39a.png deleted file mode 100644 index 010c6af..0000000 Binary files a/docs/assets/examples/SPCA/39a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/39b.png b/docs/assets/examples/SPCA/39b.png deleted file mode 100644 index e50fe16..0000000 Binary files a/docs/assets/examples/SPCA/39b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/3a.png b/docs/assets/examples/SPCA/3a.png deleted file mode 100644 index aa9745f..0000000 Binary files a/docs/assets/examples/SPCA/3a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/3b.png b/docs/assets/examples/SPCA/3b.png deleted file mode 100644 index a9f90b6..0000000 Binary files a/docs/assets/examples/SPCA/3b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/40a.png b/docs/assets/examples/SPCA/40a.png deleted file mode 100644 index 50c2d65..0000000 Binary files a/docs/assets/examples/SPCA/40a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/40b.png b/docs/assets/examples/SPCA/40b.png deleted file mode 100644 index 5b9614e..0000000 Binary files a/docs/assets/examples/SPCA/40b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/41a.png b/docs/assets/examples/SPCA/41a.png deleted file mode 100644 index f6fcc36..0000000 Binary files a/docs/assets/examples/SPCA/41a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/41b.png b/docs/assets/examples/SPCA/41b.png deleted file mode 100644 index d6e1959..0000000 Binary files a/docs/assets/examples/SPCA/41b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/42a.png b/docs/assets/examples/SPCA/42a.png deleted file mode 100644 index ed98ca6..0000000 Binary files a/docs/assets/examples/SPCA/42a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/4a.png b/docs/assets/examples/SPCA/4a.png deleted file mode 100644 index 88e97e4..0000000 Binary files a/docs/assets/examples/SPCA/4a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/4b.png b/docs/assets/examples/SPCA/4b.png deleted file mode 100644 index 1319c70..0000000 Binary files a/docs/assets/examples/SPCA/4b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/5a.png b/docs/assets/examples/SPCA/5a.png deleted file mode 100644 index dafe973..0000000 Binary files a/docs/assets/examples/SPCA/5a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/5b.png b/docs/assets/examples/SPCA/5b.png deleted file mode 100644 index 8a1a4cc..0000000 Binary files a/docs/assets/examples/SPCA/5b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/6a.png b/docs/assets/examples/SPCA/6a.png deleted file mode 100644 index 4e6b2a8..0000000 Binary files a/docs/assets/examples/SPCA/6a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/6b.png b/docs/assets/examples/SPCA/6b.png deleted file mode 100644 index 1ad355d..0000000 Binary files a/docs/assets/examples/SPCA/6b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/7a.png b/docs/assets/examples/SPCA/7a.png deleted file mode 100644 index adb30cc..0000000 Binary files a/docs/assets/examples/SPCA/7a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/7b.png b/docs/assets/examples/SPCA/7b.png deleted file mode 100644 index 85641a5..0000000 Binary files a/docs/assets/examples/SPCA/7b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/8a.png b/docs/assets/examples/SPCA/8a.png deleted file mode 100644 index 28c2c85..0000000 Binary files a/docs/assets/examples/SPCA/8a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/8b.png b/docs/assets/examples/SPCA/8b.png deleted file mode 100644 index bc6e5d6..0000000 Binary files a/docs/assets/examples/SPCA/8b.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/9a.png b/docs/assets/examples/SPCA/9a.png deleted file mode 100644 index ffc1075..0000000 Binary files a/docs/assets/examples/SPCA/9a.png and /dev/null differ diff --git a/docs/assets/examples/SPCA/9b.png b/docs/assets/examples/SPCA/9b.png deleted file mode 100644 index 50b38c9..0000000 Binary files a/docs/assets/examples/SPCA/9b.png and /dev/null differ diff --git a/docs/assets/examples/ex1.png b/docs/assets/examples/ex1.png deleted file mode 100644 index 5050b2b..0000000 Binary files a/docs/assets/examples/ex1.png and /dev/null differ diff --git a/docs/assets/examples/ex2.png b/docs/assets/examples/ex2.png deleted file mode 100644 index 229b06c..0000000 Binary files a/docs/assets/examples/ex2.png and /dev/null differ diff --git a/docs/assets/icon.png b/docs/assets/icon.png deleted file mode 100644 index a8cba16..0000000 Binary files a/docs/assets/icon.png and /dev/null differ diff --git a/docs/batch-management.md b/docs/batch-management.md index 459b54a..2ebbc23 100644 --- a/docs/batch-management.md +++ b/docs/batch-management.md @@ -3,7 +3,7 @@ Keep your image and masks sized, cropped and ordered how ever you like without having to recreate the masks or mess with connections. Note the resizing capability doesn't provide any settings, we simply crop and resize to maximise the image size with your provided width/height. -![Batch Management](../examples/BatchManagers.jpg) +![Batch Management](https://s3.ap-southeast-2.amazonaws.com/superbeasts.ai/repos/ComfyUI-SuperBeasts/examples/BatchManagers.jpg) ## ImageBatchManagement diff --git a/docs/hdr-effects.md b/docs/hdr-effects.md index 6678e43..7384267 100644 --- a/docs/hdr-effects.md +++ b/docs/hdr-effects.md @@ -13,10 +13,10 @@ The HDR Effects is an image processing application that enhances the dynamic ran ## Examples -![Example 1](assets/examples/ex1.png) +![Example 1](https://s3.ap-southeast-2.amazonaws.com/superbeasts.ai/repos/ComfyUI-SuperBeasts/examples/ex1.png) -![Example 2](assets/examples/ex2.png) +![Example 2](https://s3.ap-southeast-2.amazonaws.com/superbeasts.ai/repos/ComfyUI-SuperBeasts/examples/ex2.png) ### Parameter details diff --git a/docs/super-pop-color-adjustment.md b/docs/super-pop-color-adjustment.md index 198fcf9..27e4712 100644 --- a/docs/super-pop-color-adjustment.md +++ b/docs/super-pop-color-adjustment.md @@ -1,9 +1,11 @@ + + # Super Pop Color Adjustment (SuperBeasts.AI) Bring the **SuperBeasts “pop”** to any image in one click using a new custom residual color correction model trained to push **vibrant palette, deep blacks, crisp highlights, and HDR bite** with ease. - -Downlaod the ComfyUI Workflow. + +Download the ComfyUI Workflow. ### Why I built it Everything I post gets graded: lift blacks, pull whites, bend curves, punch colour. Existing auto tools (and many AI outputs — looking at you, muddy yellow “whites”) just weren’t landing, especially across batches or iterative generations. Training my own lightweight correction model let me reclaim hours of post and lock in a consistent SuperBeasts look I could reuse — and share. @@ -19,6 +21,8 @@ Happy generating!! ### Using Super Pop +Before starting please review the the licence and [System Requirements](#System-Requirements) to run onnx models. + 1. Add **SB Load Model (SuperBeasts.AI)** node. 2. Pick **SuperPopColorAdjustment/latest**. 3a. Connect to **Super Pop Color Adjustment (SuperBeasts.AI)** node. @@ -45,21 +49,21 @@ Local/client use OK; no public SaaS redistribution. - + - + - - - - + + + + - - - - + + + +
OriginalS.P.C.ASPCA OriginalS.P.C.ASPCA
Original 1S.P.C.A 1Original 2S.P.C.A 2Original 1SPCA 1Original 2SPCA 2
Original 3S.P.C.A 3Original 4S.P.C.A 4Original 3SPCA 3Original 4SPCA 4
@@ -68,21 +72,21 @@ Local/client use OK; no public SaaS redistribution. - + - + - - - - + + + + - - - - + + + +
OriginalS.P.C.ASPCA OriginalS.P.C.ASPCA
Original 1S.P.C.A 1Original 2S.P.C.A 2Original 5SPCA 5Original 6SPCA 6
Original 3S.P.C.A 3Original 4S.P.C.A 4Original 7SPCA 7Original 8SPCA 8
@@ -92,21 +96,21 @@ Local/client use OK; no public SaaS redistribution. - + - + - - - - + + + + - - - - + + + +
OriginalS.P.C.ASPCA OriginalS.P.C.ASPCA
Original 1S.P.C.A 1Original 2S.P.C.A 2Original 9SPCA 9Original 10SPCA 10
Original 3S.P.C.A 3Original 4S.P.C.A 4Original 11SPCA 11Original 12SPCA 12
@@ -116,21 +120,21 @@ Local/client use OK; no public SaaS redistribution. - + - + - - - - + + + + - - - - + + + +
OriginalS.P.C.ASPCA OriginalS.P.C.ASPCA
Original 1S.P.C.A 1Original 2S.P.C.A 2Original 13SPCA 13Original 14SPCA 14
Original 3S.P.C.A 3Original 4S.P.C.A 4Original 15SPCA 15Original 16SPCA 16
@@ -140,21 +144,21 @@ Local/client use OK; no public SaaS redistribution. - + - + - - - - + + + + - - - - + + + +
OriginalS.P.C.ASPCA OriginalS.P.C.ASPCA
Original 1S.P.C.A 1Original 2S.P.C.A 2Original 17SPCA 17Original 18SPCA 18
Original 3S.P.C.A 3Original 4S.P.C.A 4Original 19SPCA 19Original 20SPCA 20
@@ -164,21 +168,21 @@ Local/client use OK; no public SaaS redistribution. - + - + - - - - + + + + - - - - + + + +
OriginalS.P.C.ASPCA OriginalS.P.C.ASPCA
Original 1S.P.C.A 1Original 2S.P.C.A 2Original 21SPCA 21Original 22SPCA 22
Original 3S.P.C.A 3Original 4S.P.C.A 4Original 23SPCA 23Original 24SPCA 24
@@ -244,5 +248,23 @@ down-scaled or abstract images work great but this only plays a small role in th ### System Requirements -Works out-of-the-box on CPU via `onnxruntime`. For **10×** speed improvement ensure ComfyUI can use the `CUDAExecutionProvider` by using onnxruntime-gpu with the hardware that allows this. -This node print out the currently available providers into the ComfyUI server logs if you need to check this. \ No newline at end of file +Works out-of-the-box on CPU via `onnxruntime` but it is extremely slow and I wouldn't recommend it. +For **10×** speed improvement ensure ComfyUI can use the `CUDAExecutionProvider` by installing `onnxruntime-gpu` into the ComfyUI environment. + +This node currently prints out the available providers for onnxruntime into the ComfyUI server logs when you run the node. So if you are not seeing `CUDAExecutionProvider` it's likely the model is using CPU instead. + +For reference I use: +Python version: 3.12.1 +pytorch version: 2.5.0+cu124 +NVIDIA GeForce RTX 4090 +with onnxruntime-gpu installed in the ComfyUI Environment +and C:\Program Files\NVIDIA\CUDNN\vX.X\bin on your System Path (where x.x is your version) +For more on debugging try tips from the user here: https://github.com/SuperBeastsAI/ComfyUI-SuperBeasts/issues/10#issuecomment-3113530795 + +Results - Server Logs show: +[SuperBeasts] ONNX Runtime providers available: ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'] +[SuperBeasts] Using ORT providers: ['CUDAExecutionProvider', 'CPUExecutionProvider'] +1024x1024px with 0.5 overlap = 1.63s +2048x2048px with 0.5 overlap = 5.23s + +Note: I am aware the extra installation of `onnxruntime-gpu` isn't ideal and I'm attempting to convert the model into .safetensors to work within the ComfyUI Pytorch environment better with less issues accessing the GPU. diff --git a/pyproject.toml b/pyproject.toml index 88a1eff..f5001e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ Repository = "https://github.com/SuperBeastsAI/ComfyUI-SuperBeasts" [tool.comfy] PublisherId = "superbeastsai" DisplayName = "ComfyUI-SuperBeasts" -Icon = "https://raw.githubusercontent.com/SuperBeastsAI/ComfyUI-SuperBeasts/refs/heads/master/docs/assets/icon.png" +Icon = "https://s3.ap-southeast-2.amazonaws.com/superbeasts.ai/repos/ComfyUI-SuperBeasts/icon.png" excludes = [ "docs" ] \ No newline at end of file diff --git a/scripts/dump_onnx_tensors.py b/scripts/dump_onnx_tensors.py new file mode 100644 index 0000000..ef1394f --- /dev/null +++ b/scripts/dump_onnx_tensors.py @@ -0,0 +1,19 @@ +from pathlib import Path +import csv, onnx + +def main(): + onnx_path = Path("models/SuperBeasts_ColorAdjustment_512px_V1.onnx") + if not onnx_path.exists(): + raise FileNotFoundError(onnx_path) + model = onnx.load(str(onnx_path)) + out_path = Path(__file__).with_name("tensor_blueprint.csv") + with out_path.open("w", newline="") as f: + wr = csv.writer(f) + wr.writerow(["index", "name", "shape"]) + for idx, init in enumerate(model.graph.initializer): + dims = list(init.dims) + wr.writerow([idx, init.name, dims]) + print(f"[dump] wrote {out_path} with {len(model.graph.initializer)} tensors") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/tensor_blueprint.csv b/scripts/tensor_blueprint.csv new file mode 100644 index 0000000..fcec913 --- /dev/null +++ b/scripts/tensor_blueprint.csv @@ -0,0 +1,178 @@ +index,name,shape +0,new_shape__1590,[5] +1,new_shape__1589,[5] +2,new_shape__1588,[5] +3,new_shape__1587,[5] +4,new_shape__1586,[5] +5,new_shape__1585,[5] +6,new_shape__1013,[4] +7,new_shape__1010,[4] +8,new_shape__1008,[4] +9,einsum124219583030736_ba_batch_axes__1283,[2] +10,einsum124219582908752_ba_left_set__1510,[1] +11,einsum124219582908752_ba_final_axes__1513,[1] +12,einsum124218983837520_id_axes__1281,[1] +13,einsum124218981442064_ba__1__1213,[1] +14,einsum124218844259344_ba_ones__1436,[1] +15,contextual_residual_unet_1/up_sampling2d_1_2/Repeat_1/concat/values_1__9,[1] +16,contextual_residual_unet_1/skip_proj_2_1/convolution/Cast:0,"[32, 64, 1, 1]" +17,contextual_residual_unet_1/skip_proj_2_1/Squeeze:0,[32] +18,contextual_residual_unet_1/skip_proj_1_1/convolution/Cast:0,"[64, 96, 1, 1]" +19,contextual_residual_unet_1/skip_proj_1_1/Squeeze:0,[64] +20,contextual_residual_unet_1/skip_proj_0_1/convolution/Cast:0,"[96, 128, 1, 1]" +21,contextual_residual_unet_1/skip_proj_0_1/Squeeze:0,[96] +22,contextual_residual_unet_1/reshape_8_1/Reshape_shape__1125,[4] +23,contextual_residual_unet_1/reshape_7_1/Reshape_shape__1129,[3] +24,contextual_residual_unet_1/reshape_5_1/Reshape_shape__1130,[4] +25,contextual_residual_unet_1/reshape_4_1/Reshape_shape__1127,[3] +26,contextual_residual_unet_1/reshape_2_2/Reshape_shape__1128,[4] +27,contextual_residual_unet_1/reshape_1_1/Reshape_shape__1126,[3] +28,contextual_residual_unet_1/group_normalization_9_1/Cast_2:0,"[1, 1, 1, 8, 12]" +29,contextual_residual_unet_1/group_normalization_9_1/Cast_1:0,"[1, 1, 1, 8, 12]" +30,contextual_residual_unet_1/group_normalization_8_1/Cast_2:0,"[1, 1, 1, 8, 16]" +31,contextual_residual_unet_1/group_normalization_8_1/Cast_1:0,"[1, 1, 1, 8, 16]" +32,contextual_residual_unet_1/group_normalization_7_1/Cast_2:0,"[1, 1, 1, 8, 16]" +33,contextual_residual_unet_1/group_normalization_7_1/Cast_1:0,"[1, 1, 1, 8, 16]" +34,contextual_residual_unet_1/group_normalization_6_1/Cast_2:0,"[1, 1, 1, 8, 16]" +35,contextual_residual_unet_1/group_normalization_6_1/Cast_1:0,"[1, 1, 1, 8, 16]" +36,contextual_residual_unet_1/group_normalization_5_1/Cast_2:0,"[1, 1, 1, 8, 16]" +37,contextual_residual_unet_1/group_normalization_5_1/Cast_1:0,"[1, 1, 1, 8, 16]" +38,contextual_residual_unet_1/group_normalization_4_1/Cast_2:0,"[1, 1, 1, 8, 12]" +39,contextual_residual_unet_1/group_normalization_4_1/Cast_1:0,"[1, 1, 1, 8, 12]" +40,contextual_residual_unet_1/group_normalization_3_1/Cast_2:0,"[1, 1, 1, 8, 12]" +41,contextual_residual_unet_1/group_normalization_3_1/Cast_1:0,"[1, 1, 1, 8, 12]" +42,contextual_residual_unet_1/group_normalization_2_1/add/y:0,[] +43,contextual_residual_unet_1/group_normalization_2_1/Cast_2:0,"[1, 1, 1, 8, 8]" +44,contextual_residual_unet_1/group_normalization_2_1/Cast_1:0,"[1, 1, 1, 8, 8]" +45,contextual_residual_unet_1/group_normalization_1_2/Cast_2:0,"[1, 1, 1, 8, 8]" +46,contextual_residual_unet_1/group_normalization_1_2/Cast_1:0,"[1, 1, 1, 8, 8]" +47,contextual_residual_unet_1/group_normalization_14_1/Cast_2:0,"[1, 1, 1, 8, 4]" +48,contextual_residual_unet_1/group_normalization_14_1/Cast_1:0,"[1, 1, 1, 8, 4]" +49,contextual_residual_unet_1/group_normalization_13_1/Cast_2:0,"[1, 1, 1, 8, 4]" +50,contextual_residual_unet_1/group_normalization_13_1/Cast_1:0,"[1, 1, 1, 8, 4]" +51,contextual_residual_unet_1/group_normalization_12_1/Cast_2:0,"[1, 1, 1, 8, 8]" +52,contextual_residual_unet_1/group_normalization_12_1/Cast_1:0,"[1, 1, 1, 8, 8]" +53,contextual_residual_unet_1/group_normalization_11_1/Cast_2:0,"[1, 1, 1, 8, 8]" +54,contextual_residual_unet_1/group_normalization_11_1/Cast_1:0,"[1, 1, 1, 8, 8]" +55,contextual_residual_unet_1/group_normalization_10_1/Cast_2:0,"[1, 1, 1, 8, 12]" +56,contextual_residual_unet_1/group_normalization_10_1/Cast_1:0,"[1, 1, 1, 8, 12]" +57,contextual_residual_unet_1/group_normalization_1/Cast_2:0,"[1, 1, 1, 8, 4]" +58,contextual_residual_unet_1/group_normalization_1/Cast_1:0,"[1, 1, 1, 8, 4]" +59,contextual_residual_unet_1/decoder_q_proj_2_1/convolution/Cast:0,"[32, 64, 1, 1]" +60,contextual_residual_unet_1/decoder_q_proj_1_1/convolution/Cast:0,"[64, 96, 1, 1]" +61,contextual_residual_unet_1/decoder_q_proj_1_1/Squeeze:0,[64] +62,contextual_residual_unet_1/decoder_q_proj_0_1/convolution/Cast:0,"[96, 128, 1, 1]" +63,contextual_residual_unet_1/decoder_mha_2_1/value_1/Reshape:0,"[32, 128]" +64,contextual_residual_unet_1/decoder_mha_2_1/value_1/Add/Cast:0,"[4, 32]" +65,contextual_residual_unet_1/decoder_mha_2_1/query_1/Reshape:0,"[32, 128]" +66,contextual_residual_unet_1/decoder_mha_2_1/key_1/Reshape:0,"[32, 128]" +67,contextual_residual_unet_1/decoder_mha_2_1/attention_output_1/Reshape_1:0,"[128, 32]" +68,contextual_residual_unet_1/decoder_mha_2_1/attention_output_1/BiasAdd/Cast:0,[32] +69,contextual_residual_unet_1/decoder_mha_2_1/Cast_1/x:0,[] +70,contextual_residual_unet_1/decoder_mha_1_1/value_1/Reshape:0,"[64, 256]" +71,contextual_residual_unet_1/decoder_mha_1_1/value_1/Add/Cast:0,"[4, 64]" +72,contextual_residual_unet_1/decoder_mha_1_1/query_1/Reshape:0,"[64, 256]" +73,contextual_residual_unet_1/decoder_mha_1_1/key_1/Reshape:0,"[64, 256]" +74,contextual_residual_unet_1/decoder_mha_1_1/attention_output_1/Reshape_1:0,"[256, 64]" +75,contextual_residual_unet_1/decoder_mha_1_1/Cast_1/x:0,[] +76,contextual_residual_unet_1/decoder_mha_0_1/value_1/Reshape:0,"[96, 384]" +77,contextual_residual_unet_1/decoder_mha_0_1/value_1/Add/Cast:0,"[4, 96]" +78,contextual_residual_unet_1/decoder_mha_0_1/query_1/Reshape:0,"[96, 384]" +79,contextual_residual_unet_1/decoder_mha_0_1/key_1/Reshape:0,"[96, 384]" +80,contextual_residual_unet_1/decoder_mha_0_1/attention_output_1/Reshape_1:0,"[384, 96]" +81,contextual_residual_unet_1/decoder_mha_0_1/attention_output_1/BiasAdd/Cast:0,[96] +82,contextual_residual_unet_1/decoder_mha_0_1/Cast_1/x:0,[] +83,contextual_residual_unet_1/decoder_ctx_kv_proj_2_1/Cast/Cast:0,"[128, 32]" +84,contextual_residual_unet_1/decoder_ctx_kv_proj_2_1/BiasAdd/Cast:0,[32] +85,contextual_residual_unet_1/decoder_ctx_kv_proj_1_1/Cast/Cast:0,"[128, 64]" +86,contextual_residual_unet_1/decoder_ctx_kv_proj_1_1/BiasAdd/Cast:0,[64] +87,contextual_residual_unet_1/decoder_ctx_kv_proj_0_1/Cast/Cast:0,"[128, 96]" +88,contextual_residual_unet_1/decoder_ctx_kv_proj_0_1/BiasAdd/Cast:0,[96] +89,contextual_residual_unet_1/conv2d_9_1/convolution/Cast:0,"[128, 96, 1, 1]" +90,contextual_residual_unet_1/conv2d_9_1/Squeeze:0,[128] +91,contextual_residual_unet_1/conv2d_8_1/convolution/Cast:0,"[96, 96, 3, 3]" +92,contextual_residual_unet_1/conv2d_8_1/Squeeze:0,[96] +93,contextual_residual_unet_1/conv2d_7_1/convolution/Cast:0,"[96, 64, 3, 3]" +94,contextual_residual_unet_1/conv2d_7_1/Squeeze:0,[96] +95,contextual_residual_unet_1/conv2d_6_1/convolution/Cast:0,"[96, 64, 1, 1]" +96,contextual_residual_unet_1/conv2d_6_1/Squeeze:0,[96] +97,contextual_residual_unet_1/conv2d_5_1/convolution/Cast:0,"[64, 64, 3, 3]" +98,contextual_residual_unet_1/conv2d_5_1/Squeeze:0,[64] +99,contextual_residual_unet_1/conv2d_4_1/convolution/Cast:0,"[64, 32, 3, 3]" +100,contextual_residual_unet_1/conv2d_4_1/Squeeze:0,[64] +101,contextual_residual_unet_1/conv2d_3_1/convolution/Cast:0,"[64, 32, 1, 1]" +102,contextual_residual_unet_1/conv2d_3_1/Squeeze:0,[64] +103,contextual_residual_unet_1/conv2d_2_1/convolution/Cast:0,"[32, 7, 3, 3]" +104,contextual_residual_unet_1/conv2d_2_1/Squeeze:0,[32] +105,contextual_residual_unet_1/conv2d_26_1/convolution/ReadVariableOp:0,"[3, 32, 3, 3]" +106,contextual_residual_unet_1/conv2d_26_1/Squeeze:0,[3] +107,contextual_residual_unet_1/conv2d_25_1/convolution/Cast:0,"[32, 32, 3, 3]" +108,contextual_residual_unet_1/conv2d_25_1/Squeeze:0,[32] +109,contextual_residual_unet_1/conv2d_24_1/convolution/Cast:0,"[32, 64, 3, 3]" +110,contextual_residual_unet_1/conv2d_24_1/Squeeze:0,[32] +111,contextual_residual_unet_1/conv2d_23_1/convolution/Cast:0,"[32, 64, 1, 1]" +112,contextual_residual_unet_1/conv2d_23_1/Squeeze:0,[32] +113,contextual_residual_unet_1/conv2d_22_1/convolution/Cast:0,"[1, 2, 7, 7]" +114,contextual_residual_unet_1/conv2d_22_1/Reshape:0,[1] +115,contextual_residual_unet_1/conv2d_21_1/convolution/Cast:0,"[64, 64, 3, 3]" +116,contextual_residual_unet_1/conv2d_21_1/Squeeze:0,[64] +117,contextual_residual_unet_1/conv2d_20_1/convolution/Cast:0,"[64, 128, 3, 3]" +118,contextual_residual_unet_1/conv2d_20_1/Squeeze:0,[64] +119,contextual_residual_unet_1/conv2d_1_2/convolution/Cast:0,"[64, 32, 3, 3]" +120,contextual_residual_unet_1/conv2d_1_2/Squeeze:0,[64] +121,contextual_residual_unet_1/conv2d_19_1/convolution/Cast:0,"[64, 128, 1, 1]" +122,contextual_residual_unet_1/conv2d_19_1/Squeeze:0,[64] +123,contextual_residual_unet_1/conv2d_18_1/convolution/Cast:0,"[1, 2, 7, 7]" +124,contextual_residual_unet_1/conv2d_18_1/Reshape:0,[1] +125,contextual_residual_unet_1/conv2d_17_1/convolution/Cast:0,"[96, 96, 3, 3]" +126,contextual_residual_unet_1/conv2d_17_1/Squeeze:0,[96] +127,contextual_residual_unet_1/conv2d_16_1/convolution/Cast:0,"[96, 192, 3, 3]" +128,contextual_residual_unet_1/conv2d_16_1/Squeeze:0,[96] +129,contextual_residual_unet_1/conv2d_15_1/convolution/Cast:0,"[96, 192, 1, 1]" +130,contextual_residual_unet_1/conv2d_15_1/Squeeze:0,[96] +131,contextual_residual_unet_1/conv2d_14_1/convolution/Cast:0,"[1, 2, 7, 7]" +132,contextual_residual_unet_1/conv2d_14_1/Reshape:0,[1] +133,contextual_residual_unet_1/conv2d_13_1/convolution/Cast:0,"[128, 128, 3, 3]" +134,contextual_residual_unet_1/conv2d_13_1/Squeeze:0,[128] +135,contextual_residual_unet_1/conv2d_12_1/convolution/Cast:0,"[128, 128, 3, 3]" +136,contextual_residual_unet_1/conv2d_12_1/Squeeze:0,[128] +137,contextual_residual_unet_1/conv2d_11_1/convolution/Cast:0,"[128, 128, 3, 3]" +138,contextual_residual_unet_1/conv2d_11_1/Squeeze:0,[128] +139,contextual_residual_unet_1/conv2d_10_1/convolution/Cast:0,"[128, 96, 3, 3]" +140,contextual_residual_unet_1/conv2d_10_1/Squeeze:0,[128] +141,contextual_residual_unet_1/conv2d_1/convolution/Cast:0,"[32, 7, 3, 3]" +142,contextual_residual_unet_1/conv2d_1/Squeeze:0,[32] +143,contextual_residual_unet_1/conv2d_1/Gelu/add/x:0,[] +144,contextual_residual_unet_1/context_embedding_1/Gelu/mul/x:0,[] +145,contextual_residual_unet_1/context_embedding_1/Cast/Cast:0,"[64, 128]" +146,contextual_residual_unet_1/context_embedding_1/BiasAdd/Cast:0,[128] +147,const_fold_opt__1120,[4] +148,const_fold_opt__1119,[3] +149,const_fold_opt__1115,[5] +150,const_fold_opt__1114,[1] +151,const_fold_opt__1106,[1] +152,const_fold_opt__1104,[1] +153,const_fold_opt__1103,[5] +154,const_fold_opt__1100,[1] +155,const_fold_opt__1099,[3] +156,const_fold_opt__1091,[1] +157,const_fold_opt__1079,[1] +158,const_fold_opt__1077,[1] +159,const_fold_opt__1073,[4] +160,const_fold_opt__1072,[1] +161,const_fold_opt__1070__1578,[5] +162,const_fold_opt__1065,[5] +163,const_fold_opt__1063,[5] +164,const_fold_opt__1062,[5] +165,const_fold_opt__1059,[4] +166,const_fold_opt__1054,[4] +167,const_fold_opt__1049,[3] +168,const_fold_opt__1036,[5] +169,const_fold_opt__1034,[4] +170,const_fold_opt__1021,[5] +171,const_fold_opt__1015,[4] +172,const_ends__372,[1] +173,const_axes__508,[1] +174,const_axes__1131,[2] +175,ConstantFolding/contextual_residual_unet_1/context_embedding_1/Gelu/truediv_recip:0,[] +176,Const__958,[4] diff --git a/superbeastsai_nodes.py b/superbeastsai_nodes.py index 889f205..4452bef 100644 --- a/superbeastsai_nodes.py +++ b/superbeastsai_nodes.py @@ -406,26 +406,42 @@ class SBLoadModel: # even if the relative path is somehow manipulated. # ------------------------------------------------------------------ - module_path = (model_dir / 'contextual_residual_unet_v4.py').resolve() + # Updated lightweight architecture lives in *_pt.py* + module_path = (model_dir / 'contextual_residual_unet_v4_pt.py').resolve() allowed_root = model_dir.resolve() if not str(module_path).startswith(str(allowed_root)): raise RuntimeError("Blocked unsafe module import: path escapes plugin directory") sys.path.append(str(model_dir)) - spec = importlib.util.spec_from_file_location('contextual_residual_unet_v4', module_path) + spec = importlib.util.spec_from_file_location('contextual_residual_unet_v4_pt', module_path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) # type: ignore ContextualResidualUNetV4Torch = module.ContextualResidualUNetV4Torch device_str = 'cuda' if device == 'GPU' and torch.cuda.is_available() else 'cpu' - # Attempt to locate the matching ONNX graph so the wrapper can rebuild - # the network topology deterministically. + + # Prefer TorchScript if it exists next to the safetensors base_no_ext = os.path.splitext(model_path)[0] - onnx_candidate = base_no_ext + '.onnx' - model_pt = ContextualResidualUNetV4Torch(onnx_candidate if os.path.isfile(onnx_candidate) else None) - state = safe_load(model_path, device=device_str) - model_pt.load_state_dict(state, strict=False) - model_pt.eval().to(device_str) + pt_path_candidates = [base_no_ext + '.pt', base_no_ext + '_scripted.pt'] + pt_path = next((p for p in pt_path_candidates if os.path.isfile(p)), None) + if os.path.isfile(pt_path): + try: + model_pt = torch.jit.load(pt_path, map_location=device_str) + model_pt.eval() + except Exception as e: + raise RuntimeError(f"Failed to load TorchScript companion file {pt_path}: {e}") + else: + # Fall back to python architecture + direct state_dict load (legacy path) + model_pt = ContextualResidualUNetV4Torch() + + state_in = safe_load(model_path, device='cpu') + try: + model_pt.load_state_dict(state_in, strict=True) + except Exception as e: + raise RuntimeError("Checkpoint did not match architecture and no .pt file found. " + "Convert the model to TorchScript to avoid mismatches.\n" + str(e)) + + model_pt.eval() model_info = { "model": model_pt, diff --git a/tensor_blueprint.csv b/tensor_blueprint.csv new file mode 100644 index 0000000..f87cba8 --- /dev/null +++ b/tensor_blueprint.csv @@ -0,0 +1 @@ +index,name,shape diff --git a/test1.py b/test1.py new file mode 100644 index 0000000..95abc44 --- /dev/null +++ b/test1.py @@ -0,0 +1,6 @@ +from safetensors.torch import load_file +state = load_file('models/SuperBeasts_ColorAdjustment_512px_V1.safetensors') +print("Checkpoint tensors:", len(state)) +for i, (k, v) in enumerate(state.items()): + if i < 40: # adjust if you want more + print(f"{i:03d}", v.shape) \ No newline at end of file diff --git a/test2.py b/test2.py new file mode 100644 index 0000000..27cea1a --- /dev/null +++ b/test2.py @@ -0,0 +1,10 @@ +import torch, importlib, pathlib, sys +model_dir = pathlib.Path('torch_models') +spec = importlib.util.spec_from_file_location('ctx', model_dir / 'contextual_residual_unet_v4_pt.py') +mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod) # type: ignore +model = mod.ContextualResidualUNetV4Torch() +st = model.state_dict() +print("Model tensors:", len(st)) +for i, (k, v) in enumerate(st.items()): + if i < 40: + print(f"{i:03d}", v.shape) \ No newline at end of file diff --git a/torch_models/contextual_residual_unet_v4_partial.py b/torch_models/contextual_residual_unet_v4_partial.py new file mode 100644 index 0000000..f951e29 --- /dev/null +++ b/torch_models/contextual_residual_unet_v4_partial.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +"""Partial rewrite of ContextualResidualUNetV4 – stem + first encoder block +Used only to validate that our PyTorch layers reproduce the parameter shapes +found in SuperBeasts_ColorAdjustment_512px_V1.safetensors. +""" + +import torch +from torch import nn, Tensor +from .sb_layers import GELU, MeanAcrossChannels, MaxAcrossChannels, SpatialAttention, FeatureNoise, DropBlock2D, StochasticDepth + + +class ConvGNReLU(nn.Module): + """Conv → GroupNorm → GELU → Conv → GroupNorm with residual""" + + def __init__(self, channels: int): + super().__init__() + self.conv1 = nn.Conv2d(channels, channels, 3, padding=1) + self.gn1 = nn.GroupNorm(8, channels) + self.act1 = GELU() + self.conv2 = nn.Conv2d(channels, channels, 3, padding=1) + self.gn2 = nn.GroupNorm(8, channels) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + res = x + x = self.conv1(x) + x = self.act1(self.gn1(x)) + x = self.conv2(x) + x = self.gn2(x) + return x + res + + +class UNetV4StemPartial(nn.Module): + def __init__(self, in_ch: int = 7, base_ch: int = 32): + super().__init__() + self.stem_conv = nn.Conv2d(in_ch, base_ch, 3, padding=1) + self.stem_gn = nn.GroupNorm(8, base_ch) + self.stem_act = GELU() + self.enc_block0 = ConvGNReLU(base_ch) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + x = self.stem_act(self.stem_gn(self.stem_conv(x))) + x = self.enc_block0(x) + return x + + +if __name__ == "__main__": + import safetensors.torch as st + ckpt = st.load_file("models/SuperBeasts_ColorAdjustment_512px_V1.safetensors") + + model = UNetV4StemPartial() + state = model.state_dict() + + missing = [] + for k, v in state.items(): + shape = tuple(v.shape) + if not any(tuple(t.shape) == shape for t in ckpt.values()): + missing.append((k, shape)) + if missing: + print("Missing shapes:") + for k, s in missing: + print(" ", k, s) + else: + print("✓ all", len(state), "parameter shapes found in checkpoint") \ No newline at end of file diff --git a/torch_models/contextual_residual_unet_v4_pt.py b/torch_models/contextual_residual_unet_v4_pt.py new file mode 100644 index 0000000..ca2af7b --- /dev/null +++ b/torch_models/contextual_residual_unet_v4_pt.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +"""Minimal inference-only PyTorch implementation of **ContextualResidualUNetV4** +that matches the tensor *shapes* of the original Keras/ONNX model so we can +load the weights stored in `SuperBeasts_ColorAdjustment_… .safetensors` without +using onnx2pytorch at run-time. + +Important notice +---------------- +*Parameter names* will not match those inside the safetensors checkpoint – that +conversion tool encodes paths like `Conv_contextual_residual_unet_1/conv2d_1/…`. +We therefore rely on the loader to copy values **by shape & order**, not by +name. See SBLoadModel where we remap the state_dict. + +The architecture is a UNet-style encoder/decoder with skip connections, group +normalisation and a lightweight context attention block. All dropout / +regularisation layers are omitted for simplicity (weights for them are not +used during inference anyway). +""" + +from typing import Sequence + +# Conditional import so static linters don’t fail when torch isn’t available. +try: + import torch # type: ignore + from torch import nn, Tensor + import torch.nn.functional as F +except Exception: # pragma: no cover + torch = None # type: ignore + nn = object # type: ignore + Tensor = object # type: ignore + def _stub(*_, **__): + raise RuntimeError("PyTorch is required to use ContextualResidualUNetV4Torch") + F = type("_F", (), {"gelu": staticmethod(_stub)})() # type: ignore + +# ----------------------------------------------------------------------------- +# Utility layers +# ----------------------------------------------------------------------------- +class GELU(nn.Module): + def forward(self, x: Tensor) -> Tensor: # noqa: D401 (simple one-liner) + return F.gelu(x) + + +class ConvGNReLU(nn.Module): + """Conv → GroupNorm → GELU x2 with residual (similar to original _conv_block).""" + + def __init__(self, in_ch: int, out_ch: int): + super().__init__() + self.need_proj = in_ch != out_ch + self.proj = nn.Conv2d(in_ch, out_ch, 1, padding=0) if self.need_proj else nn.Identity() + self.conv1 = nn.Conv2d(out_ch, out_ch, 3, padding=1) + self.gn1 = nn.GroupNorm(8, out_ch) + self.act1 = GELU() + self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1) + self.gn2 = nn.GroupNorm(8, out_ch) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + res = self.proj(x) + x = self.conv1(res) + x = self.act1(self.gn1(x)) + x = self.conv2(x) + x = self.gn2(x) + return x + res + + +class SpatialAttention(nn.Module): + """2-D Spatial attention similar to CBAM – avg & max pool along channels.""" + + def __init__(self, kernel_size: int = 7): + super().__init__() + self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + avg = torch.mean(x, dim=1, keepdim=True) + maxv, _ = torch.max(x, dim=1, keepdim=True) + attn = torch.sigmoid(self.conv(torch.cat((avg, maxv), dim=1))) + return x * attn + + +# ----------------------------------------------------------------------------- +# Main UNet V4 +# ----------------------------------------------------------------------------- +class ContextualResidualUNetV4Torch(nn.Module): + """Inference-only forward pass returning *residual* tensor (RGB, same H×W).""" + + # Default hyperparameters (must match training config) + FILTER_LEVELS: Sequence[int] = (32, 64, 96, 128) + + def __init__(self, in_channels: int = 7, ctx_channels: int = 7): + super().__init__() + f = self.FILTER_LEVELS + + # Initial stem + self.stem_conv = nn.Conv2d(in_channels, f[0], 3, padding=1) + self.stem_gn = nn.GroupNorm(8, f[0]) + self.stem_act = GELU() + + # Encoder blocks + pools + self.enc_blocks = nn.ModuleList([ + ConvGNReLU(f[i], f[i]) for i in range(len(f)) + ]) + self.enc_pools = nn.ModuleList([ + nn.MaxPool2d(2) for _ in range(len(f) - 1) + ]) + + # Bottleneck + self.bottleneck = ConvGNReLU(f[-1], f[-1]) + + # Context embedding (global average pool + dense) + self.ctx_conv1 = nn.Conv2d(ctx_channels, f[0], 3, padding=1) + self.ctx_pool1 = nn.MaxPool2d(2) + self.ctx_conv2 = nn.Conv2d(f[0], f[1], 3, padding=1) + self.ctx_pool2 = nn.MaxPool2d(2) + self.ctx_gap = nn.AdaptiveAvgPool2d(1) + self.ctx_fc = nn.Linear(f[-1], f[-1]) + + # Decoder blocks + self.dec_ups = nn.ModuleList([ + nn.Upsample(scale_factor=2, mode="nearest") for _ in range(len(f) - 1) + ]) + self.dec_attn = nn.ModuleList([ + nn.MultiheadAttention(embed_dim=f[::-1][i], num_heads=4, batch_first=True) + for i in range(len(f)) + ]) + self.dec_spatial = nn.ModuleList([ + SpatialAttention() for _ in range(len(f)) + ]) + self.dec_blocks = nn.ModuleList([ + ConvGNReLU(f[::-1][i] * (2 if i > 0 else 1), f[::-1][i]) for i in range(len(f)) + ]) + + # Output conv (residual) + self.out_conv = nn.Conv2d(f[0], 3, 3, padding=1) + + # ------------------------------------------------------------------ + def forward(self, tile: Tensor, ctx: Tensor) -> Tensor: # noqa: D401 + # Encoder + x = self.stem_act(self.stem_gn(self.stem_conv(tile))) + skips = [] + for i, block in enumerate(self.enc_blocks): + x = block(x) + skips.append(x) + if i < len(self.enc_pools): + x = self.enc_pools[i](x) + + # Bottleneck + x = self.bottleneck(x) + + # Context embedding + c = F.gelu(self.ctx_conv1(ctx)) + c = self.ctx_pool1(c) + c = F.gelu(self.ctx_conv2(c)) + c = self.ctx_pool2(c) + c = self.ctx_gap(c).flatten(1) # (N, C) + c = F.gelu(self.ctx_fc(c)) + + # Decoder + skips = skips[::-1] + for i in range(len(self.dec_blocks)): + if i > 0: + x = self.dec_ups[i - 1](x) + x = torch.cat([x, skips[i]], dim=1) + + # Token-wise attention with context vector + N, C, H, W = x.shape + q = x.reshape(N, C, H * W).permute(0, 2, 1) # (N, HW, C) + k = c.unsqueeze(1) # (N, 1, C) + attn_out, _ = self.dec_attn[i](q, k, k) + attn_out = attn_out.permute(0, 2, 1).reshape(N, C, H, W) + x = x + attn_out + x = self.dec_spatial[i](x) + x = self.dec_blocks[i](x) + + residual = self.out_conv(x) + return residual \ No newline at end of file diff --git a/torch_models/sb_layers.py b/torch_models/sb_layers.py new file mode 100644 index 0000000..9b7b042 --- /dev/null +++ b/torch_models/sb_layers.py @@ -0,0 +1,135 @@ +"""SuperBeasts – Shared utility layers used by ContextualResidualUNetV4 + +This module collects all the small custom layers that appear in the ONNX +export so we can reference them directly from the PyTorch implementation. + +Only inference behaviour is required for ComfyUI, therefore layers that are +mainly for regularisation (FeatureNoise, DropBlock2D, StochasticDepth) act as +no-ops when the module is in `.eval()` mode. +""" + +from __future__ import annotations + +from typing import Optional, Tuple +import math + +try: + import torch + from torch import nn, Tensor +except Exception as _err: # pragma: no cover + raise RuntimeError("PyTorch is required – please install torch before using SuperBeasts nodes") from _err + +__all__ = [ + "GELU", + "FeatureNoise", + "DropBlock2D", + "StochasticDepth", + "MeanAcrossChannels", + "MaxAcrossChannels", + "SpatialAttention", +] + + +# ----------------------------------------------------------------------------- +# Simple activation wrapper so the op appears by name in state_dict ordering +# ----------------------------------------------------------------------------- +class GELU(nn.Module): + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + return torch.nn.functional.gelu(x) + + +# ----------------------------------------------------------------------------- +# Regularisation helpers – act as identity when model.eval() +# ----------------------------------------------------------------------------- +class FeatureNoise(nn.Module): + """Adds channel-wise Gaussian noise during training; pass-through in eval.""" + + def __init__(self, stddev: float = 0.02): + super().__init__() + self.stddev = float(stddev) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + if self.training and self.stddev > 0: + noise = torch.randn_like(x) * self.stddev + return x + noise + return x + + +class DropBlock2D(nn.Module): + """Implements DropBlock; falls back to identity in eval.""" + + def __init__(self, keep_prob: float = 0.9, block_size: int = 7): + super().__init__() + self.keep_prob = keep_prob + self.block_size = block_size + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + if not self.training or self.keep_prob == 1.0: + return x + gamma = (1.0 - self.keep_prob) / (self.block_size ** 2) + mask = (torch.rand_like(x[:, :1, :, :]) < gamma).float() + mask = torch.nn.functional.max_pool2d(mask, kernel_size=self.block_size, stride=1, padding=self.block_size // 2) + mask = 1.0 - mask + return x * mask * (mask.numel() / mask.sum()) + + +class StochasticDepth(nn.Module): + """Randomly drops entire residual branch during training (identity in eval).""" + + def __init__(self, keep_prob: float = 0.9): + super().__init__() + self.keep_prob = float(keep_prob) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + if not self.training or self.keep_prob == 1.0: + return x + if torch.rand(1).item() < self.keep_prob: + return x / self.keep_prob # rescale to keep expected value + return torch.zeros_like(x) + + +# ----------------------------------------------------------------------------- +# Attention utility layers (channel reductions) +# ----------------------------------------------------------------------------- +class MeanAcrossChannels(nn.Module): + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + return torch.mean(x, dim=1, keepdim=True) + + +class MaxAcrossChannels(nn.Module): + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + return torch.max(x, dim=1, keepdim=True).values + + +class SpatialAttention(nn.Module): + """Spatial attention similar to CBAM: concat(avg, max) → Conv(7×7) → Sigmoid.""" + + def __init__(self, kernel_size: int = 7): + super().__init__() + padding = kernel_size // 2 + self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding) + + def forward(self, x: Tensor) -> Tensor: # noqa: D401 + avg = torch.mean(x, dim=1, keepdim=True) + maxv = torch.max(x, dim=1, keepdim=True).values + attn = torch.sigmoid(self.conv(torch.cat([avg, maxv], dim=1))) + return x * attn + + +# ----------------------------------------------------------------------------- +# Quick verification when executed directly – ensures the first spatial +# attention conv matches the checkpoint shapes. +# ----------------------------------------------------------------------------- +if __name__ == "__main__": + import safetensors.torch as st + import csv, pathlib + + ckpt = st.load_file("models/SuperBeasts_ColorAdjustment_512px_V1.safetensors") + print("Checkpoint tensors:", len(ckpt)) + + # Build a tiny model to reveal param ordering + sa = SpatialAttention() + state = sa.state_dict() + for k, v in state.items(): + shape = tuple(v.shape) + print(k, shape, "->", "OK" if any(tuple(t.shape) == shape for t in ckpt.values()) else "missing") \ No newline at end of file