Stability SDK for #405

__init__.py

@@ -110,6 +110,6 @@ def unregister():
         for tool in TOOLS:
             bpy.utils.unregister_tool(tool)
 
-        # unregister_render_pass()
+        unregister_render_pass()
 
         kill_generator()

generator_process/actions/image_to_image.py

@@ -36,8 +36,11 @@ def image_to_image(
 
     step_preview_mode: StepPreviewMode,
 
+    # Stability SDK
+    key: str | None = None,
+
     **kwargs
-) -> Generator[NDArray, None, None]:
+) -> Generator[ImageGenerationResult, None, None]:
     match pipeline:
         case Pipeline.STABLE_DIFFUSION:
             import diffusers
@@ -232,7 +235,40 @@ def __call__(
                         step_preview_mode=step_preview_mode
                     )
         case Pipeline.STABILITY_SDK:
-            import stability_sdk
-            raise NotImplementedError()
+            import stability_sdk.client
+            import stability_sdk.interfaces.gooseai.generation.generation_pb2
+            from PIL import Image, ImageOps
+            import io
+
+            if key is None:
+                raise ValueError("DreamStudio key not provided. Enter your key in the add-on preferences.")
+            client = stability_sdk.client.StabilityInference(key=key, engine=model)
+
+            if seed is None:
+                seed = random.randrange(0, np.iinfo(np.uint32).max)
+
+            answers = client.generate(
+                prompt=prompt,
+                width=width,
+                height=height,
+                cfg_scale=cfg_scale,
+                sampler=scheduler.stability_sdk(),
+                steps=steps,
+                seed=seed,
+                init_image=(Image.open(image) if isinstance(image, str) else Image.fromarray(image)).convert('RGB'),
+                start_schedule=strength,
+            )
+            for answer in answers:
+                for artifact in answer.artifacts:
+                    if artifact.finish_reason == stability_sdk.interfaces.gooseai.generation.generation_pb2.FILTER:
+                        raise ValueError("Your request activated DreamStudio's safety filter. Please modify your prompt and try again.")
+                    if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE:
+                        image = Image.open(io.BytesIO(artifact.binary))
+                        yield ImageGenerationResult(
+                            np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255.,
+                            seed,
+                            steps,
+                            True
+                        )
         case _:
             raise Exception(f"Unsupported pipeline {pipeline}.")

generator_process/actions/inpaint.py

@@ -36,6 +36,9 @@ def inpaint(
 
     step_preview_mode: StepPreviewMode,
 
+    # Stability SDK
+    key: str | None = None,
+
     **kwargs
 ) -> Generator[NDArray, None, None]:
     match pipeline:
@@ -272,7 +275,44 @@ def __call__(
                         step_preview_mode=step_preview_mode
                     )
         case Pipeline.STABILITY_SDK:
-            import stability_sdk
-            raise NotImplementedError()
+            import stability_sdk.client
+            import stability_sdk.interfaces.gooseai.generation.generation_pb2
+            from PIL import Image, ImageOps
+            import io
+
+            if key is None:
+                raise ValueError("DreamStudio key not provided. Enter your key in the add-on preferences.")
+            client = stability_sdk.client.StabilityInference(key=key, engine=model)
+
+            if seed is None:
+                seed = random.randrange(0, np.iinfo(np.uint32).max)
+
+            init_image = Image.open(image) if isinstance(image, str) else Image.fromarray(image)
+
+            answers = client.generate(
+                prompt=prompt,
+                width=width,
+                height=height,
+                cfg_scale=cfg_scale,
+                sampler=scheduler.stability_sdk(),
+                steps=steps,
+                seed=seed,
+                init_image=init_image.convert('RGB'),
+                mask_image=init_image.getchannel('A'),
+                start_schedule=strength,
+            )
+            for answer in answers:
+                for artifact in answer.artifacts:
+                    if artifact.finish_reason == stability_sdk.interfaces.gooseai.generation.generation_pb2.FILTER:
+                        raise ValueError("Your request activated DreamStudio's safety filter. Please modify your prompt and try again.")
+                    if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE:
+                        image = Image.open(io.BytesIO(artifact.binary))
+                        yield ImageGenerationResult(
+                            np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255.,
+                            seed,
+                            steps,
+                            True
+                        )
+
         case _:
             raise Exception(f"Unsupported pipeline {pipeline}.")

generator_process/actions/prompt_to_image.py

@@ -96,6 +96,22 @@ def scheduler_class():
             return scheduler_class().from_pretrained(pretrained['model_path'], subfolder=pretrained['subfolder'])
         else:
             return scheduler_class().from_config(pipeline.scheduler.config)
+
+    def stability_sdk(self):
+        import stability_sdk.interfaces.gooseai.generation.generation_pb2
+        match self:
+            case Scheduler.LMS_DISCRETE:
+                return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_K_LMS
+            case Scheduler.DDIM:
+                return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_DDIM
+            case Scheduler.DDPM:
+                return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_DDPM
+            case Scheduler.EULER_DISCRETE:
+                return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_K_EULER
+            case Scheduler.EULER_ANCESTRAL_DISCRETE:
+                return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_K_EULER_ANCESTRAL
+            case _:
+                raise ValueError(f"{self} cannot be used with DreamStudio.")
 
 @dataclass(eq=True)
 class Optimizations:
@@ -230,6 +246,9 @@ def prompt_to_image(
 
     step_preview_mode: StepPreviewMode,
 
+    # Stability SDK
+    key: str | None = None,
+
     **kwargs
 ) -> Generator[ImageGenerationResult, None, None]:
     match pipeline:
@@ -429,8 +448,39 @@ def __call__(
                         step_preview_mode=step_preview_mode
                     )
         case Pipeline.STABILITY_SDK:
-            import stability_sdk
-            raise NotImplementedError()
+            import stability_sdk.client
+            import stability_sdk.interfaces.gooseai.generation.generation_pb2
+            from PIL import Image, ImageOps
+            import io
+
+            if key is None:
+                raise ValueError("DreamStudio key not provided. Enter your key in the add-on preferences.")
+            client = stability_sdk.client.StabilityInference(key=key, engine=model)
+
+            if seed is None:
+                seed = random.randrange(0, np.iinfo(np.uint32).max)
+
+            answers = client.generate(
+                prompt=prompt,
+                width=width,
+                height=height,
+                cfg_scale=cfg_scale,
+                sampler=scheduler.stability_sdk(),
+                steps=steps,
+                seed=seed
+            )
+            for answer in answers:
+                for artifact in answer.artifacts:
+                    if artifact.finish_reason == stability_sdk.interfaces.gooseai.generation.generation_pb2.FILTER:
+                        raise ValueError("Your request activated DreamStudio's safety filter. Please modify your prompt and try again.")
+                    if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE:
+                        image = Image.open(io.BytesIO(artifact.binary))
+                        yield ImageGenerationResult(
+                            np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255.,
+                            seed,
+                            steps,
+                            True
+                        )
         case _:
             raise Exception(f"Unsupported pipeline {pipeline}.")
 

render_pass.py

@@ -2,8 +2,9 @@
 import cycles
 import numpy as np
 import os
-from .generator_process.actions.prompt_to_image import Pipeline, StepPreviewMode
+from .generator_process.actions.prompt_to_image import Pipeline, StepPreviewMode, ImageGenerationResult
 from .generator_process import Generator
+import threading
 
 update_render_passes_original = cycles.CyclesRender.update_render_passes
 render_original = cycles.CyclesRender.render
@@ -33,91 +34,27 @@ def render(self, depsgraph):
                     self.report({"ERROR"}, f"Image dimensions must be multiples of 64 (e.x. 512x512, 512x768, ...) closest is {round(size_x/64)*64}x{round(size_y/64)*64}")
                     return result
                 render_result = self.begin_result(0, 0, size_x, size_y)
-                for original_layer in original_result.layers:
-                    layer = None
-                    for layer_i in render_result.layers:
-                        if layer_i.name == original_layer.name:
-                            layer = layer_i
-                    for original_render_pass in original_layer.passes:
-                        render_pass = None
-                        for pass_i in layer.passes:
-                            if pass_i.name == original_render_pass.name:
-                                render_pass = pass_i
+                for layer in render_result.layers:
+                    for render_pass in layer.passes:
                         if render_pass.name == "Dream Textures":
-                            self.update_stats("Dream Textures", "Starting")
-
-                            # step_count = int(scene.dream_textures_render_properties_prompt.strength * scene.dream_textures_render_properties_prompt.steps)
-
-                            self.update_stats("Dream Textures", "Creating temporary image")
-                            combined_pass_image = bpy.data.images.new("dream_textures_post_processing_temp", width=size_x, height=size_y)
-
-                            rect = layer.passes["Combined"].rect
-
-                            combined_pixels = np.empty((size_x * size_y, 4), dtype=np.float32)
-                            rect.foreach_get(combined_pixels)
-
-                            gen = Generator.shared()
-                            self.update_stats("Dream Textures", "Applying color management transforms")
-                            combined_pixels = gen.ocio_transform(
-                                combined_pixels,
-                                config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'),
-                                exposure=scene.view_settings.exposure,
-                                gamma=scene.view_settings.gamma,
-                                view_transform=scene.view_settings.view_transform,
-                                display_device=scene.display_settings.display_device,
-                                look=scene.view_settings.look,
-                                inverse=False,
-                                _block=True
-                            ).result()
-
-                            self.update_stats("Dream Textures", "Generating...")
-
-                            generated_args = scene.dream_textures_render_properties_prompt.generate_args()
-                            generated_args['step_preview_mode'] = StepPreviewMode.NONE
-                            generated_args['width'] = size_x
-                            generated_args['height'] = size_y
-                            pixels = gen.image_to_image(
-                                image=(combined_pixels.reshape((size_x, size_y, 4)) * 255).astype(np.uint8),
-                                **generated_args,
-                                _block=True
-                            ).result()
-
-                            # Perform an inverse transform so when Blender applies its transform everything looks correct.
-                            self.update_stats("Dream Textures", "Applying inverse color management transforms")
-                            pixels = gen.ocio_transform(
-                                pixels,
-                                config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'),
-                                exposure=scene.view_settings.exposure,
-                                gamma=scene.view_settings.gamma,
-                                view_transform=scene.view_settings.view_transform,
-                                display_device=scene.display_settings.display_device,
-                                look=scene.view_settings.look,
-                                inverse=True,
-                                _block=True
-                            ).result()
-
-                            reshaped = pixels.reshape((size_x * size_y, 4))
-                            render_pass.rect.foreach_set(reshaped)
-
-                            # delete pointers before closing shared memory
-                            del pixels
-                            del combined_pixels
-                            del reshaped
-
-                            def cleanup():
-                                bpy.data.images.remove(combined_pass_image)
-                            bpy.app.timers.register(cleanup)
-                            self.update_stats("Dream Textures", "Finished")
+                            self._render_dream_textures_pass(layer, (size_x, size_y), scene, render_pass, render_result)
                         else:
-                            pixels = np.empty((len(original_render_pass.rect), len(original_render_pass.rect[0])), dtype=np.float32)
-                            original_render_pass.rect.foreach_get(pixels)
+                            source_pass = None
+                            for original_layer in original_result.layers:
+                                if layer.name == original_layer.name:
+                                    for original_pass in original_layer.passes:
+                                        if original_pass.name == render_pass.name:
+                                            source_pass = original_pass
+                            pixels = np.empty((len(source_pass.rect), len(source_pass.rect[0])), dtype=np.float32)
+                            source_pass.rect.foreach_get(pixels)
                             render_pass.rect[:] = pixels
                 self.end_result(render_result)
             except Exception as e:
                 print(e)
             return result
         return render
     cycles.CyclesRender.render = render_decorator(cycles.CyclesRender.render)
+    cycles.CyclesRender._render_dream_textures_pass = _render_dream_textures_pass
 
     # def del_decorator(original):
     #     def del_patch(self):
@@ -132,5 +69,74 @@ def unregister_render_pass():
     cycles.CyclesRender.update_render_passes = update_render_passes_original
     global render_original
     cycles.CyclesRender.render = render_original
+    del cycles.CyclesRender._render_dream_textures_pass
     # global del_original
     # cycles.CyclesRender.__del__ = del_original
+
+def _render_dream_textures_pass(self, layer, size, scene, render_pass, render_result):
+    self.update_stats("Dream Textures", "Starting")
+
+    rect = layer.passes["Combined"].rect
+
+    combined_pixels = np.empty((size[0] * size[1], 4), dtype=np.float32)
+    rect.foreach_get(combined_pixels)
+
+    gen = Generator.shared()
+    self.update_stats("Dream Textures", "Applying color management transforms")
+    combined_pixels = gen.ocio_transform(
+        combined_pixels,
+        config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'),
+        exposure=scene.view_settings.exposure,
+        gamma=scene.view_settings.gamma,
+        view_transform=scene.view_settings.view_transform,
+        display_device=scene.display_settings.display_device,
+        look=scene.view_settings.look,
+        inverse=False
+    ).result()
+
+    self.update_stats("Dream Textures", "Generating...")
+
+    generated_args = scene.dream_textures_render_properties_prompt.generate_args()
+    generated_args['width'] = size[0]
+    generated_args['height'] = size[1]
+    f = gen.image_to_image(
+        image=np.flipud(combined_pixels.reshape((size[1], size[0], 4)) * 255).astype(np.uint8),
+        **generated_args
+    )
+    event = threading.Event()
+    def on_step(_, step: ImageGenerationResult):
+        if step.final:
+            return
+        self.update_progress(step.step / generated_args['steps'])
+        if step.image is not None:
+            combined_pixels = step.image
+            render_pass.rect.foreach_set(combined_pixels.reshape((size[0] * size[1], 4)))
+            self.update_result(render_result) # This does not seem to have an effect.
+    def on_done(future):
+        nonlocal combined_pixels
+        result = future.result()
+        if isinstance(result, list):
+            result = result[-1]
+        combined_pixels = result.image
+        event.set()
+    f.add_response_callback(on_step)
+    f.add_done_callback(on_done)
+    event.wait()
+
+    # Perform an inverse transform so when Blender applies its transform everything looks correct.
+    self.update_stats("Dream Textures", "Applying inverse color management transforms")
+    combined_pixels = gen.ocio_transform(
+        combined_pixels.reshape((size[0] * size[1], 4)),
+        config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'),
+        exposure=scene.view_settings.exposure,
+        gamma=scene.view_settings.gamma,
+        view_transform=scene.view_settings.view_transform,
+        display_device=scene.display_settings.display_device,
+        look=scene.view_settings.look,
+        inverse=True
+    ).result()
+
+    combined_pixels = combined_pixels.reshape((size[0] * size[1], 4))
+    render_pass.rect.foreach_set(combined_pixels)
+
+    self.update_stats("Dream Textures", "Finished")

requirements/dreamstudio.txt

@@ -1,2 +1,2 @@
-stability-sdk==0.2.6
+stability-sdk
 opencolorio