Add CoreML backend

.gitignore

@@ -1,4 +1,5 @@
 .DS_Store
+__packaged__
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

__init__.py

@@ -49,6 +49,8 @@ def clear_modules():
 
     from . import engine
 
+    from .diffusers_backend import DiffusersBackend
+
     requirements_path_items = (
         ('requirements/win-linux-cuda.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'),
         ('requirements/mac-mps-cpu.txt', 'Apple Silicon', 'Apple M1/M2'),
@@ -127,6 +129,9 @@ def project_use_controlnet(self, context):
         register_render_pass()
 
         register_default_presets()
+
+        # Register the default backend.
+        bpy.utils.register_class(DiffusersBackend)
 
     def unregister():
         for cls in PREFERENCE_CLASSES:
@@ -143,4 +148,7 @@ def unregister():
 
         unregister_render_pass()
 
+        # Unregister the default backend
+        bpy.utils.unregister_class(DiffusersBackend)
+
         kill_generator()

api/__init__.py

@@ -0,0 +1,2 @@
+from .models import *
+from .backend import *

api/backend/__init__.py

@@ -0,0 +1 @@
+from .backend import *

api/backend/backend.py

@@ -0,0 +1,112 @@
+try:
+    import bpy
+    from typing import Callable, List, Tuple
+    from ..models.generation_result import GenerationResult
+    from ..models.task import Task
+    from ..models.model import Model
+    from ..models.prompt import Prompt
+    from ..models.seamless_axes import SeamlessAxes
+    from ..models.step_preview_mode import StepPreviewMode
+
+    StepCallback = Callable[[GenerationResult], None]
+    Callback = Callable[[List[GenerationResult] | Exception], None]
+
+    class Backend(bpy.types.PropertyGroup):
+        """A backend for Dream Textures.
+
+        Provide the following methods to create a valid backend.
+
+        ```python
+        def list_models(self) -> List[Model]
+        def generate(
+            self,
+            task: Task,
+            model: Model,
+            prompt: Prompt,
+            size: Tuple[int, int] | None,
+            seamless_axes: SeamlessAxes,
+
+            step_callback: StepCallback,
+            callback: Callback
+        )
+        ```
+        """
+
+        @classmethod
+        def register(cls):
+            from ...property_groups.dream_prompt import DreamPrompt
+            setattr(DreamPrompt, cls._attribute(), bpy.props.PointerProperty(type=cls))
+
+        @classmethod
+        def unregister(cls):
+            from ...property_groups.dream_prompt import DreamPrompt
+            delattr(DreamPrompt, cls._attribute())
+
+        @classmethod
+        def _id(cls) -> str:
+            return f"{cls.__module__}.{cls.__name__}"
+
+        @classmethod
+        def _attribute(cls) -> str:
+            return cls._id().replace('.', '_')
+
+        @classmethod
+        def _lookup(cls, id):
+            return next(backend for backend in cls._list_backends() if backend._id() == id)
+
+        @classmethod
+        def _list_backends(cls):
+            return cls.__subclasses__()
+
+        def list_models(self, context) -> List[Model]:
+            """Provide a list of available models.
+
+            The `id` of the model will be provided 
+            """
+            ...
+
+        def list_schedulers(self, context) -> List[str]:
+            """Provide a list of available schedulers."""
+            ...
+
+        def draw_prompt(self, layout, context):
+            """Draw additional UI in the 'Prompt' panel"""
+            ...
+
+        def draw_advanced(self, layout, context):
+            """Draw additional UI in the 'Advanced' panel"""
+            ...
+
+        def draw_speed_optimizations(self, layout, context):
+            """Draw additional UI in the 'Speed Optimizations' panel"""
+            ...
+
+        def draw_memory_optimizations(self, layout, context):
+            """Draw additional UI in the 'Memory Optimizations' panel"""
+            ...
+
+        def draw_extra(self, layout, context):
+            """Draw additional UI in the panel"""
+            ...
+
+        def generate(
+            self,
+            task: Task,
+            model: Model,
+            prompt: Prompt,
+            size: Tuple[int, int] | None,
+            seed: int,
+            steps: int,
+            guidance_scale: float,
+            scheduler: str,
+            seamless_axes: SeamlessAxes,
+            step_preview_mode: StepPreviewMode,
+            iterations: int,
+
+            step_callback: StepCallback,
+            callback: Callback
+        ):
+            """A request to generate an image."""
+            ...
+except:
+    pass

api/models/__init__.py

@@ -0,0 +1,6 @@
+from .generation_result import *
+from .model import *
+from .prompt import *
+from .seamless_axes import *
+from .step_preview_mode import *
+from .task import *

api/models/generation_result.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from numpy.typing import NDArray
+
+@dataclass
+class GenerationResult:
+    """The output of a `Backend`.
+
+    Create a result with an `image` and a `seed`.
+
+    ```python
+    result = GenerationResult(
+        progress=3,
+        total=5,
+        image=np.zeros((512, 512, 3)),
+        seed=42
+    )
+    ```
+
+    Alternatively, create a result with just a `title` and progress values.
+
+    ```python
+    result = GenerationResult(
+        progress=3,
+        total=5,
+        title="Loading model"
+    )
+    ```
+    """
+
+    progress: int
+    """The amount out of `total` that has been completed"""
+
+    total: int
+    """The number of steps to complete"""
+
+    title: str | None = None
+    """The name of the currently executing task"""
+
+    image: NDArray | None = None
+    """The generated image as a Numpy array.
+
+    The shape should be `(height, width, channels)`, where `channels` is 3 or 4.
+    """
+
+    seed: int | None = None
+    """The seed used to generate the image."""

api/models/model.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+@dataclass
+class Model:
+    name: str
+    description: str
+    id: str

api/models/prompt.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+from typing import List
+
+@dataclass
+class Prompt:
+    positive: str | List[str]
+    negative: str | List[str] | None

api/models/seamless_axes.py

@@ -0,0 +1,75 @@
+from enum import Enum
+
+class SeamlessAxes(Enum):
+    """Unified handling of seamless axes.
+    Can be converted from str (id or text) or bool tuple/list (x, y).
+    Each enum is equal to their respective convertible values.
+    Special cases:
+        AUTO: None
+        OFF: False, empty str
+        BOTH: True
+    """
+
+    AUTO =       'auto', 'Auto-detect', None,  None
+    OFF =        'off',  'Off',         False, False
+    HORIZONTAL = 'x',    'X',           True,  False
+    VERTICAL =   'y',    'Y',           False, True
+    BOTH =       'xy',   'Both',        True,  True
+
+    def __init__(self, id, text, x, y):
+        self.id = id
+        self.text = text
+        self.x = x
+        self.y = y
+
+    def __eq__(self, other):
+        if isinstance(other, type(self)):
+            return self is other
+        if isinstance(other, str):
+            return self.id == other or self.text == other or (other == '' and self is self.OFF)
+        if isinstance(other, (tuple, list)) and len(other) == 2:
+            return self.x == other[0] and self.y == other[1]
+        if other is True and self is self.BOTH:
+            return True
+        if other is False and self is self.OFF:
+            return True
+        if other is None and self is self.AUTO:
+            return True
+        return False
+
+    def __and__(self, other):
+        return SeamlessAxes((self.x and other.x, self.y and other.y))
+
+    def __or__(self, other):
+        return SeamlessAxes((self.x or other.x, self.y or other.y))
+
+    def __xor__(self, other):
+        return SeamlessAxes((self.x != other.x, self.y != other.y))
+
+    def __invert__(self):
+        return SeamlessAxes((not self.x, not self.y))
+
+    @classmethod
+    def _missing_(cls, value):
+        if isinstance(value, str):
+            if value == '':
+                return cls.OFF
+            for e in cls:
+                if e.id == value or e.text == value:
+                    return e
+            raise ValueError(f'no {cls.__name__} with id {repr(id)}')
+        elif isinstance(value, (tuple, list)) and len(value) == 2:
+            for e in cls:
+                if e.x == value[0] and e.y == value[1]:
+                    return e
+            raise ValueError(f'no {cls.__name__} with x {value[0]} and y {value[1]}')
+        elif value is True:
+            return cls.BOTH
+        elif value is False:
+            return cls.OFF
+        elif value is None:
+            return cls.AUTO
+        raise TypeError(f'expected str, bool, tuple[bool, bool], or None, got {repr(value)}')
+
+    def bpy_enum(self, *args):
+        return self.id, self.text, *args

api/models/step_preview_mode.py

@@ -0,0 +1,8 @@
+import enum
+
+class StepPreviewMode(enum.Enum):
+    NONE = "None"
+    FAST = "Fast"
+    FAST_BATCH = "Fast (Batch Tiled)"
+    ACCURATE = "Accurate"
+    ACCURATE_BATCH = "Accurate (Batch Tiled)"

api/models/task.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+from typing import Tuple
+from numpy.typing import NDArray
+from enum import IntEnum
+
+class Task:
+    """A specific task type.
+
+    Access the properties of the task using dot notation.
+
+    ```python
+    # Task.ImageToImage
+    task.image
+    task.strength
+    task.fit
+    ```
+
+    Switch over the task to perform the correct actions.
+
+    ```python
+    match type(task):
+        case PromptToImage:
+            ...
+        case ImageToImage:
+            ...
+        case Inpaint:
+            ...
+        case DepthToImage:
+            ...
+        case Outpaint:
+            ...
+    ```
+    """
+    pass
+
+@dataclass
+class PromptToImage(Task):
+    pass
+
+@dataclass
+class ImageToImage(Task):
+    image: NDArray
+    strength: float
+    fit: bool
+
+@dataclass
+class Inpaint(Task):
+    class MaskSource(IntEnum):
+        ALPHA = 0
+        PROMPT = 1
+
+    image: NDArray
+    strength: float
+    fit: bool
+    mask_source: MaskSource
+    mask_prompt: str
+    confidence: float
+
+@dataclass
+class DepthToImage(Task):
+    depth: NDArray | None
+    image: NDArray | None
+    strength: float
+
+@dataclass
+class Outpaint(Task):
+    image: NDArray
+    origin: Tuple[int, int]

community_backends/coreml/README.md

@@ -0,0 +1,4 @@
+# CoreML Backend
+Faster inference on Apple Silicon with [apple/ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion).
+
+Converted mlpackages are stored in the directory specified by `DREAM_TEXTURES_COREML_HOME`, or `~/.cache/dream_textures_coreml` by default.