Improved Image Handling

api/models/generation_result.py

@@ -58,14 +58,15 @@ def tile_images(results: list['GenerationResult']) -> NDArray:
         height = images[0].shape[0]
         tiles_x = math.ceil(math.sqrt(len(images)))
         tiles_y = math.ceil(len(images) / tiles_x)
-        tiles = np.zeros((height * tiles_y, width * tiles_x, 4), dtype=np.float32)
+        tiles = np.zeros((height * tiles_y, width * tiles_x, images[0].shape[2]), dtype=images[0].dtype)
         bottom_offset = (tiles_x*tiles_y-len(images)) * width // 2
+        bottom = (tiles_y - 1) * height
         for i, image in enumerate(images):
             x = i % tiles_x
-            y = tiles_y - 1 - int((i - x) / tiles_x)
+            y = int((i - x) / tiles_x)
             x *= width
             y *= height
-            if y == 0:
+            if y == bottom:
                 x += bottom_offset
             tiles[y: y + height, x: x + width] = image
         return tiles

diffusers_backend.py

@@ -8,7 +8,6 @@
 from .api.models.fix_it_error import FixItError
 
 from .generator_process import Generator
-from .generator_process.actions.prompt_to_image import ImageGenerationResult
 from .generator_process.future import Future
 from .generator_process.models import CPUOffload, ModelType, Optimizations, Scheduler
 
@@ -236,24 +235,13 @@ def generate(self, arguments: GenerationArguments, step_callback: StepCallback,
                 )
             case _:
                 raise NotImplementedError()
-        def on_step(future: Future, step_image: ImageGenerationResult):
-            if len(step_image.images) == 0:
-                results = [(GenerationResult(progress=step_image.step, total=step_image.total or arguments.steps, seed=step_image.seeds[-1]))]
-            else:
-                results = [
-                    GenerationResult(progress=step_image.step, total=step_image.total or arguments.steps, image=step_image.images[i], seed=step_image.seeds[i])
-                    for i in range(len(step_image.images))
-                ]
-            should_continue = step_callback(results)
+        def on_step(future: Future, step_image: [GenerationResult]):
+            should_continue = step_callback(step_image)
             if not should_continue:
                 future.cancel()
                 callback(InterruptedError())
         def on_done(future: Future):
-            result: ImageGenerationResult = future.result(last_only=True)
-            callback([
-                GenerationResult(progress=result.step, total=arguments.steps, image=result.images[i], seed=result.seeds[i])
-                for i in range(len(result.images))
-            ])
+            callback(future.result(last_only=True))
         def on_exception(_, exception):
             callback(exception)
         future.add_response_callback(on_step)

engine/__init__.py

@@ -101,8 +101,7 @@ def register():
     bpy.utils.register_class(NodeString)
     bpy.utils.register_class(NodeCollection)
     bpy.utils.register_class(NodeImage)
-    if bpy.app.version >= (3, 5, 0):
-        bpy.utils.register_class(NodeImageFile)
+    bpy.utils.register_class(NodeImageFile)
     bpy.utils.register_class(NodeRenderProperties)
 
     bpy.utils.register_class(NodeAnnotationDepth)
@@ -118,8 +117,7 @@ def register():
     bpy.utils.register_class(NodeClamp)
     bpy.utils.register_class(NodeFramePath)
     bpy.utils.register_class(NodeCropImage)
-    if bpy.app.version >= (3, 5, 0):
-        bpy.utils.register_class(NodeResizeImage)
+    bpy.utils.register_class(NodeResizeImage)
     bpy.utils.register_class(NodeJoinImages)
     bpy.utils.register_class(NodeColorCorrect)
     bpy.utils.register_class(NodeSeparateColor)
@@ -152,8 +150,7 @@ def unregister():
     bpy.utils.unregister_class(NodeString)
     bpy.utils.unregister_class(NodeCollection)
     bpy.utils.unregister_class(NodeImage)
-    if bpy.app.version >= (3, 5, 0):
-        bpy.utils.unregister_class(NodeImageFile)
+    bpy.utils.unregister_class(NodeImageFile)
     bpy.utils.unregister_class(NodeRenderProperties)
 
     bpy.utils.unregister_class(NodeAnnotationDepth)
@@ -169,8 +166,7 @@ def unregister():
     bpy.utils.unregister_class(NodeClamp)
     bpy.utils.unregister_class(NodeFramePath)
     bpy.utils.unregister_class(NodeCropImage)
-    if bpy.app.version >= (3, 5, 0):
-        bpy.utils.unregister_class(NodeResizeImage)
+    bpy.utils.unregister_class(NodeResizeImage)
     bpy.utils.unregister_class(NodeJoinImages)
     bpy.utils.unregister_class(NodeColorCorrect)
     bpy.utils.unregister_class(NodeSeparateColor)

engine/engine.py

@@ -11,6 +11,7 @@
 from ..property_groups.dream_prompt import backend_options
 from .nodes.pipeline_nodes import NodeStableDiffusion
 from ..generator_process import actor
+from .. import image_utils
 
 class DreamTexturesRenderEngine(bpy.types.RenderEngine):
     """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles."""
@@ -29,19 +30,6 @@ def __del__(self):
 
     def render(self, depsgraph):
         scene = depsgraph.scene
-
-        def prepare_result(result):
-            if len(result.shape) == 2:
-                return np.concatenate(
-                    (
-                        np.stack((result,)*3, axis=-1),
-                        np.ones((*result.shape, 1))
-                    ),
-                    axis=-1
-                )
-            else:
-                return result
-
         result = self.begin_result(0, 0, scene.render.resolution_x, scene.render.resolution_y)
         layer = result.layers[0].passes["Combined"]
         self.update_result(result)
@@ -53,8 +41,7 @@ def node_begin(node):
             def node_update(response):
                 if isinstance(response, np.ndarray):
                     try:
-                        node_result = prepare_result(response)
-                        layer.rect = node_result.reshape(-1, node_result.shape[-1])
+                        image_utils.np_to_render_pass(response, layer, top_to_bottom=False)
                         self.update_result(result)
                     except:
                         pass
@@ -67,16 +54,15 @@ def node_end(_):
             for k, v in group_outputs:
                 if type(v) == int or type(v) == str or type(v) == float:
                     self.get_result().stamp_data_add_field(k, str(v))
-            node_result = prepare_result(node_result)
         except Exception as error:
             self.report({'ERROR'}, str(error))
             raise error
 
-        layer.rect = node_result.reshape(-1, node_result.shape[-1])
+        image_utils.np_to_render_pass(node_result, layer, top_to_bottom=False)
 
         if "Depth" in result.layers[0].passes:
             z = depth.render_depth_map(depsgraph, invert=True)
-            result.layers[0].passes["Depth"].rect = z.reshape((scene.render.resolution_x * scene.render.resolution_y, 1))
+            image_utils.np_to_render_pass(z, result.layers[0].passes["Depth"], top_to_bottom=False)
 
         self.end_result(result)
 

engine/nodes/annotation_nodes.py

@@ -13,7 +13,8 @@
 )
 
 def _update_annotation_inputs(self, context):
-    self.inputs['Collection'].enabled = self.src == 'collection'
+    inputs = {socket.name: socket for socket in self.inputs}
+    inputs['Collection'].enabled = self.src == 'collection'
 
 class NodeAnnotationDepth(DreamTexturesNode):
     bl_idname = "dream_textures.node_annotation_depth"

engine/nodes/input_nodes.py

@@ -5,6 +5,7 @@
 from ..node import DreamTexturesNode
 from ..annotations import openpose
 from ..annotations import depth
+from ... import image_utils
 
 class NodeString(DreamTexturesNode):
     bl_idname = "dream_textures.node_string"
@@ -68,16 +69,17 @@ def init(self, context):
 
     def draw_buttons(self, context, layout):
         layout.template_ID(self, "value", open="image.open")
+        if self.value is not None:
+            layout.prop(self.value.colorspace_settings, "name", text="Color Space")
 
     def execute(self, context):
-        result = np.array(self.value.pixels).reshape((*self.value.size, self.value.channels))
+        result = image_utils.bpy_to_np(self.value, color_space="Linear", top_to_bottom=False)
         context.update(result)
         return {
             'Image': result
         }
 
 class NodeImageFile(DreamTexturesNode):
-    """Requires Blender 3.5+ for OpenImageIO"""
     bl_idname = "dream_textures.node_image_file"
     bl_label = "Image File"
 
@@ -90,10 +92,7 @@ def draw_buttons(self, context, layout):
         pass
 
     def execute(self, context, path):
-        import OpenImageIO as oiio
-        image = oiio.ImageInput.open(path)
-        pixels = np.flipud(image.read_image('float'))
-        image.close()
+        pixels = image_utils.image_to_np(path, default_color_space="sRGB", to_color_space="Linear", top_to_bottom=False)
         context.update(pixels)
         return {
             'Image': pixels

engine/nodes/pipeline_nodes.py

@@ -14,6 +14,7 @@
 from ... import api
 from ...property_groups.seamless_result import SeamlessAxes
 import threading
+from ... import image_utils
 
 class NodeSocketControlNet(bpy.types.NodeSocket):
     bl_idname = "NodeSocketControlNet"
@@ -57,12 +58,13 @@ def control(self, context):
                     return np.flipud(ade20k.render_ade20k_map(context, collection=self.collection))
 
 def _update_stable_diffusion_sockets(self, context):
-    self.inputs['Source Image'].enabled = self.task in {'image_to_image', 'depth_to_image', 'inpaint'}
-    self.inputs['Noise Strength'].enabled = self.task in {'image_to_image', 'depth_to_image'}
+    inputs = {socket.name: socket for socket in self.inputs}
+    inputs['Source Image'].enabled = self.task in {'image_to_image', 'depth_to_image', 'inpaint'}
+    inputs['Noise Strength'].enabled = self.task in {'image_to_image', 'depth_to_image'}
     if self.task == 'depth_to_image':
-        self.inputs['Noise Strength'].default_value = 1.0
-    self.inputs['Depth Map'].enabled = self.task == 'depth_to_image'
-    self.inputs['ControlNets'].enabled = self.task != 'depth_to_image'
+        inputs['Noise Strength'].default_value = 1.0
+    inputs['Depth Map'].enabled = self.task == 'depth_to_image'
+    inputs['ControlNets'].enabled = self.task != 'depth_to_image'
 class NodeStableDiffusion(DreamTexturesNode):
     bl_idname = "dream_textures.node_stable_diffusion"
     bl_label = "Stable Diffusion"
@@ -109,14 +111,20 @@ def draw_buttons(self, context, layout):
     def execute(self, context, prompt, negative_prompt, width, height, steps, seed, cfg_scale, controlnets, depth_map, source_image, noise_strength):
         backend: api.Backend = self.prompt.get_backend()
 
+        if np.array(source_image).shape == (4,):
+            # the source image is a default color, ignore it.
+            source_image = None
+        else:
+            source_image = image_utils.color_transform(np.flipud(source_image), "Linear", "sRGB")
+
         def get_task():
             match self.task:
                 case 'prompt_to_image':
                     return api.PromptToImage()
                 case 'image_to_image':
                     return api.ImageToImage(source_image, noise_strength, fit=False)
                 case 'depth_to_image':
-                    return api.DepthToImage(depth_map, source_image, noise_strength)
+                    return api.DepthToImage(image_utils.grayscale(depth_map), source_image, noise_strength)
                 case 'inpaint':
                     return api.Inpaint(source_image, noise_strength, fit=False, mask_source=api.Inpaint.MaskSource.ALPHA, mask_prompt="", confidence=0)
 
@@ -140,16 +148,12 @@ def map_controlnet(c):
             iterations=1,
             control_nets=[map_controlnet(c) for c in controlnets] if isinstance(controlnets, list) else ([map_controlnet(controlnets)] if controlnets is not None else [])
         )
-
-        # the source image is a default color, ignore it.
-        if np.array(source_image).shape == (4,):
-            source_image = None
 
         event = threading.Event()
         result = None
         exception = None
         def step_callback(progress: List[api.GenerationResult]) -> bool:
-            context.update(progress[-1].image)
+            context.update(image_utils.image_to_np(progress[-1].image, default_color_space="sRGB", to_color_space="Linear", top_to_bottom=False))
             return True
             # if context.test_break():
             #     nonlocal result
@@ -163,7 +167,7 @@ def callback(results: List[api.GenerationResult] | Exception):
                 event.set()
             else:
                 nonlocal result
-                result = results[-1].image
+                result = image_utils.image_to_np(results[-1].image, default_color_space="sRGB", to_color_space="Linear", top_to_bottom=False)
                 event.set()
 
         backend = self.prompt.get_backend()
@@ -177,8 +181,9 @@ def callback(results: List[api.GenerationResult] | Exception):
         }
 
 def _update_control_net_sockets(self, context):
-    self.inputs['Collection'].enabled = self.input_type == 'collection'
-    self.inputs['Image'].enabled = self.input_type == 'image'
+    inputs = {socket.name: socket for socket in self.inputs}
+    inputs['Collection'].enabled = self.input_type == 'collection'
+    inputs['Image'].enabled = self.input_type == 'image'
 class NodeControlNet(DreamTexturesNode):
     bl_idname = "dream_textures.node_control_net"
     bl_label = "ControlNet"

engine/nodes/utility_nodes.py

@@ -3,6 +3,7 @@
 import random
 from ..node import DreamTexturesNode
 from ...property_groups.dream_prompt import seed_clamp
+from ... import image_utils
 
 class NodeMath(DreamTexturesNode):
     bl_idname = "dream_textures.node_math"
@@ -173,8 +174,7 @@ def draw_buttons(self, context, layout):
         pass
 
     def execute(self, context, image, width, height):
-        import OpenImageIO as oiio
-        result = oiio.ImageBufAlgo.resize(oiio.ImageBuf(image), roi=oiio.ROI(0, int(width), 0, int(height))).get_pixels()
+        result = image_utils.resize(image, (width, height))
         context.update(result)
         return {
             'Resized Image': result,

generator_process/__init__.py

@@ -1,4 +1,63 @@
-from .actor import Actor
+from typing import Callable
+
+from .actor import Actor, is_actor_process
+
+class RunInSubprocess(Exception):
+    """
+    Decorators to support running functions that are not defined under the Generator class in its subprocess.
+    This is to reduce what would otherwise be duplicate function definitions that logically don't belong to
+    the Generator, but require something in its subprocess (such as access to installed dependencies).
+    """
+
+    def __new__(cls, func=None):
+        if func is None:
+            # support `raise RunInSubprocess`
+            return super().__new__(cls)
+        return cls.always(func)
+
+    @staticmethod
+    def always(func):
+        if is_actor_process:
+            return func
+        def wrapper(*args, **kwargs):
+            return Generator.shared().call(wrapper, *args, **kwargs).result()
+        RunInSubprocess._copy_attributes(func, wrapper)
+        return wrapper
+
+    @staticmethod
+    def when(condition: bool | Callable[..., bool]):
+        if not isinstance(condition, Callable):
+            if condition:
+                return RunInSubprocess.always
+            return lambda x: x
+        def decorator(func):
+            if is_actor_process:
+                return func
+            def wrapper(*args, **kwargs):
+                if condition(*args, **kwargs):
+                    return Generator.shared().call(wrapper, *args, **kwargs).result()
+                return func(*args, **kwargs)
+            RunInSubprocess._copy_attributes(func, wrapper)
+            return wrapper
+        return decorator
+
+    @staticmethod
+    def when_raised(func):
+        if is_actor_process:
+            return func
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except RunInSubprocess:
+                return Generator.shared().call(wrapper, *args, **kwargs).result()
+        RunInSubprocess._copy_attributes(func, wrapper)
+        return wrapper
+
+    @staticmethod
+    def _copy_attributes(src, dst):
+        for n in ["__annotations__", "__doc__", "__name__", "__module__", "__qualname__"]:
+            if hasattr(src, n):
+                setattr(dst, n, getattr(src, n))
 
 class Generator(Actor):
     """
@@ -15,6 +74,9 @@ class Generator(Actor):
     from .actions.depth_to_image import depth_to_image
     from .actions.control_net import control_net
     from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models
-    from .actions.ocio_transform import ocio_transform
     from .actions.convert_original_stable_diffusion_to_diffusers import convert_original_stable_diffusion_to_diffusers
     from .actions.detect_seamless import detect_seamless
+
+    @staticmethod
+    def call(func, *args, **kwargs):
+        return func(*args, **kwargs)