Blender Script to bake Displacement map from high poly model to low poly model

bake_displacement.py

import bpy

# uses method described here:
# https://blender.stackexchange.com/questions/228646/how-to-bake-displacement-map-from-another-mesh-not-multires/228659#228659

IMAGE_RES = 1024


def err(msg):
    
    def oops(self, context):
        self.layout.label(text=msg)
    
    bpy.context.window_manager.popup_menu(oops, title="Error", icon='ERROR')
    print("Error:", msg)


def create_material(obj, name):
    mat = bpy.data.materials.get("bake-" + name)
    if mat is None:
        mat = bpy.data.materials.new(name="bake-" + name)
        

    slot = obj.material_slots.find("bake-" + name)
    if slot == -1:
        obj.data.materials.append(mat)

    
    print("assigning material " + mat.name + " to " + obj.name)

    obj.active_material = mat
            
    mat.use_nodes = True
    
    nodes = mat.node_tree.nodes
    
    for node in nodes:
        nodes.remove(node)
    return mat

def prepare_material_highpoly(material, hp_pos):
    
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    
    out = nodes.new(type = 'ShaderNodeOutputMaterial')
    out.location = 0, 0
    em = nodes.new(type="ShaderNodeEmission")
    em.location = -200, 0
    math = nodes.new(type="ShaderNodeVectorMath")
    math.location = -400, 0
    geo = nodes.new(type="ShaderNodeNewGeometry")
    geo.location = -600, 0
    
    
    links.new(em.outputs[0], out.inputs[0])
    links.new(geo.outputs[0], math.inputs[0])
    links.new(math.outputs[0], em.inputs[0])
    math.inputs[1].default_value = hp_pos
    math.operation = "SUBTRACT"


def prepare_material_lowpoly1(material):
    nodes = material.node_tree.nodes
    links = material.node_tree.links

    img = nodes.new(type="ShaderNodeTexImage")
    img.location = 0, -200
    
    imgtex = bpy.data.images.new("tmp-highpoly-pos", IMAGE_RES, IMAGE_RES , float_buffer=True, is_data=True)
    img.image = imgtex
    nodes.active = img
    
    return imgtex

def prepare_material_lowpoly2(material, hp_image, lp_pos):
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    
    out = nodes.new(type = 'ShaderNodeOutputMaterial')
    out.location = 0, 0
    em = nodes.new(type="ShaderNodeEmission")
    em.location = -200, 0
    img = nodes.new(type="ShaderNodeTexImage")
    img.location = -1000, 400
    math3 = nodes.new(type="ShaderNodeVectorMath")
    math3.location = -400, 0
    math = nodes.new(type="ShaderNodeVectorMath")
    math.location = -600, 0
    math2 = nodes.new(type="ShaderNodeVectorMath")
    math2.location = -800, 0
    geo = nodes.new(type="ShaderNodeNewGeometry")
    geo.location = -1000, 0
    
    
    img.image = hp_image
    math.operation = 'SUBTRACT'
    math2.operation = 'SUBTRACT'
    math3.operation = 'LENGTH'
    math2.inputs[1].default_value = lp_pos
    
    
    links.new(em.outputs[0], out.inputs[0])
    links.new(math.outputs[0], math3.inputs[0])
    links.new(img.outputs[0], math.inputs[0])
    links.new(geo.outputs[0], math2.inputs[0])
    links.new(math2.outputs[0], math.inputs[1])
    links.new(math3.outputs[1], em.inputs[0])
    
    img = nodes.new(type="ShaderNodeTexImage")
    img.location = 0, -200
    
    imgtex = bpy.data.images.new("bake-displacement", IMAGE_RES, IMAGE_RES , float_buffer=True, is_data=True)
    img.image = imgtex
    nodes.active = img
    
    return img


def main():
    
    if len(bpy.context.selected_objects) != 2:
        err("Exactly one active and one selected object must exist")
        return
    
    lowpoly = bpy.context.active_object
    highpoly = next(x for x in bpy.context.selected_objects if x.name != lowpoly.name)
    
    if lowpoly is None:
        print("no object selected")
        return
    
    if lowpoly.location != highpoly.location:
        err("Lowpoly and highpoly objects are not at the same location!")
    
    render_engine_old =  bpy.context.scene.render.engine
    bpy.context.scene.render.engine = 'CYCLES'
    
    mat_lp = create_material(lowpoly, "lowpoly")
    lowpoly.active_material = mat_lp
    mat_hp = create_material(highpoly, "highpoly")
    highpoly.active_material = mat_hp
    prepare_material_highpoly(mat_hp, highpoly.location)
    img_high_poly_normals = prepare_material_lowpoly1(mat_lp)
    
    print("Bake 1 Starting")
    bpy.ops.object.bake(type="EMIT", use_selected_to_active=True)
    
    
    prepare_material_lowpoly2(create_material(lowpoly, "lowpoly"), img_high_poly_normals, lowpoly.location)
    
    highpoly.select_set(False)
    
    print("Bake 2 Starting")
    bpy.ops.object.bake(type="EMIT", use_selected_to_active=False, use_clear=True)
    
    
    bpy.context.scene.render.engine = render_engine_old
    print("Done")
    
main()