Blender Python OpenGL compute shaders

gl_compute.py

import bgl
import bpy
import numpy as np

def install_lib(libname):
    from subprocess import call
    pp = bpy.app.binary_path_python
    call([pp, "-m", "ensurepip", "--user"])
    call([pp, "-m", "pip", "install", "--user", libname])
    
# uncomment to install required lib
# install_lib("moderngl")
import moderngl

def get_teximage(context):
    teximage = None
    for area in context.screen.areas:
        if area.type == "IMAGE_EDITOR":
            teximage = area.spaces.active.image
            break
    if teximage is not None and teximage.size[1] != 0:
        return teximage
    else:
        return None

print("OpenGL supported version (by Blender):", bgl.glGetString(bgl.GL_VERSION))
ctx = moderngl.create_context(require=430)
print("GL context version code:", ctx.version_code)
assert ctx.version_code >= 430
print("Compute max work group size:", ctx.info['GL_MAX_COMPUTE_WORK_GROUP_SIZE'], end='\n\n')

basic_shader = """
#version 430 
#define TILE_WIDTH 8
#define TILE_HEIGHT 8

const ivec2 tileSize = ivec2(TILE_WIDTH, TILE_HEIGHT);

layout(local_size_x=TILE_WIDTH, local_size_y=TILE_HEIGHT) in;
layout(binding=0) writeonly buffer out_0 { vec4 outp[][1]; };
layout(binding=1) readonly buffer in_0 { vec4 inp[][1]; };

uniform uint img_size;

void main() {
    const ivec2 tile_xy = ivec2(gl_WorkGroupID);
    const ivec2 thread_xy = ivec2(gl_LocalInvocationID);
    const ivec2 pixel_xy = tile_xy * tileSize + thread_xy;
    
    const float tx = pixel_xy.x;
    const float ty = pixel_xy.y;
    
    vec4 outc = vec4(tx/img_size, ty/img_size, int(tx) & int(ty), 1.0);
    //outc = inp[pixel_xy.x+5][(pixel_xy.y) * img_size] * 0.9 + outc * 0.1;
    outc = outc;
    
    outp[pixel_xy.x][pixel_xy.y * img_size] = outc;
        
} 
"""

image = get_teximage(bpy.context)

sourcepixels = np.float32(np.array(image.pixels).reshape(image.size[0], image.size[1], 4))

compute_shader = ctx.compute_shader(basic_shader)

print("start compute")
buffer = ctx.buffer(np.empty((image.size[0], image.size[1], 4), dtype=np.float32))
buffer.bind_to_storage_buffer(0)
in_buf = ctx.buffer(sourcepixels)
in_buf.bind_to_storage_buffer(1)
compute_shader.get("img_size", 5).value = image.size[0]
compute_shader.run(group_x=image.size[0]//8, group_y=image.size[1]//8)
print("end compute")

image.pixels = np.frombuffer(buffer.read(), dtype=np.float32)

print("fin.")