ControlNet + Inpainting Projection

__init__.py

@@ -25,7 +25,7 @@
 
 if current_process().name != "__actor__":
     import bpy
-    from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty, CollectionProperty, FloatProperty
+    from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty, CollectionProperty
     import sys
     import os
 
@@ -45,8 +45,9 @@ def clear_modules():
     from .operators.dream_texture import DreamTexture, kill_generator
     from .property_groups.dream_prompt import DreamPrompt
     from .property_groups.seamless_result import SeamlessResult
-    from .preferences import StableDiffusionPreferences
     from .ui.presets import register_default_presets
+
+    from . import engine
 
     requirements_path_items = (
         ('requirements/win-linux-cuda.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'),
@@ -101,13 +102,26 @@ def get_selection_preview(self):
         bpy.types.Scene.dream_textures_project_prompt = PointerProperty(type=DreamPrompt)
         bpy.types.Scene.dream_textures_project_framebuffer_arguments = EnumProperty(name="Inputs", items=framebuffer_arguments)
         bpy.types.Scene.dream_textures_project_bake = BoolProperty(name="Bake", default=False, description="Re-maps the generated texture onto the specified UV map")
+        def project_use_controlnet(self, context):
+            if self.dream_textures_project_use_control_net:
+                if len(self.dream_textures_project_prompt.control_nets) < 1:
+                    self.dream_textures_project_prompt.control_nets.add()
+            else:
+                self.dream_textures_project_prompt.control_nets.clear()
+        bpy.types.Scene.dream_textures_project_use_control_net = BoolProperty(name="Use ControlNet", default=False, description="Use a depth ControlNet instead of a depth model", update=project_use_controlnet)
+
+        engine.register()
 
         for cls in CLASSES:
             bpy.utils.register_class(cls)
 
         for tool in TOOLS:
             bpy.utils.register_tool(tool)
 
+        bpy.types.Scene.dream_textures_render_engine = PointerProperty(type=engine.DreamTexturesRenderEngineProperties)
+
+        bpy.types.RENDER_PT_context.append(engine.draw_device)
+
         # Monkey patch cycles render passes
         register_render_pass()
 
@@ -121,6 +135,10 @@ def unregister():
             bpy.utils.unregister_class(cls)
         for tool in TOOLS:
             bpy.utils.unregister_tool(tool)
+
+        bpy.types.RENDER_PT_context.remove(engine.draw_device)
+
+        engine.unregister()
 
         unregister_render_pass()
 

classes.py

@@ -6,13 +6,16 @@
 from .operators.upscale import Upscale
 from .operators.project import ProjectDreamTexture, dream_texture_projection_panels
 from .operators.notify_result import NotifyResult
+from .property_groups.control_net import ControlNet, SCENE_UL_ControlNetList, ControlNetsAdd, ControlNetsRemove
 from .property_groups.dream_prompt import DreamPrompt
 from .property_groups.seamless_result import SeamlessResult
 from .ui.panels import dream_texture, history, upscaling, render_properties
 from .preferences import OpenURL, StableDiffusionPreferences, ImportWeights, Model, ModelSearch, InstallModel, PREFERENCES_UL_ModelList
 
 from .ui.presets import DREAM_PT_AdvancedPresets, DREAM_MT_AdvancedPresets, AddAdvancedPreset, RestoreDefaultPresets
 
+from . import engine
+
 CLASSES = (
     *render_properties.render_properties_panels(),
 
@@ -29,12 +32,21 @@
     InpaintAreaBrushActivated,
     Upscale,
     ProjectDreamTexture,
+
+    SCENE_UL_ControlNetList,
+    ControlNetsAdd,
+    ControlNetsRemove,
 
     DREAM_PT_AdvancedPresets,
     DREAM_MT_AdvancedPresets,
     AddAdvancedPreset,
 
     NotifyResult,
+
+    engine.DreamTexturesRenderEngineProperties,
+    engine.DreamTexturesRenderEngine,
+    engine.NewEngineNodeTree,
+    *engine.engine_panels(),
 
     # The order these are registered in matters
     *dream_texture.dream_texture_panels(),
@@ -44,15 +56,17 @@
 )
 
 PREFERENCE_CLASSES = (
-                      PREFERENCES_UL_ModelList,
-                      ModelSearch,
-                      InstallModel,
-                      Model,
-                      DreamPrompt,
-                      SeamlessResult,
-                      UninstallDependencies,
-                      InstallDependencies,
-                      OpenURL,
-                      ImportWeights,
-                      RestoreDefaultPresets,
-                      StableDiffusionPreferences)
+    PREFERENCES_UL_ModelList,
+    ModelSearch,
+    InstallModel,
+    Model,
+    ControlNet,
+    DreamPrompt,
+    SeamlessResult,
+    UninstallDependencies,
+    InstallDependencies,
+    OpenURL,
+    ImportWeights,
+    RestoreDefaultPresets,
+    StableDiffusionPreferences,
+)

engine/__init__.py

@@ -0,0 +1,108 @@
+from .engine import *
+from .node_tree import *
+from .node_executor import *
+from .node import *
+from .nodes.input_nodes import *
+from .nodes.pipeline_nodes import *
+from .nodes.utility_nodes import *
+from .nodes.annotation_nodes import *
+from .annotations import openpose
+
+import bpy
+import nodeitems_utils
+
+class DreamTexturesNodeCategory(nodeitems_utils.NodeCategory):
+    @classmethod
+    def poll(cls, context):
+        return context.space_data.tree_type == DreamTexturesNodeTree.__name__
+
+categories = [
+    DreamTexturesNodeCategory("DREAM_TEXTURES_PIPELINE", "Pipeline", items = [
+        nodeitems_utils.NodeItem(NodeStableDiffusion.bl_idname),
+        nodeitems_utils.NodeItem(NodeControlNet.bl_idname),
+    ]),
+    DreamTexturesNodeCategory("DREAM_TEXTURES_INPUT", "Input", items = [
+        nodeitems_utils.NodeItem(NodeInteger.bl_idname),
+        nodeitems_utils.NodeItem(NodeString.bl_idname),
+        nodeitems_utils.NodeItem(NodeImage.bl_idname),
+        nodeitems_utils.NodeItem(NodeCollection.bl_idname),
+        nodeitems_utils.NodeItem(NodeRenderProperties.bl_idname),
+    ]),
+    DreamTexturesNodeCategory("DREAM_TEXTURES_UTILITY", "Utilities", items = [
+        nodeitems_utils.NodeItem(NodeMath.bl_idname),
+        nodeitems_utils.NodeItem(NodeRandomValue.bl_idname),
+        nodeitems_utils.NodeItem(NodeClamp.bl_idname),
+    ]),
+    DreamTexturesNodeCategory("DREAM_TEXTURES_ANNOTATIONS", "Annotations", items = [
+        nodeitems_utils.NodeItem(NodeAnnotationDepth.bl_idname),
+        nodeitems_utils.NodeItem(NodeAnnotationOpenPose.bl_idname),
+    ]),
+    DreamTexturesNodeCategory("DREAM_TEXTURES_GROUP", "Group", items = [
+        nodeitems_utils.NodeItem(bpy.types.NodeGroupOutput.__name__),
+    ]),
+]
+
+def register():
+    # Prompt
+    bpy.types.Scene.dream_textures_engine_prompt = bpy.props.PointerProperty(type=DreamPrompt)
+
+    # OpenPose
+    bpy.utils.register_class(openpose.ArmatureOpenPoseData)
+    bpy.types.Armature.dream_textures_openpose = bpy.props.PointerProperty(
+        type=openpose.ArmatureOpenPoseData
+    )
+    bpy.utils.register_class(openpose.BoneOpenPoseData)
+    bpy.types.Bone.dream_textures_openpose = bpy.props.PointerProperty(
+        type=openpose.BoneOpenPoseData
+    )
+
+    bpy.utils.register_class(DreamTexturesNodeTree)
+
+    # Nodes
+    bpy.utils.register_class(NodeSocketControlNet)
+    bpy.utils.register_class(NodeStableDiffusion)
+    bpy.utils.register_class(NodeControlNet)
+
+    bpy.utils.register_class(NodeInteger)
+    bpy.utils.register_class(NodeString)
+    bpy.utils.register_class(NodeCollection)
+    bpy.utils.register_class(NodeImage)
+    bpy.utils.register_class(NodeRenderProperties)
+
+    bpy.utils.register_class(NodeAnnotationDepth)
+    bpy.utils.register_class(NodeAnnotationOpenPose)
+
+    bpy.utils.register_class(NodeMath)
+    bpy.utils.register_class(NodeRandomValue)
+    bpy.utils.register_class(NodeClamp)
+
+    nodeitems_utils.register_node_categories("DREAM_TEXTURES_CATEGORIES", categories)
+
+def unregister():
+    # OpenPose
+    del bpy.types.Armature.dream_textures_openpose
+    bpy.utils.unregister_class(openpose.ArmatureOpenPoseData)
+    del bpy.types.Bone.dream_textures_openpose
+    bpy.utils.unregister_class(openpose.BoneOpenPoseData)
+
+    bpy.utils.unregister_class(DreamTexturesNodeTree)
+
+    # Nodes
+    bpy.utils.unregister_class(NodeSocketControlNet)
+    bpy.utils.unregister_class(NodeStableDiffusion)
+    bpy.utils.unregister_class(NodeControlNet)
+
+    bpy.utils.unregister_class(NodeInteger)
+    bpy.utils.unregister_class(NodeString)
+    bpy.utils.unregister_class(NodeCollection)
+    bpy.utils.unregister_class(NodeImage)
+    bpy.utils.unregister_class(NodeRenderProperties)
+
+    bpy.utils.unregister_class(NodeAnnotationDepth)
+    bpy.utils.unregister_class(NodeAnnotationOpenPose)
+
+    bpy.utils.unregister_class(NodeMath)
+    bpy.utils.unregister_class(NodeRandomValue)
+    bpy.utils.unregister_class(NodeClamp)
+
+    nodeitems_utils.unregister_node_categories("DREAM_TEXTURES_CATEGORIES")

engine/annotations/depth.py

@@ -0,0 +1,64 @@
+import bpy
+import gpu
+from gpu_extras.batch import batch_for_shader
+import numpy as np
+import threading
+
+def render_depth_map(context, collection=None, invert=True):
+    e = threading.Event()
+    result = None
+    def _execute():
+        nonlocal result
+        width, height = context.scene.render.resolution_x, context.scene.render.resolution_y
+        matrix = context.scene.camera.matrix_world.inverted()
+        projection_matrix = context.scene.camera.calc_matrix_camera(
+            context,
+            x=width,
+            y=height
+        )
+        offscreen = gpu.types.GPUOffScreen(width, height)
+
+        with offscreen.bind():
+            fb = gpu.state.active_framebuffer_get()
+            fb.clear(color=(0.0, 0.0, 0.0, 0.0))
+            gpu.state.depth_test_set('LESS_EQUAL')
+            gpu.state.depth_mask_set(True)
+            with gpu.matrix.push_pop():
+                gpu.matrix.load_matrix(matrix)
+                gpu.matrix.load_projection_matrix(projection_matrix)
+
+                shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR')
+
+                for object in (context.scene.objects if collection is None else collection.objects):
+                    object = object.evaluated_get(context)
+                    try:
+                        mesh = object.to_mesh(depsgraph=context).copy()
+                    except:
+                        continue
+                    if mesh is None:
+                        continue
+                    vertices = np.empty((len(mesh.vertices), 3), 'f')
+                    indices = np.empty((len(mesh.loop_triangles), 3), 'i')
+
+                    mesh.transform(object.matrix_world)
+                    mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3))
+                    mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3))
+
+                    batch = batch_for_shader(
+                        shader, 'TRIS',
+                        {"pos": vertices},
+                        indices=indices,
+                    )
+                    batch.draw(shader)
+            depth = np.array(fb.read_depth(0, 0, width, height).to_list())
+            if invert:
+                depth = 1 - depth
+                mask = np.array(fb.read_color(0, 0, width, height, 4, 0, 'UBYTE').to_list())[:, :, 3]
+                depth *= mask
+            depth = np.interp(depth, [np.ma.masked_equal(depth, 0, copy=False).min(), depth.max()], [0, 1]).clip(0, 1)
+        offscreen.free()
+        result = depth
+        e.set()
+    bpy.app.timers.register(_execute, first_interval=0)
+    e.wait()
+    return result