226 lines
7.9 KiB
Python
226 lines
7.9 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import sys
|
|
import OpenGL.GL as gl
|
|
import glfw
|
|
import ctypes
|
|
|
|
VERTEX_SHADER = """
|
|
#version 330 core
|
|
|
|
layout (location = 0) in vec3 aPos;
|
|
layout (location = 1) in vec2 aTexCoord;
|
|
|
|
out vec2 TexCoord;
|
|
|
|
void main()
|
|
{
|
|
gl_Position = vec4(aPos, 1.0);
|
|
TexCoord = aTexCoord;
|
|
}
|
|
"""
|
|
|
|
FRAGMENT_SHADER = """
|
|
#version 330 core
|
|
|
|
out vec4 FragColor;
|
|
|
|
in vec2 TexCoord;
|
|
|
|
uniform sampler2D iChannel0;
|
|
uniform vec3 iResolution;
|
|
uniform float iTime;
|
|
uniform float iScale;
|
|
uniform float iSwirl;
|
|
uniform float iSwirlStrength;
|
|
uniform float iBlur;
|
|
uniform float iIterations;
|
|
uniform float iTimeSpeed;
|
|
|
|
void main()
|
|
{
|
|
vec4 O = vec4(0.0);
|
|
vec2 U = TexCoord;
|
|
|
|
float s = 0.0, s2 = 0.0, t = iTime * iTimeSpeed;
|
|
U = U - 0.5;
|
|
U.x += 0.03 * sin(1.14 * t);
|
|
float sc = pow(iScale, -mod(t, 2.0) - 0.8);
|
|
U *= sc;
|
|
|
|
for (int i = 0; i < int(iIterations); i++) {
|
|
vec2 V = abs(U + U);
|
|
if (max(V.x, V.y) > 1.0) break;
|
|
V = smoothstep(1.0, 0.5, V);
|
|
float m = V.x * V.y;
|
|
O = mix(O, texture(iChannel0, U + 0.5), m);
|
|
s = mix(s, 1.0, m);
|
|
s2 = s2 * (1.0 - m) * (1.0 - m) + m * m;
|
|
U *= iScale;
|
|
if (iSwirl > 0.5) {
|
|
U.x = -U.x * iSwirlStrength;
|
|
}
|
|
}
|
|
|
|
vec4 mean = texture(iChannel0, U, 10.0);
|
|
O = mean + (O - s * mean) / sqrt(s2);
|
|
|
|
FragColor = O;
|
|
}
|
|
"""
|
|
|
|
|
|
class FL_InfiniteZoom:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"scale": ("FLOAT", {"default": 2.00, "min": 1.10, "max": 10.00, "step": 0.05}),
|
|
"mirror": (["on", "off"],),
|
|
"mirror_warp": ("FLOAT", {"default": 1.00, "min": 0.50, "max": 1.50, "step": 0.05}),
|
|
"iterations": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1}),
|
|
#"blur": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"speed": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
|
|
"fps": ("INT", {"default": 30, "min": 1, "max": 120, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply_shader"
|
|
CATEGORY = "🏵️Fill Nodes"
|
|
|
|
def apply_shader(self, images, scale, mirror, mirror_warp, iterations, speed, fps):
|
|
result = []
|
|
total_images = len(images)
|
|
frame_time = 1.0 / fps
|
|
|
|
for i, image in enumerate(images, start=1):
|
|
img = self.t2p(image)
|
|
result_img = self.process_image(img, scale, mirror, mirror_warp, iterations, speed,
|
|
i * frame_time)
|
|
result_img = self.p2t(result_img)
|
|
result.append(result_img)
|
|
|
|
# Update the print log
|
|
progress = i / total_images * 100
|
|
sys.stdout.write(f"\rProcessing images: {progress:.2f}%")
|
|
sys.stdout.flush()
|
|
|
|
# Print a new line after the progress log
|
|
print()
|
|
|
|
return (torch.cat(result, dim=0),)
|
|
|
|
def process_image(self, image, scale, mirror, mirror_warp, iterations, speed, time):
|
|
# Convert the PIL image to a numpy array
|
|
img_array = np.array(image).astype(np.float32) / 255.0
|
|
|
|
# Create a PyOpenGL context
|
|
if not glfw.init():
|
|
raise RuntimeError("Failed to initialize GLFW")
|
|
|
|
glfw.window_hint(glfw.VISIBLE, glfw.FALSE)
|
|
window = glfw.create_window(image.width, image.height, "Hidden Window", None, None)
|
|
if not window:
|
|
glfw.terminate()
|
|
raise RuntimeError("Failed to create GLFW window")
|
|
|
|
glfw.make_context_current(window)
|
|
|
|
# Compile the shader program
|
|
vertex_shader = gl.glCreateShader(gl.GL_VERTEX_SHADER)
|
|
gl.glShaderSource(vertex_shader, VERTEX_SHADER)
|
|
gl.glCompileShader(vertex_shader)
|
|
|
|
fragment_shader = gl.glCreateShader(gl.GL_FRAGMENT_SHADER)
|
|
gl.glShaderSource(fragment_shader, FRAGMENT_SHADER)
|
|
gl.glCompileShader(fragment_shader)
|
|
|
|
shader_program = gl.glCreateProgram()
|
|
gl.glAttachShader(shader_program, vertex_shader)
|
|
gl.glAttachShader(shader_program, fragment_shader)
|
|
gl.glLinkProgram(shader_program)
|
|
|
|
gl.glUseProgram(shader_program)
|
|
|
|
# Set up vertex buffer object (VBO) and vertex array object (VAO)
|
|
vertices = np.array([
|
|
-1.0, -1.0, 0.0, 0.0, 0.0,
|
|
1.0, -1.0, 0.0, 1.0, 0.0,
|
|
-1.0, 1.0, 0.0, 0.0, 1.0,
|
|
1.0, 1.0, 0.0, 1.0, 1.0
|
|
], dtype=np.float32)
|
|
|
|
vao = gl.glGenVertexArrays(1)
|
|
gl.glBindVertexArray(vao)
|
|
|
|
vbo = gl.glGenBuffers(1)
|
|
gl.glBindBuffer(gl.GL_ARRAY_BUFFER, vbo)
|
|
gl.glBufferData(gl.GL_ARRAY_BUFFER, vertices.nbytes, vertices, gl.GL_STATIC_DRAW)
|
|
|
|
gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, 5 * vertices.itemsize, None)
|
|
gl.glEnableVertexAttribArray(0)
|
|
|
|
gl.glVertexAttribPointer(1, 2, gl.GL_FLOAT, gl.GL_FALSE, 5 * vertices.itemsize,
|
|
ctypes.c_void_p(3 * vertices.itemsize))
|
|
gl.glEnableVertexAttribArray(1)
|
|
|
|
# Set up texture
|
|
texture = gl.glGenTextures(1)
|
|
gl.glBindTexture(gl.GL_TEXTURE_2D, texture)
|
|
gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_REPEAT)
|
|
gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_REPEAT)
|
|
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.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, image.width, image.height, 0, gl.GL_RGB, gl.GL_FLOAT, img_array)
|
|
|
|
# Set shader uniforms
|
|
gl.glUniform1i(gl.glGetUniformLocation(shader_program, "iChannel0"), 0)
|
|
gl.glUniform3f(gl.glGetUniformLocation(shader_program, "iResolution"), image.width, image.height, 0.0)
|
|
gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iTime"), time)
|
|
gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iScale"), scale)
|
|
gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iSwirl"), 1.0 if mirror == "on" else 0.0)
|
|
gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iSwirlStrength"), mirror_warp)
|
|
# gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iBlur"), blur)
|
|
gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iIterations"), iterations)
|
|
gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iTimeSpeed"), speed)
|
|
|
|
# Render the shader
|
|
gl.glClear(gl.GL_COLOR_BUFFER_BIT)
|
|
gl.glDrawArrays(gl.GL_TRIANGLE_STRIP, 0, 4)
|
|
|
|
# Read the rendered image from the framebuffer
|
|
img_data = gl.glReadPixels(0, 0, image.width, image.height, gl.GL_RGB, gl.GL_FLOAT)
|
|
img_array = np.frombuffer(img_data, dtype=np.float32).reshape((image.height, image.width, 3))
|
|
|
|
# Clean up OpenGL resources
|
|
gl.glDeleteTextures(1, [texture])
|
|
gl.glDeleteBuffers(1, [vbo])
|
|
gl.glDeleteVertexArrays(1, [vao])
|
|
gl.glDeleteProgram(shader_program)
|
|
gl.glDeleteShader(vertex_shader)
|
|
gl.glDeleteShader(fragment_shader)
|
|
|
|
glfw.destroy_window(window)
|
|
glfw.terminate()
|
|
|
|
# Convert the processed image back to a PIL image
|
|
processed_image = Image.fromarray((img_array * 255).astype(np.uint8))
|
|
|
|
return processed_image
|
|
|
|
def t2p(self, t):
|
|
if t is not None:
|
|
i = 255.0 * t.cpu().numpy().squeeze()
|
|
p = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
return p
|
|
|
|
def p2t(self, p):
|
|
if p is not None:
|
|
i = np.array(p).astype(np.float32) / 255.0
|
|
t = torch.from_numpy(i).unsqueeze(0)
|
|
return t |