Extract Color Palette

# Python 3.10+: pip install comfy-sdk==0.4.0
import os

from comfy_sdk import Comfy

api_key = os.environ.get("COMFY_API_KEY", "").strip()
if not api_key: raise RuntimeError("Set COMFY_API_KEY")
os.environ.setdefault("COMFY_BASE_URL", "https://cloud.comfy.org")
client = Comfy(api_key=api_key)
workflow = client.workflows.from_json({
    "148": {
        "class_type": "LoadImage",
        "inputs": {
            "image": "UPLOADED_image_FILENAME"
        },
        "_meta": {
            "title": "Load Image"
        }
    },
    "149": {
        "class_type": "GLSLShader",
        "inputs": {
            "fragment_shader": "#version 300 es\n// kmeans_palette.glsl — GPU k-means dominant colors -> vertical stripes.\n//\n// INPUTS\n//   u_image0  — source image (required)\n//   u_int0    — stripe count K (try this first)\n//   u_float0  — stripe count K (use this if u_int0 doesn't update; this node\n//               forwards floats more reliably than ints)\n//   If neither input forwards a value, K falls back to DEFAULT_K below.\n//\n// OUTPUT SIZE: set on the node via size_mode -> custom (width / height),\n//   not in the shader. Keep it >= ~64x64.\n\n#pragma passes 10\n\nprecision highp float;\n\nconst int DEFAULT_K  = 8;     // used only if no count input forwards\nconst int MAX_K      = 16;\nconst int NUM_PASSES = 10;    // keep equal to `#pragma passes` above\nconst int GRID       = 48;    // cached source samples per axis\n\nuniform sampler2D u_image0;   // pass 0: SOURCE; pass 1+: ping-pong buffer\nuniform vec2  u_resolution;\nuniform int   u_int0;         // K (preferred)\nuniform float u_float0;       // K (fallback path)\nuniform int   u_pass;\n\nin vec2 v_texCoord;\nlayout(location = 0) out vec4 fragColor0;\n\nint resolveK() {\n    int k = 0;\n    if (u_int0 > 0)             k = u_int0;            // try the int input\n    else if (u_float0 > 0.5)    k = int(u_float0 + 0.5); // else the float input\n    else                        k = DEFAULT_K;         // else the constant\n    return clamp(k, 2, MAX_K);\n}\n\nvec3 readCentroid(int i)        { return texelFetch(u_image0, ivec2(i, 0), 0).rgb; }\nvec3 readCache(int sx, int sy)  { return texelFetch(u_image0, ivec2(sx, sy + 1), 0).rgb; }\n\nint nearestCentroid(vec3 px, int K) {\n    float best = 1e9;\n    int bi = 0;\n    for (int j = 0; j < MAX_K; j++) {\n        if (j >= K) break;\n        vec3 d = px - readCentroid(j);\n        float dist = dot(d, d);\n        if (dist < best) { best = dist; bi = j; }\n    }\n    return bi;\n}\n\nvoid main() {\n    int K = resolveK();\n    ivec2 frag = ivec2(gl_FragCoord.xy);\n\n    // PASS 0: cache a GRID x GRID downsample of the source (rows 1..GRID) and\n    //         seed K centroids (row 0). On this pass u_image0 IS the source.\n    if (u_pass == 0) {\n        if (frag.y == 0 && frag.x < K) {\n            float u = (float(frag.x) + 0.5) / float(K);\n            float v = fract(u * 2.6180339887);\n            fragColor0 = vec4(texture(u_image0, vec2(u, v)).rgb, 1.0);\n        } else if (frag.y >= 1 && frag.y <= GRID && frag.x < GRID) {\n            vec2 uv = (vec2(float(frag.x), float(frag.y - 1)) + 0.5) / float(GRID);\n            fragColor0 = vec4(texture(u_image0, uv).rgb, 1.0);\n        } else {\n            fragColor0 = vec4(0.0, 0.0, 0.0, 1.0);\n        }\n        return;\n    }\n\n    // PASSES 1..N-2: one k-means iteration\n    if (u_pass < NUM_PASSES - 1) {\n        if (frag.y == 0 && frag.x < K) {\n            int c = frag.x;\n            vec3 sum = vec3(0.0);\n            float cnt = 0.0;\n            for (int sy = 0; sy < GRID; sy++) {\n                for (int sx = 0; sx < GRID; sx++) {\n                    vec3 px = readCache(sx, sy);\n                    if (nearestCentroid(px, K) == c) { sum += px; cnt += 1.0; }\n                }\n            }\n            fragColor0 = vec4(cnt > 0.0 ? sum / cnt : readCentroid(c), 1.0);\n        } else {\n            fragColor0 = texelFetch(u_image0, frag, 0);\n        }\n        return;\n    }\n\n    // FINAL PASS: K vertical stripes\n    int band = clamp(int(v_texCoord.x * float(K)), 0, K - 1);\n    fragColor0 = vec4(readCentroid(band), 1.0);\n}",
            "images.image0": [
                "148",
                0
            ],
            "ints.u_int0": [
                "150",
                0
            ],
            "size_mode": "custom",
            "size_mode.height": [
                "153",
                1
            ],
            "size_mode.width": [
                "153",
                0
            ]
        },
        "_meta": {
            "title": "GLSL Shader"
        }
    },
    "150": {
        "class_type": "PrimitiveInt",
        "inputs": {
            "value": 12
        },
        "_meta": {
            "title": "Palette Size"
        }
    },
    "151": {
        "class_type": "GLSLShader",
        "inputs": {
            "fragment_shader": "#version 300 es\n// kmeans_palette_weighted.glsl — GPU k-means dominant colors as vertical stripes,\n// with each stripe's WIDTH proportional to that color's share of the image.\n// Widths are snapped to 2.5% increments (largest-remainder), so a colour that's\n// a tiny fraction rounds to 0% (disappears) or 2.5% min — no negligible slivers.\n//\n// INPUTS\n//   u_image0 — source image (required)\n//   u_int0   — max number of colors K (drive with the stock \"Int\" node)\n//   u_float0 — same, fallback if u_int0 doesn't forward\n//   If neither forwards, K = DEFAULT_K below.\n//\n// OUTPUT SIZE: set on the node via size_mode -> custom. Keep >= ~64x64.\n//\n// NOTE: stripes are NOT equal width here, so ShaderPaletteToHex (which samples\n// equal bands) won't line up — use the equal-width kmeans_palette.glsl for that.\n\n#pragma passes 10\n\nprecision highp float;\n\nconst int   DEFAULT_K  = 8;\nconst int   MAX_K      = 16;\nconst int   NUM_PASSES = 10;     // keep equal to `#pragma passes`\nconst int   GRID       = 48;     // cached source samples per axis\nconst int   STEPS      = 40;     // 100% / 2.5%  -> stripe widths are multiples of 2.5%\nconst float TOTAL      = float(GRID * GRID);\n\nuniform sampler2D u_image0;      // pass 0: SOURCE; pass 1+: ping-pong buffer\nuniform vec2  u_resolution;\nuniform int   u_int0;\nuniform float u_float0;\nuniform int   u_pass;\n\nin vec2 v_texCoord;\nlayout(location = 0) out vec4 fragColor0;\n\nint resolveK() {\n    int k = u_int0 > 0 ? u_int0 : (u_float0 > 0.5 ? int(u_float0 + 0.5) : DEFAULT_K);\n    return clamp(k, 2, MAX_K);\n}\n\n// row 0: centroids — rgb = colour, a = cluster's share of samples (0..1)\n// rows 1..GRID: cached source downsample\nvec3  readColor(int i)         { return texelFetch(u_image0, ivec2(i, 0), 0).rgb; }\nfloat readShare(int i)         { return texelFetch(u_image0, ivec2(i, 0), 0).a; }\nvec3  readCache(int sx, int sy){ return texelFetch(u_image0, ivec2(sx, sy + 1), 0).rgb; }\n\nint nearestCentroid(vec3 px, int K) {\n    float best = 1e9;\n    int bi = 0;\n    for (int j = 0; j < MAX_K; j++) {\n        if (j >= K) break;\n        vec3 d = px - readColor(j);\n        float dist = dot(d, d);\n        if (dist < best) { best = dist; bi = j; }\n    }\n    return bi;\n}\n\nvoid main() {\n    int K = resolveK();\n    ivec2 frag = ivec2(gl_FragCoord.xy);\n\n    // PASS 0: cache source (rows 1..GRID) + seed centroids (row 0)\n    if (u_pass == 0) {\n        if (frag.y == 0 && frag.x < K) {\n            float u = (float(frag.x) + 0.5) / float(K);\n            float v = fract(u * 2.6180339887);\n            fragColor0 = vec4(texture(u_image0, vec2(u, v)).rgb, 0.0);\n        } else if (frag.y >= 1 && frag.y <= GRID && frag.x < GRID) {\n            vec2 uv = (vec2(float(frag.x), float(frag.y - 1)) + 0.5) / float(GRID);\n            fragColor0 = vec4(texture(u_image0, uv).rgb, 1.0);\n        } else {\n            fragColor0 = vec4(0.0);\n        }\n        return;\n    }\n\n    // PASSES 1..N-2: k-means iteration; store cluster share in alpha\n    if (u_pass < NUM_PASSES - 1) {\n        if (frag.y == 0 && frag.x < K) {\n            int c = frag.x;\n            vec3 sum = vec3(0.0);\n            float cnt = 0.0;\n            for (int sy = 0; sy < GRID; sy++) {\n                for (int sx = 0; sx < GRID; sx++) {\n                    vec3 px = readCache(sx, sy);\n                    if (nearestCentroid(px, K) == c) { sum += px; cnt += 1.0; }\n                }\n            }\n            vec3 col = cnt > 0.0 ? sum / cnt : readColor(c);\n            fragColor0 = vec4(col, cnt / TOTAL);   // alpha = share of samples\n        } else {\n            fragColor0 = texelFetch(u_image0, frag, 0);\n        }\n        return;\n    }\n\n    // FINAL PASS: weighted vertical stripes, widths snapped to 2.5% slots.\n    // Largest-remainder allocation of STEPS slots so they sum to exactly 100%.\n    int   floorSlots[MAX_K];\n    float rem[MAX_K];\n    int   used = 0;\n    for (int j = 0; j < MAX_K; j++) {\n        float share = (j < K) ? readShare(j) : 0.0;\n        if (share > 0.0) {\n            float rs = share * float(STEPS);\n            int fs = int(floor(rs));\n            floorSlots[j] = fs;\n            rem[j] = rs - float(fs);\n            used += fs;\n        } else {\n            floorSlots[j] = 0;\n            rem[j] = -1.0;             // empty/unused cluster never gets a slot\n        }\n    }\n    int remaining = STEPS - used;     // leftover slots (always fills to STEPS)\n\n    // Each leftover slot goes to a largest-remainder cluster (rank-based, stable).\n    int slots[MAX_K];\n    for (int j = 0; j < MAX_K; j++) {\n        int rank = 0;\n        for (int m = 0; m < MAX_K; m++) {\n            if (rem[m] > rem[j] || (rem[m] == rem[j] && m < j)) rank++;\n        }\n        slots[j] = floorSlots[j] + (rank < remaining ? 1 : 0);\n    }\n\n    // Render order: largest weight on the left -> smallest on the right.\n    // Rank each cluster by slot count (desc, tie-break by index) to get a stable\n    // left->right order, then place each cluster at its rank position.\n    int order_[MAX_K];\n    for (int j = 0; j < MAX_K; j++) {\n        int rnk = 0;\n        for (int m = 0; m < MAX_K; m++) {\n            if (slots[m] > slots[j] || (slots[m] == slots[j] && m < j)) rnk++;\n        }\n        order_[rnk] = j;\n    }\n\n    // Which stripe does this pixel's x fall into, walking left->right by weight?\n    float xslot = v_texCoord.x * float(STEPS);\n    int cum = 0;\n    int chosen = order_[0];           // widest, leftmost (safety default)\n    for (int r = 0; r < MAX_K; r++) {\n        cum += slots[order_[r]];\n        if (xslot < float(cum)) { chosen = order_[r]; break; }\n    }\n    fragColor0 = vec4(readColor(chosen), 1.0);\n}\n",
            "images.image0": [
                "148",
                0
            ],
            "ints.u_int0": [
                "150",
                0
            ],
            "size_mode": "custom",
            "size_mode.height": [
                "153",
                1
            ],
            "size_mode.width": [
                "153",
                0
            ]
        },
        "_meta": {
            "title": "GLSL Shader"
        }
    },
    "152": {
        "class_type": "ComfySwitchNode",
        "inputs": {
            "on_false": [
                "149",
                0
            ],
            "on_true": [
                "151",
                0
            ],
            "switch": True
        },
        "_meta": {
            "title": "If/Else Switch"
        }
    },
    "153": {
        "class_type": "ResolutionSelector",
        "inputs": {
            "aspect_ratio": "16:9 (Widescreen)",
            "megapixels": 1,
            "multiple": 8
        },
        "_meta": {
            "title": "Resolution Selector"
        }
    },
    "154": {
        "class_type": "SaveImageAdvanced",
        "inputs": {
            "filename_prefix": "extracted_color_palette",
            "format": "png",
            "format.bit_depth": "8-bit",
            "format.input_color_space": "sRGB",
            "images": [
                "152",
                0
            ]
        },
        "_meta": {
            "title": "Save Image (Advanced)"
        }
    }
})
# Replace each example.com URL with a public URL of your file.
workflow.set_input("148", "image", client.assets.from_url("https://example.com/replace-with-a-url/image"))

# run() also passes api_key to the partner nodes in this workflow.
job = client.run(workflow, api_key=api_key)