diff --git a/manimlib/camera/camera.py b/manimlib/camera/camera.py index abd0bf9a..1bd3c4be 100644 --- a/manimlib/camera/camera.py +++ b/manimlib/camera/camera.py @@ -15,6 +15,7 @@ from manimlib.constants import FRAME_WIDTH from manimlib.mobject.mobject import Mobject from manimlib.mobject.mobject import Point from manimlib.utils.color import color_to_rgba +from manimlib.utils.shaders import get_texture_id from typing import TYPE_CHECKING @@ -67,7 +68,6 @@ class Camera(object): self.perspective_uniforms = dict() self.init_frame(**frame_config) self.init_context(window) - self.init_textures() self.init_light_source() self.refresh_perspective_uniforms() self.init_fill_fbo(self.ctx) # Experimental @@ -136,10 +136,8 @@ class Camera(object): ''', ) - tid = self.n_textures - self.fill_texture.use(tid) - self.fill_prog['Texture'].value = tid - self.n_textures += 1 + self.fill_prog['Texture'].value = get_texture_id(self.fill_texture) + verts = np.array([[0, 0], [0, 1], [1, 0], [1, 1]]) self.fill_texture_vao = ctx.simple_vertex_array( self.fill_prog, @@ -289,9 +287,8 @@ class Camera(object): def render(self, render_group: dict[str, Any]) -> None: shader_wrapper = render_group["shader_wrapper"] - shader_program = render_group["prog"] primitive = int(shader_wrapper.render_primitive) - self.set_shader_uniforms(shader_program, shader_wrapper) + shader_wrapper.update_program_uniforms(self.perspective_uniforms) self.set_ctx_depth_test(shader_wrapper.depth_test) self.set_ctx_clip_plane(shader_wrapper.use_clip_plane()) @@ -344,69 +341,20 @@ class Camera(object): shader_wrapper: ShaderWrapper, single_use: bool = True ) -> dict[str, Any]: - # Data buffer - vert_data = shader_wrapper.vert_data - indices = shader_wrapper.vert_indices - if len(indices) == 0: - ibo = None - elif single_use: - ibo = self.ctx.buffer(indices.astype(np.uint32)) - else: - ibo = self.ctx.buffer(indices.astype(np.uint32)) - # # The vao.render call is strangely longer - # # when an index buffer is used, so if the - # # mobject is not changing, meaning only its - # # uniforms are being updated, just create - # # a larger data array based on the indices - # # and don't bother with the ibo - # vert_data = vert_data[indices] - # ibo = None - vbo = self.ctx.buffer(vert_data) - - # Program and vertex array - shader_program = shader_wrapper.program - vert_format = shader_wrapper.vert_format - attributes = shader_wrapper.vert_attributes - vao = self.ctx.vertex_array( - program=shader_program, - content=[(vbo, vert_format, *attributes)], - index_buffer=ibo, - ) return { - "vbo": vbo, - "ibo": ibo, - "vao": vao, - "prog": shader_program, "shader_wrapper": shader_wrapper, + "vao": shader_wrapper.get_vao(single_use), "single_use": single_use, } def release_render_group(self, render_group: dict[str, Any]) -> None: - for key in ["vbo", "ibo", "vao"]: - if render_group[key] is not None: - render_group[key].release() + render_group["shader_wrapper"].release() def refresh_static_mobjects(self) -> None: for render_group in it.chain(*self.mob_to_render_groups.values()): self.release_render_group(render_group) self.mob_to_render_groups = {} - # Shaders - - def set_shader_uniforms( - self, - shader: moderngl.Program, - shader_wrapper: ShaderWrapper - ) -> None: - for name, path in shader_wrapper.texture_paths.items(): - tid = self.get_texture_id(path) - shader[name].value = tid - for name, value in it.chain(self.perspective_uniforms.items(), shader_wrapper.uniforms.items()): - if name in shader: - if isinstance(value, np.ndarray) and value.ndim > 0: - value = tuple(value) - shader[name].value = value - def refresh_perspective_uniforms(self) -> None: frame = self.frame view_matrix = frame.get_view_matrix() @@ -422,32 +370,6 @@ class Camera(object): focal_distance=frame.get_focal_distance(), ) - def init_textures(self) -> None: - self.n_textures: int = 0 - self.path_to_texture: dict[ - str, tuple[int, moderngl.Texture] - ] = {} - - def get_texture_id(self, path: str) -> int: - if path not in self.path_to_texture: - tid = self.n_textures - self.n_textures += 1 - im = Image.open(path).convert("RGBA") - texture = self.ctx.texture( - size=im.size, - components=len(im.getbands()), - data=im.tobytes(), - ) - texture.use(location=tid) - self.path_to_texture[path] = (tid, texture) - return self.path_to_texture[path][0] - - def release_texture(self, path: str): - tid_and_texture = self.path_to_texture.pop(path, None) - if tid_and_texture: - tid_and_texture[1].release() - return self - # Mostly just defined so old scenes don't break class ThreeDCamera(Camera): diff --git a/manimlib/mobject/mobject.py b/manimlib/mobject/mobject.py index 14b23d3b..f5de726d 100644 --- a/manimlib/mobject/mobject.py +++ b/manimlib/mobject/mobject.py @@ -1861,7 +1861,7 @@ class Mobject(object): self.shader_wrapper.vert_data = self.get_shader_data() self.shader_wrapper.vert_indices = self.get_shader_vert_indices() - self.shader_wrapper.uniforms = self.get_uniforms() + self.shader_wrapper.uniforms.update(self.get_uniforms()) self.shader_wrapper.depth_test = self.depth_test return self.shader_wrapper diff --git a/manimlib/mobject/types/vectorized_mobject.py b/manimlib/mobject/types/vectorized_mobject.py index 4fa3df1b..d9cb4d2c 100644 --- a/manimlib/mobject/types/vectorized_mobject.py +++ b/manimlib/mobject/types/vectorized_mobject.py @@ -1250,7 +1250,7 @@ class VMobject(Mobject): for sw in shader_wrappers: # Assume uniforms of the first family member - sw.uniforms = family[0].get_uniforms() + sw.uniforms.update(family[0].get_uniforms()) sw.depth_test = family[0].depth_test return [sw for sw in shader_wrappers if len(sw.vert_data) > 0] diff --git a/manimlib/shader_wrapper.py b/manimlib/shader_wrapper.py index 16ca214a..8f02a977 100644 --- a/manimlib/shader_wrapper.py +++ b/manimlib/shader_wrapper.py @@ -10,11 +10,14 @@ import numpy as np from manimlib.utils.iterables import resize_array from manimlib.utils.shaders import get_shader_code_from_file from manimlib.utils.shaders import get_shader_program +from manimlib.utils.shaders import image_path_to_texture +from manimlib.utils.shaders import get_texture_id +from manimlib.utils.shaders import release_texture from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import List + from typing import List, Optional # Mobjects that should be rendered with @@ -29,10 +32,10 @@ class ShaderWrapper(object): self, context: moderngl.context.Context, vert_data: np.ndarray, - vert_indices: np.ndarray | None = None, - shader_folder: str | None = None, - uniforms: dict[str, float | np.ndarray] | None = None, # A dictionary mapping names of uniform variables - texture_paths: dict[str, str] | None = None, # A dictionary mapping names to filepaths for textures. + vert_indices: Optional[np.ndarray] = None, + shader_folder: Optional[str] = None, + uniforms: Optional[dict[str, float | np.ndarray]] = None, # A dictionary mapping names of uniform variables + texture_paths: Optional[dict[str, str]] = None, # A dictionary mapping names to filepaths for textures. depth_test: bool = False, render_primitive: int = moderngl.TRIANGLE_STRIP, is_fill: bool = False, @@ -43,12 +46,18 @@ class ShaderWrapper(object): self.vert_attributes = vert_data.dtype.names self.shader_folder = shader_folder self.uniforms = uniforms or dict() - self.texture_paths = texture_paths or dict() self.depth_test = depth_test self.render_primitive = str(render_primitive) self.is_fill = is_fill + + self.vbo = None + self.ibo = None + self.vao = None + self.init_program_code() self.init_program() + if texture_paths is not None: + self.init_textures(texture_paths) self.refresh_id() def init_program_code(self) -> None: @@ -71,6 +80,12 @@ class ShaderWrapper(object): self.program = get_shader_program(self.ctx, **self.program_code) self.vert_format = moderngl.detect_format(self.program, self.vert_attributes) + def init_textures(self, texture_paths: dict[str, str]): + for name, path in texture_paths.items(): + texture = image_path_to_texture(path, self.ctx) + tid = get_texture_id(texture) + self.uniforms[name] = tid + def __eq__(self, shader_wrapper: ShaderWrapper): return all(( np.all(self.vert_data == shader_wrapper.vert_data), @@ -80,10 +95,6 @@ class ShaderWrapper(object): np.all(self.uniforms[key] == shader_wrapper.uniforms[key]) for key in self.uniforms ), - all( - self.texture_paths[key] == shader_wrapper.texture_paths[key] - for key in self.texture_paths - ), self.depth_test == shader_wrapper.depth_test, self.render_primitive == shader_wrapper.render_primitive, )) @@ -94,8 +105,6 @@ class ShaderWrapper(object): result.vert_indices = self.vert_indices.copy() if self.uniforms: result.uniforms = {key: np.array(value) for key, value in self.uniforms.items()} - if self.texture_paths: - result.texture_paths = dict(self.texture_paths) return result def is_valid(self) -> bool: @@ -116,7 +125,6 @@ class ShaderWrapper(object): return "|".join(map(str, [ self.program_id, self.uniforms, - self.texture_paths, self.depth_test, self.render_primitive, ])) @@ -186,3 +194,47 @@ class ShaderWrapper(object): n_verts = new_n_verts n_points += len(data) return self + + def update_program_uniforms(self, camera_uniforms: dict): + if self.program is None: + return + for name, value in (*camera_uniforms.items(), *self.uniforms.items()): + if name in self.program: + if isinstance(value, np.ndarray) and value.ndim > 0: + value = tuple(value) + self.program[name].value = value + + def get_vao(self, single_use: bool = False): + # Data buffer + vert_data = self.vert_data + indices = self.vert_indices + if len(indices) == 0: + self.ibo = None + elif single_use or self.is_fill: + self.ibo = self.ctx.buffer(indices.astype(np.uint32)) + else: + # The vao.render call is strangely longer + # when an index buffer is used, so if the + # mobject is not changing, meaning only its + # uniforms are being updated, just create + # a larger data array based on the indices + # and don't bother with the ibo + vert_data = vert_data[indices] + self.ibo = None + self.vbo = self.ctx.buffer(vert_data) + + # Vertex array object + self.vao = self.ctx.vertex_array( + program=self.program, + content=[(self.vbo, self.vert_format, *self.vert_attributes)], + index_buffer=self.ibo, + ) + return self.vao + + def release(self): + for obj in (self.vbo, self.ibo, self.vao): + if obj is not None: + obj.release() + self.vbo = None + self.ibo = None + self.vao = None diff --git a/manimlib/utils/shaders.py b/manimlib/utils/shaders.py index 2b4647f0..eaddcaea 100644 --- a/manimlib/utils/shaders.py +++ b/manimlib/utils/shaders.py @@ -4,6 +4,7 @@ import os import re from functools import lru_cache import moderngl +from PIL import Image from manimlib.utils.directories import get_shader_dir from manimlib.utils.file_ops import find_file @@ -14,6 +15,34 @@ if TYPE_CHECKING: from typing import Sequence, Optional +ID_TO_TEXTURE: dict[int, moderngl.Texture] = dict() + + +@lru_cache() +def image_path_to_texture(path: str, ctx: moderngl.Context) -> moderngl.Texture: + im = Image.open(path).convert("RGBA") + return ctx.texture( + size=im.size, + components=len(im.getbands()), + data=im.tobytes(), + ) + + +def get_texture_id(texture: moderngl.Texture) -> int: + tid = 0 + while tid in ID_TO_TEXTURE: + tid += 1 + ID_TO_TEXTURE[tid] = texture + texture.use(location=tid) + return tid + + +def release_texture(texture_id: int): + texture = ID_TO_TEXTURE.pop(texture_id, None) + if texture is not None: + texture.release() + + @lru_cache() def get_shader_program( ctx: moderngl.context.Context,