ap-canvas-creation-module-old/03_blender/sd_blender/__init__.py

# ----------------------------------------------------------
# File __init__.py
# ----------------------------------------------------------

# Imports
# Check if this add-on is being reloaded
if "bpy" in locals():
    # reloading .py files
    import bpy
    import importlib

    # from . import zs_renderscene as zsrs
    # importlib.reload(zsrs)


# or if this is the first load of this add-on
else:
    import bpy
    import json
    from pathlib import Path

    # from . import zs_renderscene as zsrs  # noqa


# Addon info
bl_info = {
    "name": "ZS Stable Diffusion Connection V0.0.1",
    "author": "Sergiu <sergiu@zixelise.com>",
    "Version": (0, 0, 1),
    "blender": (4, 00, 0),
    "category": "Scene",
    "description": "Stable Diffusion Connection",
}


# -------------------------------------------------------------------
# Functions
# -------------------------------------------------------------------


def load_scene():
    print("Loading Scene")
    # load scene data
    scene_data = load_scene_data()
    # create parent collections
    create_parent_collections("01_Products")
    create_parent_collections("02_Elements")
    create_parent_collections("03_Shapes")
    # append products
    products_data = load_objects_data(scene_data, "product")
    for index, product in enumerate(products_data):
        append_objects(product, index)
    # append elements
    # append shapes
    # set lighting
    # set camera


def load_scene_data():
    print("Loading Scene Data")
    # load scene data

    # to be replaced with actual data
    # open scene_info.json
    script_path = Path(__file__).resolve()
    scene_data_path = script_path.parent / "sample_scene" / "scene_info.json"
    with scene_data_path.open() as file:
        scene_data = json.load(file)
        print(scene_data)
    return scene_data


def load_objects_data(scene_data, object_type: str):
    print("Loading Assets Data")
    # load assets data
    objects_data = []
    for object in scene_data["scene"]["objects"]:
        if object["group_type"] == object_type:
            # get additional object data by id and combine with object data
            object_data = get_object_data_by_id(object["id"])
            object.update(object_data)
            objects_data.append(object)
    return objects_data

    # to be replaced with actual data


def get_object_data_by_id(object_id: str):
    print("Getting Object Data")
    # open assets_database.json
    script_path = Path(__file__).resolve()
    assets_data_path = script_path.parent / "sample_scene" / "assets_database.json"

    with assets_data_path.open() as file:
        assets_data = json.load(file)
    # get object data by id
    for object in assets_data:
        if object["id"] == object_id:
            return object


def append_objects(productInfo, index):
    print("Appending Objects")

    blendFileName = productInfo["name"]

    # visibleLayersJSONName = productInfo["Version"]

    # activeSwitchMaterials = json.dumps(productInfo["ActiveMaterials"])

    collectionName = blendFileName + "_" + str(index)

    append_active_layers(collectionName, productInfo)

    # replace_switch_materials(shotInfo, productInfo["ActiveMaterials"])

    link_collection_to_collection("01_Products", bpy.data.collections[collectionName])


def append_active_layers(newCollectionName, product_info):
    utility_collections = ["MaterialLibrary", "Animation_Controller", "Animation_Rig"]

    # visible_layers = utility_collections + product_info["VisibleLayers"]
    visible_layers = utility_collections

    script_path = Path(__file__).resolve().parent
    filePath = str(
        script_path / product_info["path"] / (product_info["name"] + ".blend")
    )

    print(filePath)

    # delete all objects and collections from product collection

    create_parent_collections(newCollectionName)

    # append active collections
    with bpy.data.libraries.load(filePath) as (data_from, data_to):
        data_to.collections = []
        for name in data_from.collections:

            if name in visible_layers:
                data_to.collections.append(name)
            if "NonConfigurable" in name:
                data_to.collections.append(name)

    # link appended colections to newCollection
    for collection in data_to.collections:
        # try:
        #     bpy.context.scene.collection.children.link(collection)
        # except:
        #     print(collection)
        link_collection_to_collection(newCollectionName, collection)

        # hide utility collections
        for utilityCollectionName in utility_collections:
            if utilityCollectionName in collection.name:
                # rename utility collections
                collection.name = newCollectionName + "_" + utilityCollectionName
                hide_collection(collection)

    #
    # make all objects in collection local
    for obj in bpy.data.collections[newCollectionName].all_objects:
        if obj.type == "MESH":
            obj.make_local()
            obj.data.make_local()

    # remove duplicated material slots
    mats = bpy.data.materials

    for obj in bpy.data.collections[newCollectionName].all_objects:
        for slot in obj.material_slots:
            part = slot.name.rpartition(".")
            if part[2].isnumeric() and part[0] in mats:
                slot.material = mats.get(part[0])


def create_parent_collections(group_name: str):

    if collection_exists(group_name):

        remove_collection_and_objects(group_name)
    else:
        create_collection(group_name)


# -------------------------------------------------------------------
# Utilities
# -------------------------------------------------------------------


def remove_collection_and_objects(collection_name):
    oldObjects = list(bpy.data.collections[collection_name].all_objects)
    for obj in oldObjects:
        bpy.data.objects.remove(obj, do_unlink=True)

    old_collection = bpy.data.collections[collection_name]
    if old_collection is not None:
        old_collection_names = get_subcollection_names(old_collection)
    else:
        print("Collection not found.")

    # print line break
    print("-----------------------------------------------------------------")
    print(old_collection_names)
    print("-----------------------------------------------------------------")
    for old_collection_name in old_collection_names:
        for collection in bpy.data.collections:
            if collection.name == old_collection_name:
                bpy.data.collections.remove(collection)

    bpy.ops.outliner.orphans_purge(
        do_local_ids=True, do_linked_ids=True, do_recursive=True
    )


def get_subcollection_names(collection):
    subcollection_names = []

    for child in collection.children:
        subcollection_names.append(child.name)
        subcollection_names.extend(get_subcollection_names(child))

    return subcollection_names


def link_collection_to_collection(parentCollectionName, childCollection):

    if bpy.context.scene.collection.children:
        parentCollection = bpy.context.scene.collection.children.get(
            parentCollectionName
        )
        # Add it to the main collection

        # childCollection = bpy.context.scene.collection.children.get(childCollection)

        parentCollection.children.link(childCollection)
        # if child collection is in scene collection unlink it
        if bpy.context.scene.collection.children.get(childCollection.name):
            bpy.context.scene.collection.children.unlink(childCollection)
        # bpy.context.scene.collection.children.unlink(childCollection)


# link collection to collection
def link_collection_to_collection_old(parentCollectionName, childCollection):

    if bpy.context.scene.collection.children:
        parentCollection = bpy.context.scene.collection.children.get(
            parentCollectionName
        )
        # Add it to the main collection

        try:
            childCollection = bpy.data.collections[childCollection]
        except:
            print("Collection not found.")
            return

        parentCollection.children.link(childCollection)
        bpy.context.scene.collection.children.unlink(childCollection)


# function that checks if a collection exists
def collection_exists(collection_name):
    return collection_name in bpy.data.collections


# function that creates a new collection and adds it to the scene
def create_collection(collection_name):
    new_collection = bpy.data.collections.new(collection_name)
    bpy.context.scene.collection.children.link(new_collection)


def hide_collection(collection):
    collection.hide_render = True
    collection.hide_viewport = True


# -------------------------------------------------------------------
# Operators
# -------------------------------------------------------------------
# load scene operator
class ZSSD_OT_LoadScene(bpy.types.Operator):
    bl_idname = "zs_sd_loader.load_scene"
    bl_label = "Load Scene"
    bl_description = "Load Scene"

    def execute(self, context):
        load_scene()
        return {"FINISHED"}


# parent class for panels
class ZSSDPanel:
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "ZS SD Loader"


# -------------------------------------------------------------------
#   Draw
# -------------------------------------------------------------------

# Panels


class ZSSD_PT_Main(ZSSDPanel, bpy.types.Panel):
    bl_label = "SD Loader"

    def draw(self, context):
        layout = self.layout

        scene = context.scene

        col = layout.column()

        self.is_connected = False

        col.label(text="Stable Diffusion Connection")

        # load scene button
        col.operator("zs_sd_loader.load_scene", text="Load Scene")


# modify after making products
blender_classes = [
    ZSSD_PT_Main,
    ZSSD_OT_LoadScene,
]


def register():

    # register classes
    for blender_class in blender_classes:
        bpy.utils.register_class(blender_class)

    # Has to be afqter class registering to correctly register property

    # register global properties

    # register list

    # list data

    # bpy.types.Scene.zs_product_list = bpy.props.CollectionProperty(
    #    type=ZS_Product_ListItem)

    # current item in list


def unregister():

    # unregister classes
    for blender_class in blender_classes:
        bpy.utils.unregister_class(blender_class)
    # unregister global properties

    # unregister list items
    # del bpy.types.Scene.my_list
    # del bpy.types.Scene.product_product_list_index