bake_cubemap.py

import bpy
import os
import sys

# Add the current directory to Python path so we can import config
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

from config import CUBEMAP_CONFIGS, get_dynamic_mip_levels, get_dynamic_cubemap_configs

def setup_output_directory():
    """Create and return the output directory path."""
    output_dir = os.path.join(os.getcwd(), "output")
    os.makedirs(output_dir, exist_ok=True)
    return output_dir

def get_cube_probe():
    """Get the CubeProbe object or exit if not found."""
    cube_probe = bpy.data.objects.get("CubeProbe")
    if not cube_probe:
        print("Error: CubeProbe object not found!")
        sys.exit(1)
    return cube_probe

def set_environment_texture(texture_path):
    """Set the environment texture for the world."""
    # Make sure we have a world
    if not bpy.data.worlds:
        print("Error: No world found in the scene!")
        return False
    
    world = bpy.data.worlds[0]
    
    # Verify world uses nodes
    if not world.use_nodes:
        print("Error: World does not use nodes!")
        return False
    
    nodes = world.node_tree.nodes
    
    # Find Environment Texture node
    env_tex_node = None
    for node in nodes:
        if node.type == 'TEX_ENVIRONMENT':
            env_tex_node = node
            break
    
    if not env_tex_node:
        print("Error: Environment Texture node not found in the world node tree!")
        return False
    
    # Load and assign the texture
    try:
        # Check if the image is already loaded
        image = None
        for img in bpy.data.images:
            if img.filepath == texture_path:
                image = img
                break
        
        if not image:
            # Load the image
            image = bpy.data.images.load(texture_path)
        
        # Set the texture
        env_tex_node.image = image
        print(f"Successfully set environment texture to: {texture_path}")
        return True
        
    except Exception as e:
        print(f"Error setting environment texture: {e}")
        return False

def set_white_point(white_point_value):
    """Set the value of the white point node in the world shader."""
    # Make sure we have a world
    if not bpy.data.worlds:
        print("Error: No world found in the scene!")
        return False
    
    world = bpy.data.worlds[0]
    
    # Verify world uses nodes
    if not world.use_nodes:
        print("Error: World does not use nodes!")
        return False
    
    nodes = world.node_tree.nodes
    
    # Find white point node
    white_point_node = None
    for node in nodes:
        if node.label == "WhitePoint":
            white_point_node = node
            break
    
    if not white_point_node:
        print("Error: WhitePoint node not found in the world node tree!")
        return False
    
    # Set the value
    try:
        # Check if it's a Value node
        if white_point_node.type == 'VALUE':
            # Set the value
            white_point_node.outputs[0].default_value = white_point_value
            print(f"Successfully set white point value to: {white_point_value}")
            return True
        else:
            print(f"Error: white point node is not a Value node, it's a {white_point_node.type} node")
            return False
        
    except Exception as e:
        print(f"Error setting white point value: {e}")
        return False

def set_tonemap(should_tonemap):
    """Set the tonemapping factor by finding the node with label 'ShouldTonemap'."""
    # Make sure we have a world
    if not bpy.data.worlds:
        print("Error: No world found in the scene!")
        return False
    
    world = bpy.data.worlds[0]
    
    # Verify world uses nodes
    if not world.use_nodes:
        print("Error: World does not use nodes!")
        return False
    
    nodes = world.node_tree.nodes
    
    # Find ShouldTonemap node
    tonemap_node = None
    for node in nodes:
        if node.label == "ShouldTonemap":
            tonemap_node = node
            break
    
    if not tonemap_node:
        print("Error: ShouldTonemap node not found in the world node tree!")
        return False
    
    # Set the factor value
    try:
        factor_value = 1.0 if should_tonemap else 0.0
        
        # Check if it has a Factor input (common for Mix nodes)
        if "Factor" in tonemap_node.inputs:
            tonemap_node.inputs["Factor"].default_value = factor_value
            print(f"Successfully set tonemap factor to: {factor_value}")
            return True
        # Check if it has a Fac input (older Blender versions)
        elif "Fac" in tonemap_node.inputs:
            tonemap_node.inputs["Fac"].default_value = factor_value
            print(f"Successfully set tonemap factor to: {factor_value}")
            return True
        # Check if it's a Value node
        elif tonemap_node.type == 'VALUE':
            tonemap_node.outputs[0].default_value = factor_value
            print(f"Successfully set tonemap value to: {factor_value}")
            return True
        else:
            print(f"Error: ShouldTonemap node does not have a Factor, Fac input, or is not a Value node")
            print(f"Node type: {tonemap_node.type}")
            print(f"Available inputs: {[input.name for input in tonemap_node.inputs]}")
            return False
        
    except Exception as e:
        print(f"Error setting tonemap factor: {e}")
        return False

def create_bake_image(name, output_dir, mip_level=0, base_resolution=512):
    """Create a new image for baking with specified mip level."""
    resolution = max(1, base_resolution // (2 ** mip_level))  # Ensure minimum size is 1x1
    
    # Create new image with standard 1:1 aspect ratio first
    image = bpy.data.images.new(name=name, width=resolution*4, height=resolution*3, float_buffer=True)
    image.filepath = os.path.join(output_dir, f"{name}.hdr")
    image.file_format = 'HDR'
    return image

def setup_render_settings(resolution=512):
    """Configure render and bake settings."""
    # Force CPU renderer to avoid Metal GPU crashes
    # bpy.context.scene.render.engine = 'CYCLES'
    # bpy.context.scene.cycles.device = 'CPU'
    
    # bpy.context.scene.cycles.samples = 16
    # bpy.context.scene.cycles.use_denoising = False
    
    # Set render resolution
    bpy.context.scene.render.resolution_x = resolution
    bpy.context.scene.render.resolution_y = resolution

    bpy.context.scene.render.bake.use_selected_to_active = False
    bpy.context.scene.render.bake.use_cage = False
    bpy.context.scene.render.bake.use_clear = True
    
    # Disable tonemapping during baking
    bpy.context.scene.view_settings.view_transform = 'Raw'
    bpy.context.scene.view_settings.look = 'None'
    bpy.context.scene.display_settings.display_device = 'sRGB'

def select_object(obj):
    """Select the given object and make it active."""
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj

def adjust_material_roughness(material_name, roughness_value):
    """Adjust the roughness of the specified material."""
    material = bpy.data.materials.get(material_name)
    if not material:
        print(f"Error: Material '{material_name}' not found!")
        return False
    
    if not material.use_nodes:
        print(f"Error: Material '{material_name}' does not use nodes!")
        return False
    
    # Find the Principled BSDF node
    principled_node = None
    for node in material.node_tree.nodes:
        if node.type == 'BSDF_PRINCIPLED':
            principled_node = node
            break
    
    if not principled_node:
        print(f"Error: No Principled BSDF node found in material '{material_name}'!")
        return False
    
    # Set the roughness value
    principled_node.inputs['Roughness'].default_value = roughness_value
    return True

def setup_image_texture_node(obj, image):
    """Set up or create an image texture node for baking."""
    if not obj.active_material:
        print(f"Error: Object {obj.name} has no active material")
        return False
    
    if not obj.active_material.use_nodes:
        obj.active_material.use_nodes = True
    
    material = obj.active_material
    nodes = material.node_tree.nodes
    
    # Clear existing Image Texture nodes to avoid conflicts
    for node in list(nodes):
        if node.type == 'TEX_IMAGE' or node.bl_idname == 'ShaderNodeTexImage':
            nodes.remove(node)
    
    # Create a new Image Texture node using proper Blender 4.x naming
    try:
        tex_node = nodes.new(type='ShaderNodeTexImage')
    except Exception as e:
        print(f"Error creating ShaderNodeTexImage: {e}")
        try:
            # Fallback to older naming
            tex_node = nodes.new(type='TEX_IMAGE')
        except Exception as e2:
            print(f"Error creating TEX_IMAGE: {e2}")
            return False
    
    tex_node.image = image
    
    # Make this the active node for baking
    nodes.active = tex_node
    return True

def bake_single_cubemap(cube_probe, image, mip_level, output_dir):
    """Bake a single cubemap with specified roughness and mip level."""
    try:
        base_size = 512
        resolution = max(8, base_size // (2 ** mip_level))
        
        # Setup render settings for this mip level (CPU only)
        setup_render_settings(resolution)
        
        # Select the cube probe
        select_object(cube_probe)
        
        # Set up the image texture node for baking
        if not setup_image_texture_node(cube_probe, image):
            print(f"Error: Failed to set up image texture node")
            return False
        
        # Bake with CPU
        try:
            bpy.ops.object.bake(type='COMBINED', use_selected_to_active=False)
            # save the image to disk
            image.file_format = 'HDR'
            image.save_render(filepath=os.path.join(output_dir, f"{image.name}.hdr"))
        except Exception as e:
            print(f"Baking failed with error: {e}")
            return False
        
        return True
        
    except Exception as e:
        print(f"Error during baking: {e}")
        return False

def bake_and_save_cubemap(cube_probe, image_name, output_dir, mip_level, roughness, resolution=None, base_resolution=512):
    """Bake and save a single cubemap with specified parameters.
    
    Args:
        cube_probe: The cube probe object to bake
        image_name: Name for the output image
        output_dir: Directory to save the image
        mip_level: Mip level for image creation
        roughness: Material roughness value
        resolution: Optional resolution override
        base_resolution: Base resolution for calculations
        
    Returns:
        bool: Success status
    """
    # Create image for this baking operation
    image = create_bake_image(image_name, output_dir, mip_level, base_resolution)
    
    # Adjust material roughness
    adjust_material_roughness("BakeMaterial", roughness)
    
    # Setup render settings
    if resolution is not None:
        setup_render_settings(resolution)
    else:
        calc_resolution = max(8, base_resolution // (2 ** mip_level))
        setup_render_settings(calc_resolution)
    
    # Bake the cubemap
    success = False
    try:
        success = bake_single_cubemap(cube_probe, image, mip_level, output_dir)
    except Exception as e:
        print(f"Exception during baking for {image_name}: {e}")
    
    if not success:
        print(f"Baking failed for {image_name}")
        # Clean up the failed image
        try:
            bpy.data.images.remove(image)
        except:
            pass
        return False
    
    # Clean up the image after saving
    try:
        bpy.data.images.remove(image)
    except Exception as e:
        print(f"Warning: Failed to clean up image {image_name}: {e}")
    
    return True

def bake_cubemap(should_render_skybox=False, original_tonemap_setting=None, base_resolution=512):
    """Main function to bake the cubemap with multiple roughness levels."""
    try:
        output_dir = setup_output_directory()
        cube_probe = get_cube_probe()
        
        # Get dynamic mip levels and configurations based on resolution
        mip_levels = get_dynamic_mip_levels(base_resolution)
        cubemap_configs = get_dynamic_cubemap_configs(base_resolution)
        
        # Define roughness levels properly distributed from 0.0 to 1.0 across mip levels
        max_mip_levels = len(mip_levels)
        roughness_values = [i / (max_mip_levels - 1) for i in range(max_mip_levels)]
        
        # Process each roughness level
        for i, roughness in enumerate(roughness_values):
            # The mip level directly corresponds to the index
            mip_level = i
            
            # Create new image for this roughness level - use mip level in name instead of roughness
            image_name = f"cubemap_mip{mip_level}"
            
            # Bake the cubemap for this roughness level
            if not bake_and_save_cubemap(cube_probe, image_name, output_dir, mip_level, roughness, base_resolution=base_resolution):
                # Continue with next roughness level instead of failing completely
                continue
        
        # Create diffuse image using dynamic configuration
        diffuse_config = cubemap_configs["diffuse"]
        diffuse_image_name = "cubemap_diffuse"
        
        # Bake the diffuse cubemap with maximum roughness and calculated resolution
        if not bake_and_save_cubemap(cube_probe, diffuse_image_name, output_dir, 
                                   diffuse_config["mip_level"], 1.0, 
                                   diffuse_config["resolution"], base_resolution):
            print("Failed to bake diffuse cubemap")
        
        # Render skybox if requested
        if should_render_skybox:
            print("Loading cubemap_skybox")
            
            # Temporarily disable tonemap for skybox rendering
            if not set_tonemap(False):
                print("Warning: Failed to disable tonemap for skybox rendering")
            
            # Create skybox image using dynamic configuration
            skybox_config = cubemap_configs["skybox"]
            skybox_image_name = "cubemap_skybox"
            
            # Bake the skybox cubemap with minimum roughness and full resolution
            skybox_success = bake_and_save_cubemap(cube_probe, skybox_image_name, output_dir, 
                                                 skybox_config["mip_level"], 0.0, 
                                                 skybox_config["resolution"], base_resolution)
            
            if not skybox_success:
                print("Failed to bake skybox cubemap")
            
            # Restore original tonemap setting
            if original_tonemap_setting is not None:
                if not set_tonemap(original_tonemap_setting):
                    print("Warning: Failed to restore original tonemap setting")
        
        return True
        
    except Exception as e:
        print(f"Fatal error: {e}")
        return False

if __name__ == "__main__":
    try:
        # Get command line arguments passed after "--"
        argv = sys.argv
        
        # Find the index of the script
        script_index = argv.index("bake_cubemap.py")
        
        # Find the "--" separator
        try:
            separator_index = argv.index("--", script_index)
            # Get arguments after "--"
            args = argv[separator_index + 1:]
            
            # Parse texture path (first argument)
            if len(args) > 0:
                texture_path = args[0]
                print(f"Setting environment texture: {texture_path}")
                
                # Set the environment texture
                if not set_environment_texture(texture_path):
                    print("Failed to set environment texture, continuing with default")
            else:
                print("No environment texture path provided, using default")
                
            # Parse white point value (second argument)
            if len(args) > 1 and args[1] != "None":
                try:
                    white_point_value = float(args[1])
                    print(f"Setting white point value: {white_point_value}")
                    
                    # Set the white point value
                    if not set_white_point(white_point_value):
                        print("Failed to set white point value, continuing with default")
                except ValueError:
                    print(f"Invalid white point value: {args[1]}, must be a number")
            
            # Parse tonemap setting (third argument)
            if len(args) > 2:
                try:
                    tonemap_value = int(args[2])
                    should_tonemap = tonemap_value == 1
                    print(f"Setting tonemap: {'enabled' if should_tonemap else 'disabled'}")
                    
                    # Set the tonemap setting
                    if not set_tonemap(should_tonemap):
                        print("Failed to set tonemap setting, continuing with default")
                except ValueError:
                    print(f"Invalid tonemap value: {args[2]}, must be 0 or 1")
            
            # Parse skybox setting (fourth argument)
            should_render_skybox = False
            original_tonemap_setting = None
            if len(args) > 3:
                try:
                    skybox_value = int(args[3])
                    should_render_skybox = skybox_value == 1
                    print(f"Skybox rendering: {'enabled' if should_render_skybox else 'disabled'}")
                    
                    # Store the original tonemap setting if we need to render skybox
                    if should_render_skybox and len(args) > 2:
                        original_tonemap_setting = int(args[2]) == 1
                except ValueError:
                    print(f"Invalid skybox value: {args[3]}, must be 0 or 1")
            
            # Parse resolution setting (fifth argument)
            base_resolution = 512  # Default fallback value
            if len(args) > 4:
                try:
                    base_resolution = int(args[4])
                    print(f"Using base resolution: {base_resolution}x{base_resolution}")
                except ValueError:
                    print(f"Invalid resolution value: {args[4]}, must be a number")
                    base_resolution = 512
            
        except ValueError:
            print("No command line arguments provided, using default settings")
            base_resolution = 512
        
        if not bake_cubemap(should_render_skybox, original_tonemap_setting, base_resolution):
            print("Baking failed!")
            sys.exit(1)
    except Exception as e:
        print(f"Fatal error: {e}")
        sys.exit(1)