Render Pass Fix (#370)

* use args instead of headless_prompt

* set events after completion instead of first image

generator_process/block_in_use.py

@@ -2,11 +2,16 @@ def block_in_use(func):
     def block(self, *args, **kwargs):
         if self.in_use:
             raise RuntimeError(f"Can't call {func.__qualname__} while process is in use")
-        try:
-            self.in_use = True
-            yield from func(self, *args, **kwargs)
-        finally:
-            self.in_use = False
+        self.in_use = True
+
+        # generator function is separate so in_use gets set immediately rather than waiting for first next() call
+        def sub():
+            try:
+                yield from func(self, *args, **kwargs)
+            finally:
+                self.in_use = False
+        return sub()
+
     # Pass the name through so we can use it in `setattr` on `GeneratorProcess`.
     block.__name__ = func.__name__
     return block

operators/dream_texture.py

@@ -238,12 +238,12 @@ def save_temp_image(img, path=None):
         args.update(headless_args)
         if headless_init_img is not None:
             args['use_init_img'] = True
-        if headless_prompt.prompt_structure == file_batch_structure.id:
+        if args['prompt_structure'] == file_batch_structure.id:
             args['prompt'] = [line.body for line in scene.dream_textures_prompt_file.lines if len(line.body.strip()) > 0]
         args['init_img'] = init_img_path
-        if headless_prompt.use_init_img_color:
+        if args['use_init_img_color']:
             args['init_color'] = init_img_path
-        if headless_prompt.backend == BackendTarget.STABILITY_SDK.name:
+        if args['backend'] == BackendTarget.STABILITY_SDK.name:
             args['dream_studio_key'] = context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.dream_studio_key
 
         def step_callback(step, width=None, height=None, shared_memory_name=None):
@@ -257,7 +257,7 @@ def image_callback(shared_memory_name, seed, width, height, upscaled=False):
             global headless_image_callback
             info() # clear variable
             nonlocal received_noncolorized
-            if headless_prompt.use_init_img and headless_prompt.use_init_img_color and not received_noncolorized:
+            if args['use_init_img'] and args['use_init_img_color'] and not received_noncolorized:
                 received_noncolorized = True
                 return
             received_noncolorized = False

render_pass.py

@@ -69,29 +69,25 @@ def image_callback(event, shared_memory_name, seed, width, height, upscaled=Fals
                                 nonlocal target_pixels
                                 nonlocal target_pixels_memory
                                 target_pixels[:] = np.frombuffer(target_pixels_memory.buf, dtype=np.float32).copy().reshape((size_x * size_y, 4))
-                                event.set()
                             def exception_callback(fatal, msg, trace):
                                 print(fatal, msg, trace)
-                                event.set()
                             generator_advance = GeneratorProcess.shared().apply_ocio_transforms(args, functools.partial(image_callback, event), exception_callback)
                             def timer():
                                 try:
                                     next(generator_advance)
                                     return 0.01
                                 except StopIteration:
-                                    pass
+                                    event.set()
                             bpy.app.timers.register(timer)
                         if render_pass.name == "Dream Textures":
                             self.update_stats("Dream Textures", "Starting")
-                            def image_callback(event, set_pixels, shared_memory_name, seed, width, height, upscaled=False):
+                            def image_callback(set_pixels, shared_memory_name, seed, width, height, upscaled=False):
                                 # Only use the non-upscaled texture, as upscaling is currently unsupported by the addon.
                                 if not upscaled:
                                     shared_memory = SharedMemory(shared_memory_name)
                                     set_pixels(np.frombuffer(shared_memory.buf, dtype=np.float32).copy().reshape((size_x * size_y, 4)))
 
                                     shared_memory.close()
-
-                                    event.set()
 
                             step_count = int(scene.dream_textures_render_properties_prompt.strength * scene.dream_textures_render_properties_prompt.steps)
                             def step_callback(step, width=None, height=None, shared_memory_name=None):
@@ -124,7 +120,13 @@ def set_pixels(npbuf):
                                 nonlocal pixels
                                 pixels = npbuf
                             def do_dream_texture_pass():
-                                dream_texture(scene.dream_textures_render_properties_prompt, step_callback, functools.partial(image_callback, event, set_pixels), combined_pass_image, width=size_x, height=size_y, show_steps=False, use_init_img_color=False)
+                                dream_texture(scene.dream_textures_render_properties_prompt, step_callback, functools.partial(image_callback, set_pixels), combined_pass_image, width=size_x, height=size_y, show_steps=False, use_init_img_color=False)
+                                gen = GeneratorProcess.shared(None, False)
+                                def waiter():
+                                    if gen.in_use:
+                                        return 0.01
+                                    event.set()
+                                bpy.app.timers.register(waiter)
                             bpy.app.timers.register(do_dream_texture_pass)
                             event.wait()