383 lines
15 KiB
Python
383 lines
15 KiB
Python
#!/usr/bin/python3
|
|
# copyright (c) 2018- polygoniq xyz s.r.o.
|
|
|
|
# ##### BEGIN GPL LICENSE BLOCK #####
|
|
#
|
|
# This program is free software; you can redistribute it and/or
|
|
# modify it under the terms of the GNU General Public License
|
|
# as published by the Free Software Foundation; either version 2
|
|
# of the License, or (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with this program; if not, write to the Free Software Foundation,
|
|
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
|
|
#
|
|
# ##### END GPL LICENSE BLOCK #####
|
|
|
|
import bpy
|
|
import typing
|
|
import collections
|
|
import tempfile
|
|
import logging
|
|
from . import polib
|
|
|
|
logger = logging.getLogger(f"polygoniq.{__name__}")
|
|
|
|
|
|
MODULE_CLASSES: typing.List[typing.Type] = []
|
|
|
|
|
|
class DatablockMemoryUsage:
|
|
@staticmethod
|
|
def get_attribute_datatype_size(datatype: str) -> int:
|
|
if datatype == 'FLOAT':
|
|
return 4
|
|
elif datatype == 'INT':
|
|
return 4
|
|
elif datatype == 'FLOAT_VECTOR':
|
|
return 3 * 4
|
|
elif datatype == 'FLOAT_COLOR':
|
|
return 4 * 4 # RGBA
|
|
elif datatype == 'STRING':
|
|
# TODO: we just make something up here...
|
|
return 16
|
|
elif datatype == 'BOOLEAN':
|
|
return 1
|
|
elif datatype == 'FLOAT2':
|
|
return 2 * 4
|
|
else:
|
|
# TODO: warn?
|
|
return 0
|
|
|
|
@staticmethod
|
|
def get_attribute_domain_multiplier(datablock: bpy.types.Mesh, domain: str) -> int:
|
|
if domain == 'POINT':
|
|
return len(datablock.vertices)
|
|
elif domain == 'EDGE':
|
|
return len(datablock.edges)
|
|
elif domain == 'FACE':
|
|
return len(datablock.polygons)
|
|
elif domain == 'CORNER':
|
|
# TODO: no idea
|
|
return 0
|
|
elif domain == 'CURVE':
|
|
# TODO: no idea
|
|
return 0
|
|
elif domain == 'INSTANCE':
|
|
# TODO: no idea
|
|
return 0
|
|
else:
|
|
logger.warning(f"{domain} not supported in memory estimation")
|
|
return 0
|
|
|
|
def __init__(self, datablock: bpy.types.ID):
|
|
self.bytes = 0
|
|
if hasattr(datablock, "name_full"):
|
|
self.title = f"{type(datablock).__name__}:{datablock.name_full}"
|
|
else:
|
|
self.title = f"{type(datablock).__name__}:{datablock.name}"
|
|
|
|
if isinstance(datablock, bpy.types.Mesh):
|
|
# vertex position, 3x float32
|
|
self.bytes += len(datablock.vertices) * 3 * 4
|
|
|
|
if datablock.has_custom_normals:
|
|
# custom normal, 3x float32
|
|
self.bytes += len(datablock.vertices) * 3 * 4
|
|
|
|
# each attribute has to be treated separately
|
|
for attribute in datablock.attributes:
|
|
self.bytes += self.get_attribute_datatype_size(
|
|
attribute.data_type
|
|
) * self.get_attribute_domain_multiplier(datablock, attribute.domain)
|
|
|
|
# each edge indexes into two vertices, each index is int32, plus crease and bevel floats
|
|
self.bytes += len(datablock.edges) * 4 * 4
|
|
# edge index, vertex index
|
|
# TODO: normals? does blender store tangent, bitangent or is this calculated on the fly?
|
|
self.bytes += len(datablock.loops) * 2 * 4
|
|
for polygon in datablock.polygons:
|
|
# loop_start, loop_total indices
|
|
self.bytes += 2 * 4
|
|
# vertex indices
|
|
self.bytes += len(polygon.vertices) * 4
|
|
|
|
for uv_layer in datablock.uv_layers:
|
|
# we assume 2x float32
|
|
self.bytes += len(datablock.vertices) * 2 * 4
|
|
|
|
for vertex_color in datablock.vertex_colors:
|
|
# we assume 3x int8
|
|
self.bytes += len(datablock.vertices) * 3 * 1
|
|
|
|
# TODO: This model is super simplified
|
|
|
|
elif isinstance(datablock, bpy.types.Image):
|
|
width, height = datablock.size
|
|
bits_per_pixel = datablock.depth
|
|
self.bytes += (width * height * bits_per_pixel) // 8
|
|
|
|
assert self.bytes >= 0, f"The {datablock.name} datablock cannot use negative memory"
|
|
|
|
def __repr__(self) -> str:
|
|
return f"{self.bytes} B"
|
|
|
|
|
|
class MemoryUsageStatistics:
|
|
def __init__(self):
|
|
self.datablocks_memory_usage: typing.Dict[str, DatablockMemoryUsage] = {}
|
|
self.target_to_dependencies: typing.DefaultDict[str, typing.Set[str]] = (
|
|
collections.defaultdict(set)
|
|
)
|
|
self.dependency_to_targets: typing.DefaultDict[str, typing.Set[str]] = (
|
|
collections.defaultdict(set)
|
|
)
|
|
|
|
def debug_print(self):
|
|
print(self.datablocks_memory_usage)
|
|
print(self.target_to_dependencies)
|
|
print(self.dependency_to_targets)
|
|
|
|
def write_html_report(self, fp: typing.IO) -> None:
|
|
assert (
|
|
len(self.datablocks_memory_usage) > 0
|
|
), "Memory usage cannot be calculated with no datablocks"
|
|
|
|
sorted_by_usage = sorted(
|
|
self.datablocks_memory_usage.items(), key=lambda x: x[1].bytes, reverse=True
|
|
)
|
|
_, max_usage = sorted_by_usage[0]
|
|
total_usage_bytes = sum(usage.bytes for _, usage in sorted_by_usage)
|
|
max_usage_bytes = max_usage.bytes
|
|
# We want to avoid zero division when calculating percentages
|
|
if total_usage_bytes == 0:
|
|
total_usage_bytes = 1
|
|
if max_usage_bytes == 0:
|
|
max_usage_bytes = 1
|
|
|
|
# TODO: This should be improved to look better
|
|
# TODO: Missing dependencies visualization, why is a particular mesh even loaded, ...
|
|
print("<html><body>\n", file=fp)
|
|
print(f"<b>total usage:</b> {total_usage_bytes / 1024 / 1024 :.3f} MiB<br>\n", file=fp)
|
|
print("<table>\n", file=fp)
|
|
print("<tr>\n", file=fp)
|
|
print("\t\t<th>datablock</th>\n", file=fp)
|
|
print("\t\t<th>approximate usage</th>\n", file=fp)
|
|
print("\t\t<th>%</th>", file=fp)
|
|
print("\t\t<th></th>\n", file=fp)
|
|
print("\t</tr>\n", file=fp)
|
|
|
|
for _, datablock_usage in sorted_by_usage:
|
|
print("\t<tr>\n", file=fp)
|
|
print(f"\t\t<td>{datablock_usage.title}</td>\n", file=fp)
|
|
print(f"\t\t<td>{datablock_usage.bytes / 1024.0 / 1024.0 :.3f} MiB</td>\n", file=fp)
|
|
print(
|
|
f"\t\t<td>{datablock_usage.bytes * 100 / total_usage_bytes :.2f}%</td>\n", file=fp
|
|
)
|
|
red_channel = datablock_usage.bytes * 255 // max_usage_bytes
|
|
print(
|
|
f"\t\t<td><div style=\"background: rgb({red_channel}, {255 - red_channel}, 0); height: 12px; width: {datablock_usage.bytes * 200 // max_usage_bytes}px\"></div></td>\n",
|
|
file=fp,
|
|
)
|
|
print("\t</tr>\n", file=fp)
|
|
print("</table>\n", file=fp)
|
|
print("</body></html>\n", file=fp)
|
|
|
|
def _datablock_key(self, datablock: bpy.types.ID, scope: str = "") -> str:
|
|
# NodeTrees are owned by materials and between materials NodeTree names can clash, so we
|
|
# assign the scope as part of the datablock_key to make it unique.
|
|
if hasattr(datablock, "name_full"):
|
|
return f"{scope}:{type(datablock).__name__}:{datablock.name_full}"
|
|
else:
|
|
return f"{scope}:{type(datablock).__name__}:{datablock.name}"
|
|
|
|
def record_dependency(self, target: bpy.types.ID, dependency: bpy.types.ID) -> None:
|
|
target_key = self._datablock_key(target)
|
|
dependency_key = self._datablock_key(dependency)
|
|
|
|
self.target_to_dependencies[target_key].add(dependency_key)
|
|
self.dependency_to_targets[dependency_key].add(target_key)
|
|
|
|
def record_datablock(self, datablock: bpy.types.ID, scope: str = "") -> None:
|
|
datablock_key = self._datablock_key(datablock, scope)
|
|
if datablock_key in self.datablocks_memory_usage:
|
|
return # already recorded
|
|
|
|
self.datablocks_memory_usage[datablock_key] = DatablockMemoryUsage(datablock)
|
|
|
|
if isinstance(datablock, bpy.types.Scene):
|
|
self._record_scene_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.World):
|
|
self._record_world_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.Collection):
|
|
self._record_collection_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.Object):
|
|
self._record_object_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.Camera):
|
|
pass
|
|
elif isinstance(datablock, bpy.types.Light):
|
|
pass
|
|
elif isinstance(datablock, bpy.types.Mesh):
|
|
pass # meshes don't depend on anything
|
|
elif isinstance(datablock, bpy.types.Curve):
|
|
pass # meshes don't depend on anything
|
|
elif isinstance(datablock, bpy.types.Material):
|
|
self._record_material_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.NodeTree): # TODO: or only ShaderNodeTree?
|
|
self._record_node_tree_dependencies(datablock, scope)
|
|
elif isinstance(datablock, bpy.types.ShaderNodeTexImage):
|
|
self._record_shader_node_tex_image_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.ShaderNodeTexEnvironment):
|
|
self._record_shader_node_tex_environment_dependencies(datablock)
|
|
elif isinstance(datablock, bpy.types.Image):
|
|
pass # images don't depend on anything
|
|
elif isinstance(datablock, bpy.types.Armature):
|
|
pass # TODO: Support armatures, armatures are ligthweight
|
|
elif isinstance(datablock, bpy.types.ShaderNodeGroup):
|
|
self._record_shader_node_group_dependencies(datablock)
|
|
else:
|
|
logger.warning(f"{datablock} not taken into account in estimation!")
|
|
|
|
def _record_scene_dependencies(self, scene: bpy.types.Scene) -> None:
|
|
if scene.world is not None:
|
|
self.record_datablock(scene.world)
|
|
self.record_dependency(scene, scene.world)
|
|
|
|
self.record_datablock(scene.collection)
|
|
self.record_dependency(scene, scene.collection)
|
|
|
|
def _record_world_dependencies(self, world: bpy.types.World) -> None:
|
|
if world.use_nodes:
|
|
self.record_datablock(world.node_tree, self._datablock_key(world))
|
|
self.record_dependency(world, world.node_tree)
|
|
|
|
def _record_collection_dependencies(self, collection: bpy.types.Collection) -> None:
|
|
for obj in collection.objects:
|
|
if obj.hide_render:
|
|
continue # skip hidden objects
|
|
|
|
self.record_datablock(obj)
|
|
self.record_dependency(collection, obj)
|
|
|
|
for child_collection in collection.children:
|
|
if child_collection.hide_render:
|
|
continue # skip hidden collections
|
|
|
|
self.record_datablock(child_collection)
|
|
self.record_dependency(collection, child_collection)
|
|
|
|
def _record_object_dependencies(self, obj: bpy.types.Object) -> None:
|
|
if obj.data is not None:
|
|
self.record_datablock(obj.data)
|
|
self.record_dependency(obj, obj.data)
|
|
if obj.instance_type == 'COLLECTION':
|
|
if obj.instance_collection is not None:
|
|
self.record_datablock(obj.instance_collection)
|
|
self.record_dependency(obj, obj.instance_collection)
|
|
else:
|
|
logger.warning(
|
|
f"Object {obj.name} has instance_type='COLLECTION' but its "
|
|
f"instance_collection is None! Skipping..."
|
|
)
|
|
|
|
for material_slot in obj.material_slots:
|
|
if material_slot.material is not None:
|
|
self.record_datablock(material_slot.material)
|
|
self.record_dependency(obj, material_slot.material)
|
|
|
|
for child_object in obj.children: # TODO: this takes O(len(bpy.data.objects))!
|
|
self.record_datablock(child_object)
|
|
self.record_dependency(obj, child_object)
|
|
self._record_object_dependencies(child_object)
|
|
|
|
def _record_material_dependencies(self, material: bpy.types.Object) -> None:
|
|
if not material.use_nodes:
|
|
logger.warning(
|
|
f"Warning: Cannot record dependencies of {material.name}, only node based "
|
|
f"materials are supported for now."
|
|
)
|
|
return
|
|
|
|
if material.node_tree is not None:
|
|
self.record_datablock(material.node_tree, self._datablock_key(material))
|
|
self.record_dependency(material, material.node_tree)
|
|
|
|
def _record_node_tree_dependencies(
|
|
self, node_tree: bpy.types.NodeTree, scope: str = ""
|
|
) -> None:
|
|
for node in node_tree.nodes:
|
|
if isinstance(node, bpy.types.ShaderNodeTexImage):
|
|
self.record_datablock(node, scope)
|
|
self.record_dependency(node_tree, node)
|
|
self._record_shader_node_tex_image_dependencies(node)
|
|
elif isinstance(node, bpy.types.ShaderNodeTexEnvironment):
|
|
self.record_datablock(node, scope)
|
|
self.record_dependency(node_tree, node)
|
|
self._record_shader_node_tex_environment_dependencies(node)
|
|
elif isinstance(node, bpy.types.ShaderNodeGroup):
|
|
self.record_datablock(node, scope)
|
|
self.record_dependency(node_tree, node)
|
|
else:
|
|
# TODO: Tons of nodes we don't support yet
|
|
pass
|
|
|
|
def _record_shader_node_tex_image_dependencies(
|
|
self, node: bpy.types.ShaderNodeTexImage
|
|
) -> None:
|
|
if node.image is not None:
|
|
self.record_datablock(node.image)
|
|
self.record_dependency(node, node.image)
|
|
|
|
def _record_shader_node_tex_environment_dependencies(
|
|
self, node: bpy.types.ShaderNodeTexEnvironment
|
|
) -> None:
|
|
if node.image is not None:
|
|
self.record_datablock(node.image)
|
|
self.record_dependency(node, node.image)
|
|
|
|
def _record_shader_node_group_dependencies(self, node: bpy.types.ShaderNodeGroup) -> None:
|
|
if node.node_tree is not None:
|
|
self.record_datablock(node.node_tree, scope="bpy.data.node_groups")
|
|
self.record_dependency(node, node.node_tree)
|
|
|
|
|
|
@polib.log_helpers_bpy.logged_operator
|
|
class EstimateMemoryUsage(bpy.types.Operator):
|
|
bl_idname = "memsaver.estimate_memory_usage"
|
|
bl_label = "Estimate Memory Usage (Beta)"
|
|
bl_description = (
|
|
"Goes through datablocks that would have to be loaded for rendering, "
|
|
"estimates how much memory is needed for each one"
|
|
)
|
|
bl_options = {'REGISTER'}
|
|
|
|
def execute(self, context: bpy.types.Context):
|
|
stats = MemoryUsageStatistics()
|
|
stats.record_datablock(context.scene)
|
|
|
|
out_file = tempfile.NamedTemporaryFile(
|
|
mode="w", prefix="memory_usage_", suffix=".html", delete=False
|
|
)
|
|
stats.write_html_report(out_file)
|
|
logger.info(f"Wrote memory estimation HTML to: {out_file.name}")
|
|
polib.utils_bpy.xdg_open_file(out_file.name)
|
|
return {'FINISHED'}
|
|
|
|
|
|
MODULE_CLASSES.append(EstimateMemoryUsage)
|
|
|
|
|
|
def register():
|
|
for cls in MODULE_CLASSES:
|
|
bpy.utils.register_class(cls)
|
|
|
|
|
|
def unregister():
|
|
for cls in reversed(MODULE_CLASSES):
|
|
bpy.utils.unregister_class(cls)
|