Bake image textures and export GLB using bake

gltf-export.py

from cmath import exp
import bpy

##### :: Script-Specific Functions :: #####

def deselect_all():
	bpy.ops.object.select_all(action='DESELECT')

def active_object(name):
	print(name)
	bpy.data.objects[name].select_set(True)
	bpy.context.view_layer.objects.active = bpy.data.objects[name]

def create_image(path, name):
    bpy.context.scene.render.image_settings.quality = 100

    img = bpy.data.images.new(name, 1200, 750)
    img.file_format = "JPEG"
    img.filepath_raw = path

    return img

def add_image_nodes(obj_name, img, node_name):
    for mat in bpy.data.objects[obj_name].material_slots:
        mat_nodes = mat.material.node_tree.nodes

        node_image_diffuse = mat_nodes.new("ShaderNodeTexImage")
        node_image_diffuse.image = img
        node_image_diffuse.name = node_name
        mat_nodes.active = node_image_diffuse

def cleanup_bake(obj_name, node_name, img):
    for mat in bpy.data.objects[obj_name].material_slots:
        node = mat.material.node_tree.nodes.get(node_name)
        mat.material.node_tree.nodes.remove(node)

def perform_bake(obj_name, bake_type, img, filter=set()):
    active_object(obj_name)
    bpy.ops.object.bake(type=bake_type, pass_filter=filter, use_clear=True)
    img.save()

##### :: Begin Export Process :: #####

deselect_all()

export_obj_name = "Exporting"
prebake_obj_name = 'ToExport'
prebake_object = bpy.data.objects[prebake_obj_name]

active_object(prebake_obj_name)

### Copy source object that will recieve baked textures

bpy.ops.object.duplicate(linked=False)
export_obj = bpy.context.view_layer.objects.active
export_obj.name = export_obj_name

export_obj.data.materials.clear()
uvs = [uv for uv in export_obj.data.uv_layers if not uv.name == "ExportUV"]
while uvs:
    export_obj.data.uv_layers.remove(uvs.pop())

export_material = bpy.data.materials.new("PostBake")
export_material.use_nodes = True

export_obj.active_material = export_material

### Start baking materials 

diffuse_path = os.path.join(os.path.dirname(bpy.data.filepath), "bake_diffuse.jpg")
rough_path = os.path.join(os.path.dirname(bpy.data.filepath), "bake_rough.jpg")
normal_path = os.path.join(os.path.dirname(bpy.data.filepath), "bake_normal.jpg")

# Diffuse Image
bake_diffuse = create_image(diffuse_path, "Bake_Diffuse")

add_image_nodes(prebake_obj_name, bake_diffuse, "Diffuse")
perform_bake(prebake_obj_name, "DIFFUSE", bake_diffuse, {'COLOR'})
cleanup_bake(prebake_obj_name, "Diffuse", bake_diffuse)

# Roughness Image
bake_rough = create_image(rough_path, "Bake_Rough")

add_image_nodes(prebake_obj_name, bake_rough, "Rough")
perform_bake(prebake_obj_name, "ROUGHNESS", bake_rough)
cleanup_bake(prebake_obj_name, "Rough", bake_rough)

# Normal Image
bake_normal = create_image(normal_path, "Bake_Normal")

add_image_nodes(prebake_obj_name, bake_normal, "Normal")
perform_bake(prebake_obj_name, "NORMAL", bake_normal)
cleanup_bake(prebake_obj_name, "Normal", bake_normal)

bpy.data.images.update()

### Setup post bake object

bpy.data.images.update()

nodes = export_material.node_tree.nodes
links = export_material.node_tree.links

bsdf_node = None
for node in nodes:
    if node.type == "BSDF_PRINCIPLED":
        bsdf_node = node
        break

bake_color_node = nodes.new('ShaderNodeTexImage')
bake_color_node.image = bake_diffuse
links.new(bsdf_node.inputs["Base Color"], bake_color_node.outputs["Color"])

bake_rough_node = nodes.new('ShaderNodeTexImage')
bake_rough_node.image = bake_rough
bake_rough_node.image.colorspace_settings.name = "Non-Color"
separate_color_node = nodes.new('ShaderNodeSeparateColor')
links.new(separate_color_node.inputs["Color"], bake_rough_node.outputs["Color"])
links.new(bsdf_node.inputs["Roughness"], separate_color_node.outputs["Green"])

bake_normal_node = nodes.new('ShaderNodeTexImage')
bake_normal_node.image = bake_normal
bake_normal_node.image.colorspace_settings.name = "Non-Color"
normal_map_node = nodes.new('ShaderNodeNormalMap')
links.new(normal_map_node.inputs["Color"], bake_normal_node.outputs["Color"])
links.new(bsdf_node.inputs["Normal"], normal_map_node.outputs["Normal"])

bpy.context.view_layer.update()

export_path = os.path.join(os.path.dirname(bpy.data.filepath), "export.glb")

### Actually export GLB

deselect_all()
active_object(export_obj_name)

bpy.ops.export_scene.gltf(
    filepath=export_path,
    export_format='GLB',
    use_selection=True
)

### Cleanup export-copied stuff

nodes.remove(bake_color_node)
nodes.remove(bake_rough_node)
nodes.remove(separate_color_node)
nodes.remove(bake_normal_node)
nodes.remove(normal_map_node)

bpy.data.images.remove(bake_diffuse)
bpy.data.images.remove(bake_rough)
bpy.data.images.remove(bake_normal)

bpy.data.materials.remove(export_material)
bpy.data.objects.remove(export_obj)