import numpy as np import torch import os import sys import OpenGL.GL as gl import glfw SHADERTOY_HEADER = """ #version 440 precision highp float; uniform vec3 iResolution; uniform vec4 iMouse; uniform float iTime; uniform float iTimeDelta; uniform float iFrameRate; uniform int iFrame; uniform sampler2D iChannel0; uniform sampler2D iChannel1; uniform sampler2D iChannel2; uniform sampler2D iChannel3; #define texture2D texture """ SHADERTOY_FOOTER = """ layout(location = 0) out vec4 _fragColor; void main() { mainImage(_fragColor, gl_FragCoord.xy); } """ SHADERTOY_DEFAULT = """ void mainImage( out vec4 fragColor, in vec2 fragCoord ) { // Normalized pixel coordinates (from 0 to 1) vec2 uv = fragCoord/iResolution.xy; // Time varying pixel color vec3 col = 0.5 + 0.5*cos(iTime+uv.xyx+vec3(0,2,4)); // Output to screen fragColor = vec4(col,1.0); } """ def render_surface_and_context_init(width, height): if not glfw.init(): raise RuntimeError("GLFW did not init") glfw.window_hint(glfw.VISIBLE, glfw.FALSE) # hidden window = glfw.create_window(width, height, "hidden", None, None) if not window: raise RuntimeError("GLFW did not init window") glfw.make_context_current(window) return {} def render_surface_and_context_deinit(**kwargs): glfw.terminate() def compile_shader(source, shader_type): shader = gl.glCreateShader(shader_type) gl.glShaderSource(shader, source) gl.glCompileShader(shader) if gl.glGetShaderiv(shader, gl.GL_COMPILE_STATUS) != gl.GL_TRUE: raise RuntimeError(gl.glGetShaderInfoLog(shader)) return shader def compile_program(vertex_source, fragment_source): vertex_shader = compile_shader(vertex_source, gl.GL_VERTEX_SHADER) fragment_shader = compile_shader(fragment_source, gl.GL_FRAGMENT_SHADER) program = gl.glCreateProgram() gl.glAttachShader(program, vertex_shader) gl.glAttachShader(program, fragment_shader) gl.glLinkProgram(program) if gl.glGetProgramiv(program, gl.GL_LINK_STATUS) != gl.GL_TRUE: raise RuntimeError(gl.glGetProgramInfoLog(program)) return program def setup_framebuffer(width, height): texture = gl.glGenTextures(1) gl.glBindTexture(gl.GL_TEXTURE_2D, texture) gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, width, height, 0, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, None) gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR) gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR) fbo = gl.glGenFramebuffers(1) gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, fbo) gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, texture, 0) if gl.glCheckFramebufferStatus(gl.GL_FRAMEBUFFER) != gl.GL_FRAMEBUFFER_COMPLETE: raise RuntimeError("Framebuffer is not complete") return fbo, texture def setup_render_resources(width, height, fragment_source: str): ctx = render_surface_and_context_init(width, height) vertex_source = """ #version 330 core void main() { vec2 verts[3] = vec2[](vec2(-1, -1), vec2(3, -1), vec2(-1, 3)); gl_Position = vec4(verts[gl_VertexID], 0, 1); } """ shader = compile_program(vertex_source, fragment_source) fbo, texture = setup_framebuffer(width, height) textures = gl.glGenTextures(4) return (ctx, fbo, shader, textures) def render_resources_cleanup(ctx): # assume all other resources get cleaned up with the context render_surface_and_context_deinit(**ctx) def render(width, height, fbo, shader): gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, fbo) gl.glClearColor(0.0, 0.0, 0.0, 1.0) gl.glClear(gl.GL_COLOR_BUFFER_BIT) gl.glUseProgram(shader) gl.glDrawArrays(gl.GL_TRIANGLES, 0, 3) data = gl.glReadPixels(0, 0, width, height, gl.GL_RGB, gl.GL_UNSIGNED_BYTE) image = np.frombuffer(data, dtype=np.uint8).reshape(height, width, 3) image = image[::-1, :, :] image = np.array(image).astype(np.float32) / 255.0 return image def shadertoy_vars_update(shader, width, height, time, time_delta, frame_rate, frame): gl.glUseProgram(shader) iResolution_location = gl.glGetUniformLocation(shader, "iResolution") gl.glUniform3f(iResolution_location, width, height, 0) iMouse_location = gl.glGetUniformLocation(shader, "iMouse") gl.glUniform4f(iMouse_location, 0, 0, 0, 0) iTime_location = gl.glGetUniformLocation(shader, "iTime") gl.glUniform1f(iTime_location, time) iTimeDelta_location = gl.glGetUniformLocation(shader, "iTimeDelta") gl.glUniform1f(iTimeDelta_location, time_delta) iFrameRate_location = gl.glGetUniformLocation(shader, "iFrameRate") gl.glUniform1f(iFrameRate_location, frame_rate) iFrame_location = gl.glGetUniformLocation(shader, "iFrame") gl.glUniform1i(iFrame_location, frame) def shadertoy_texture_update(texture, image, frame): if len(image.shape) == 4: image = image[frame] image = image.cpu().numpy() image = image[::-1, :, :] gl.glBindTexture(gl.GL_TEXTURE_2D, texture) gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, image.shape[1], image.shape[0], 0, gl.GL_RGB, gl.GL_FLOAT, image) gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR) gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR) gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE) gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE) def shadertoy_texture_bind(shader, textures): gl.glUseProgram(shader) for i in range(4): gl.glActiveTexture(gl.GL_TEXTURE0 + i) # type: ignore gl.glBindTexture(gl.GL_TEXTURE_2D, textures[i]) iChannel_location = gl.glGetUniformLocation(shader, f"iChannel{i}") gl.glUniform1i(iChannel_location, i) class FL_Shadertoy: @classmethod def INPUT_TYPES(s): return {"required": {"width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}), "height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}), "frame_count": ("INT", {"default": 1, "min": 1, "max": 262144}), "fps": ("INT", {"default": 1, "min": 1, "max": 120}), "source": ( "STRING", {"default": SHADERTOY_DEFAULT, "multiline": True, "dynamicPrompts": False})}, "optional": {"channel_0": ("IMAGE",), "channel_1": ("IMAGE",), "channel_2": ("IMAGE",), "channel_3": ("IMAGE",)}} RETURN_TYPES = ("IMAGE",) CATEGORY = "🏵️Fill Nodes" FUNCTION = "render" def render(self, width: int, height: int, frame_count: int, fps: int, source: str, channel_0: torch.Tensor | None = None, channel_1: torch.Tensor | None = None, channel_2: torch.Tensor | None = None, channel_3: torch.Tensor | None = None): fragment_source = SHADERTOY_HEADER fragment_source += source fragment_source += SHADERTOY_FOOTER ctx, fbo, shader, textures = setup_render_resources(width, height, fragment_source) images = [] frame = 0 total_frames = frame_count for _ in range(frame_count): shadertoy_vars_update(shader, width, height, frame * (1.0 / fps), (1.0 / fps), fps, frame) if channel_0 is not None: shadertoy_texture_update(textures[0], channel_0, frame) if channel_1 is not None: shadertoy_texture_update(textures[1], channel_1, frame) if channel_2 is not None: shadertoy_texture_update(textures[2], channel_2, frame) if channel_3 is not None: shadertoy_texture_update(textures[3], channel_3, frame) shadertoy_texture_bind(shader, textures) image = render(width, height, fbo, shader) image = torch.from_numpy(image)[None,] images.append(image) frame += 1 # Print progress update progress = frame / total_frames * 100 sys.stdout.write(f"\rProcessing frames: {progress:.2f}%") sys.stdout.flush() # Print a new line after the progress update print() render_resources_cleanup(ctx) return (torch.cat(images, dim=0),)