API tools faq
paste
here2share

# tk_WebGPU_Metaballs.py

here2share
Jul 1st, 2026
61
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 13.28 KB | None | 0 0
raw download clone embed print report
  1. # tk_WebGPU_Metaballs.py
  2.  
  3. # do not remove these comments and do no add comments
  4. # do not deviate from instructions given by these comments
  5.  
  6. import time
  7. import subprocess
  8. from http.server import BaseHTTPRequestHandler, HTTPServer
  9. import tkinter as tk
  10. import collections
  11.  
  12. SZ = 640
  13.  
  14. root = tk.Tk()
  15. root.title("tk WebGPU Metaballs")
  16. root.geometry(f"{SZ}x{SZ}+0+0")
  17. root.resizable(False, False)
  18.  
  19. canvas = tk.Canvas(root, width=SZ, height=SZ, bg="black", highlightthickness=0)
  20. canvas.pack()
  21.  
  22. core_status = "-- PROCESSING BATCH FOR THE ANIMATION --"
  23. img = tk.PhotoImage(width=SZ, height=SZ)
  24. canvas_image_id = canvas.create_image(0, 0, anchor=tk.NW, image=img)
  25. status_text = canvas.create_text(
  26.     320,
  27.     320,
  28.     text=core_status,
  29.     font=("Arial", 18, "bold"),
  30.     fill="#00FF00",
  31.     justify="center"
  32. )
  33.  
  34. PORT = 54321
  35. FRAME_SIZE = SZ * SZ * 3
  36. PPM_HEADER = b"P6\n640 640\n255\n"
  37. current_frame = b"\x00" * FRAME_SIZE
  38. frame_queue = collections.deque(maxlen=300)
  39. latest_status = "Initializing..."
  40. server_ready = False
  41. chrome_proc = None
  42.  
  43. HTML_CONTENT = """<!DOCTYPE html>
  44. <html>
  45. <body style="margin: 0; background: black;">
  46.    <script type="module">
  47.        let time = 0.0;
  48.        const width = 640;
  49.        const height = 640;
  50.        const frameSize = width * height * 3;
  51.        
  52.        let localQueue = [];
  53.        let inflightFrames = 0;
  54.        let gpuContext = null;
  55.  
  56.        async function logStatus(msg) {
  57.            console.log(msg);
  58.            try {
  59.                await fetch('/log', { method: 'POST', body: msg });
  60.            } catch(e) {}
  61.        }
  62.  
  63.        const shaderSource = `
  64.            struct Uniforms {
  65.                time: f32,
  66.                dt: f32,
  67.            }
  68.            @group(0) @binding(0) var<uniform> uniforms: Uniforms;
  69.            @group(0) @binding(1) var<storage, read_write> outputBuffer: array<vec4f>;
  70.  
  71.            fn opSmoothUnion(d1: f32, d2: f32, k: f32) -> f32 {
  72.                let h = clamp(0.5 + 0.5 * (d2 - d1) / k, 0.0, 1.0);
  73.                return mix(d2, d1, h) - k * h * (1.0 - h);
  74.            }
  75.  
  76.            fn sdSphere(p: vec3f, s: f32) -> f32 {
  77.                return length(p) - s;
  78.            }
  79.  
  80.            fn map(p: vec3f, iTime: f32) -> f32 {
  81.                var d = 2.0;
  82.                for (var i = 0; i < 16; i++) {
  83.                    let fi = f32(i);
  84.                    let t = iTime * (fract(fi * 412.531 + 0.513) - 0.5) * 2.0;
  85.                    let offset = sin(t + fi * vec3f(52.5126, 64.62744, 632.25)) * vec3f(2.0, 2.0, 0.8);
  86.                    let radius = mix(0.5, 1.0, fract(fi * 412.531 + 0.5124));
  87.                    d = opSmoothUnion(sdSphere(p + offset, radius), d, 0.4);
  88.                }
  89.                return d;
  90.            }
  91.  
  92.            fn calcNormal(p: vec3f, iTime: f32) -> vec3f {
  93.                let h = 1e-5;
  94.                let k = vec2f(1.0, -1.0);
  95.                return normalize(
  96.                    k.xyy * map(p + k.xyy * h, iTime) +
  97.                    k.yyx * map(p + k.yyx * h, iTime) +
  98.                    k.yxy * map(p + k.yxy * h, iTime) +
  99.                    k.xxx * map(p + k.xxx * h, iTime)
  100.                );
  101.            }
  102.  
  103.            @compute @workgroup_size(16, 16, 1)
  104.            fn main(@builtin(global_invocation_id) id: vec3u) {
  105.                if (id.x >= 640 || id.y >= 640) {
  106.                    return;
  107.                }
  108.                let iTime = uniforms.time;
  109.                let uv = vec2f(f32(id.x) / 640.0, f32(id.y) / 640.0);
  110.                let rayOri = vec3f((uv - 0.5) * vec2f(1.0, 1.0) * 6.0, 3.0);
  111.                let rayDir = vec3f(0.0, 0.0, -1.0);
  112.  
  113.                var depth = 0.0;
  114.                var p = vec3f(0.0);
  115.                for (var i = 0; i < 64; i++) {
  116.                    p = rayOri + rayDir * depth;
  117.                    let dist = map(p, iTime);
  118.                    depth += dist;
  119.                    if (dist < 1e-6) {
  120.                        break;
  121.                    }
  122.                }
  123.                depth = min(6.0, depth);
  124.                let n = calcNormal(p, iTime);
  125.                let b = max(0.0, dot(n, vec3f(0.577)));
  126.                var col = (0.5 + 0.5 * cos((b + iTime * 0.5) + uv.xyx * 2.0 + vec3f(0.0, 2.0, 4.0))) * (0.85 + b * 0.35);
  127.                col *= exp(-depth * 0.15);
  128.  
  129.                let pixelIndex = id.y * 640 + id.x;
  130.                outputBuffer[pixelIndex] = vec4f(col, 1.0);
  131.            }
  132.        `;
  133.  
  134.        async function initWebGPU() {
  135.            await logStatus("Initializing WebGPU...");
  136.            if (!navigator.gpu) {
  137.                await logStatus("WebGPU not supported on this browser.");
  138.                return null;
  139.            }
  140.            await logStatus("Requesting adapter...");
  141.            const adapter = await navigator.gpu.requestAdapter();
  142.            if (!adapter) {
  143.                await logStatus("No WebGPU adapter available.");
  144.                return null;
  145.            }
  146.            await logStatus("Requesting device...");
  147.            const device = await adapter.requestDevice();
  148.            await logStatus("Device obtained.");
  149.            await logStatus("Compiling shader module...");
  150.            const shaderModule = device.createShaderModule({ code: shaderSource });
  151.            await logStatus("Allocating buffers...");
  152.  
  153.            const totalElements = width * height;
  154.            const storageBufferSize = totalElements * 16;
  155.            const stagingBufferSize = totalElements * 16;
  156.  
  157.            const storageBuffer = device.createBuffer({
  158.                size: storageBufferSize,
  159.                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
  160.            });
  161.            const stagingBuffer = device.createBuffer({
  162.                size: stagingBufferSize,
  163.                usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
  164.            });
  165.            const uniformBuffer = device.createBuffer({
  166.                size: 8,
  167.                usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
  168.            });
  169.            await logStatus("Creating bind group layout...");
  170.            const bindGroupLayout = device.createBindGroupLayout({
  171.                entries: [
  172.                    { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
  173.                    { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } }
  174.                ]
  175.            });
  176.            await logStatus("Creating bind group...");
  177.            const bindGroup = device.createBindGroup({
  178.                layout: bindGroupLayout,
  179.                entries: [
  180.                    { binding: 0, resource: { buffer: uniformBuffer } },
  181.                    { binding: 1, resource: { buffer: storageBuffer } }
  182.                ]
  183.            });
  184.            await logStatus("Creating pipeline layout...");
  185.            const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] });
  186.            await logStatus("Building compute pipeline...");
  187.            const pipeline = device.createComputePipeline({
  188.                layout: pipelineLayout,
  189.                compute: { module: shaderModule, entryPoint: "main" }
  190.            });
  191.            await logStatus("Pipeline configuration ready.");
  192.            return { device, pipeline, bindGroup, uniformBuffer, storageBuffer, stagingBuffer };
  193.        }
  194.  
  195.        function runComputePipeline() {
  196.            if (localQueue.length + inflightFrames >= 300 || !gpuContext) return;
  197.            inflightFrames++;
  198.  
  199.            const timeData = new Float32Array([time, 1.0]);
  200.            time += 0.015;
  201.  
  202.            gpuContext.device.queue.writeBuffer(gpuContext.uniformBuffer, 0, timeData);
  203.  
  204.            const commandEncoder = gpuContext.device.createCommandEncoder();
  205.            const passEncoder = commandEncoder.beginComputePass();
  206.            passEncoder.setPipeline(gpuContext.pipeline);
  207.            passEncoder.setBindGroup(0, gpuContext.bindGroup);
  208.            passEncoder.dispatchWorkgroups(Math.ceil(width / 16), Math.ceil(height / 16), 1);
  209.            passEncoder.end();
  210.  
  211.            commandEncoder.copyBufferToBuffer(gpuContext.storageBuffer, 0, gpuContext.stagingBuffer, 0, width * height * 16);
  212.            gpuContext.device.queue.submit([commandEncoder.finish()]);
  213.  
  214.            const targetStaging = gpuContext.stagingBuffer;
  215.            targetStaging.mapAsync(GPUMapMode.READ).then(() => {
  216.                const mapped = new Float32Array(targetStaging.getMappedRange());
  217.                const rgbData = new Uint8Array(frameSize);
  218.                let rgbIdx = 0;
  219.                for (let i = 0; i < mapped.length; i += 4) {
  220.                    rgbData[rgbIdx]     = Math.min(255, Math.max(0, mapped[i] * 255)) | 0;
  221.                    rgbData[rgbIdx + 1] = Math.min(255, Math.max(0, mapped[i + 1] * 255)) | 0;
  222.                    rgbData[rgbIdx + 2] = Math.min(255, Math.max(0, mapped[i + 2] * 255)) | 0;
  223.                    rgbIdx += 3;
  224.                }
  225.                targetStaging.unmap();
  226.                inflightFrames--;
  227.  
  228.                if (localQueue.length < 300) {
  229.                    localQueue.push(rgbData);
  230.                }
  231.            }).catch(() => {
  232.                inflightFrames--;
  233.            });
  234.        }
  235.  
  236.        async function orchestratorLoop() {
  237.            while (localQueue.length + inflightFrames < 300 && inflightFrames < 4) {
  238.                runComputePipeline();
  239.            }
  240.  
  241.            if (localQueue.length > 0) {
  242.                try {
  243.                    let checkResponse = await fetch('/queue_status');
  244.                    if (checkResponse.ok) {
  245.                        const pyQueueSize = parseInt(await checkResponse.text());
  246.                        if (pyQueueSize < 60 && localQueue.length > 0) {
  247.                            const nextFrame = localQueue.shift();
  248.                            fetch('/data', { method: 'POST', body: nextFrame }).catch(() => {});
  249.                        }
  250.                    }
  251.                } catch (err) {
  252.                    window.close();
  253.                    return;
  254.                }
  255.            }
  256.  
  257.            requestAnimationFrame(orchestratorLoop);
  258.        }
  259.  
  260.        async function boot() {
  261.            await logStatus("Starting WebGPU initialization...");
  262.            gpuContext = await initWebGPU();
  263.            if (gpuContext) {
  264.                await logStatus("WebGPU initialization complete.");
  265.                orchestratorLoop();
  266.            }
  267.        }
  268.  
  269.        window.onload = boot;
  270.    </script>
  271. </body>
  272. </html>
  273. """
  274.  
  275. def make_handler():
  276.     def handler(request, client_address, server):
  277.         class FastHandler(BaseHTTPRequestHandler):
  278.             def do_GET(self):
  279.                 if self.path == "/queue_status":
  280.                     self.send_response(200)
  281.                     self.send_header("Content-Type", "text/plain")
  282.                     self.end_headers()
  283.                     self.wfile.write(str(len(frame_queue)).encode("utf-8"))
  284.                 else:
  285.                     self.send_response(200)
  286.                     self.send_header("Content-Type", "text/html")
  287.                     self.end_headers()
  288.                     self.wfile.write(HTML_CONTENT.encode("utf-8"))
  289.  
  290.             def do_POST(self):
  291.                 global latest_status
  292.                 if self.path == "/data":
  293.                     content_length = int(self.headers["Content-Length"])
  294.                     raw_data = self.rfile.read(content_length)
  295.                     for i in range(0, content_length, FRAME_SIZE):
  296.                         frame_slice = raw_data[i:i + FRAME_SIZE]
  297.                         if len(frame_slice) == FRAME_SIZE:
  298.                             frame_queue.append(frame_slice)
  299.                     self.send_response(200)
  300.                     self.end_headers()
  301.                 elif self.path == "/log":
  302.                     content_length = int(self.headers["Content-Length"])
  303.                     latest_status = self.rfile.read(content_length).decode("utf-8")
  304.                     self.send_response(200)
  305.                     self.end_headers()
  306.  
  307.             def log_message(self, format, *args):
  308.                 pass
  309.  
  310.         return FastHandler(request, client_address, server)
  311.     return handler
  312.  
  313. HTTPServer.allow_reuse_address = True
  314. httpd = HTTPServer(("127.0.0.1", PORT), make_handler())
  315. httpd.timeout = 0.0
  316.  
  317. def launch_browser():
  318.     global chrome_proc
  319.     chrome_path = r"C:\Program Files\Google\Chrome\Application\chrome.exe"
  320.     chrome_flags = [
  321.         chrome_path,
  322.         f"http://127.0.0.1:{PORT}",
  323.         "--headless=new",
  324.         "--no-sandbox",
  325.         "--enable-unsafe-webgpu",
  326.         "--enable-features=WebGPUService"
  327.     ]
  328.     chrome_proc = subprocess.Popen(chrome_flags, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  329.  
  330. def update_gui_status():
  331.     canvas.itemconfig(status_text, text=core_status + "\n" + latest_status)
  332.  
  333. # do not change the code that follows after this line...
  334. launch_browser()
  335. while len(frame_queue) < 60:
  336.     httpd.handle_request()
  337.     update_gui_status()
  338.     root.update()
  339.  
  340. wait = 0
  341. canvas.itemconfig(status_text, text='')
  342. while 1:
  343.     httpd.handle_request()
  344.     if frame_queue and time.time() > wait:
  345.         wait = time.time() + 0.12
  346.         current_frame = frame_queue.popleft()
  347.         img.put(PPM_HEADER + current_frame)
  348.         canvas.itemconfig(canvas_image_id, image=img)
  349.     root.update()
Advertisement
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!