diff --git a/api/models/generation_result.py b/api/models/generation_result.py index 2ce5711e..a2022fdd 100644 --- a/api/models/generation_result.py +++ b/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 diff --git a/diffusers_backend.py b/diffusers_backend.py index 85e4ffd1..53758679 100644 --- a/diffusers_backend.py +++ b/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) diff --git a/engine/__init__.py b/engine/__init__.py index 4d98077c..ca52fdc8 100644 --- a/engine/__init__.py +++ b/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) diff --git a/engine/engine.py b/engine/engine.py index 2f75f748..bb241ed0 100644 --- a/engine/engine.py +++ b/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) diff --git a/engine/nodes/annotation_nodes.py b/engine/nodes/annotation_nodes.py index e98f16c7..aea223f6 100644 --- a/engine/nodes/annotation_nodes.py +++ b/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" diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index f217bdd3..256d69ce 100644 --- a/engine/nodes/input_nodes.py +++ b/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 diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index ea37f2e4..096bb253 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/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,6 +111,12 @@ 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': @@ -116,7 +124,7 @@ def get_task(): 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" diff --git a/engine/nodes/utility_nodes.py b/engine/nodes/utility_nodes.py index db5fbbbc..c71c2f77 100644 --- a/engine/nodes/utility_nodes.py +++ b/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, diff --git a/generator_process/__init__.py b/generator_process/__init__.py index 46c544a2..8afc2962 100644 --- a/generator_process/__init__.py +++ b/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) diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 9bedacde..95fcb60d 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -1,17 +1,14 @@ from typing import Union, Generator, Callable, List, Optional, Dict, Any from contextlib import nullcontext -from numpy.typing import NDArray import numpy as np import logging import os import random -from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, _configure_model_padding +from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, step_latents, step_images, _configure_model_padding from ...api.models.seamless_axes import SeamlessAxes from ..future import Future - - -logger = logging.getLogger(__name__) +from ...image_utils import image_to_np, rgb, resize, ImageOrPath def control_net( @@ -24,10 +21,10 @@ def control_net( optimizations: Optimizations, control_net: list[str | Checkpoint], - control: list[NDArray] | None, + control: list[ImageOrPath] | None, controlnet_conditioning_scale: list[float], - image: NDArray | str | None, # image to image + image: ImageOrPath | None, # image to image # inpaint inpaint: bool, inpaint_mask_src: str, @@ -57,8 +54,6 @@ def control_net( import diffusers import torch - import PIL.Image - import PIL.ImageOps device = self.choose_device(optimizations) @@ -92,29 +87,33 @@ def control_net( int(8 * (width // 8)), int(8 * (height // 8)), ) - control_image = [PIL.Image.fromarray(np.uint8(c * 255)).convert('RGB').resize(rounded_size) for c in control] if control is not None else None - init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).resize(rounded_size) + # StableDiffusionControlNetPipeline.check_image() currently fails without adding batch dimension + control_image = None if control is None else [image_to_np(c, mode="RGB", size=rounded_size)[np.newaxis] for c in control] + image = image_to_np(image, size=rounded_size) if inpaint: match inpaint_mask_src: case 'alpha': - mask_image = PIL.ImageOps.invert(init_image.getchannel('A')) + mask_image = 1-image[..., -1] + image = rgb(image) case 'prompt': + image = rgb(image) from transformers import AutoProcessor, CLIPSegForImageSegmentation - processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") + processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined", do_rescale=False) clipseg = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") - inputs = processor(text=[text_mask], images=[init_image.convert('RGB')], return_tensors="pt", padding=True) + inputs = processor(text=[text_mask], images=[image], return_tensors="pt", padding=True) outputs = clipseg(**inputs) - mask_image = PIL.Image.fromarray(np.uint8((1 - torch.sigmoid(outputs.logits).lt(text_mask_confidence).int().detach().numpy()) * 255), 'L').resize(init_image.size) + mask_image = (torch.sigmoid(outputs.logits) >= text_mask_confidence).detach().numpy().astype(np.float32) + mask_image = resize(mask_image, (width, height)) else: mask_image = None # Seamless if seamless_axes == SeamlessAxes.AUTO: - init_sa = None if init_image is None else self.detect_seamless(np.array(init_image) / 255) - control_sa = None if control_image is None else self.detect_seamless(np.array(control_image[0]) / 255) + init_sa = None if image is None else self.detect_seamless(image) + control_sa = None if control_image is None else self.detect_seamless(control_image[0][0]) if init_sa is not None and control_sa is not None: - seamless_axes = SeamlessAxes((init_sa.x and control_sa.x, init_sa.y and control_sa.y)) + seamless_axes = init_sa & control_sa elif init_sa is not None: seamless_axes = init_sa elif control_sa is not None: @@ -128,16 +127,16 @@ def control_net( def callback(step, timestep, latents): if future.check_cancelled(): raise InterruptedError() - future.add_response(ImageGenerationResult.step_preview(self, step_preview_mode, width, height, latents, generator, step)) + future.add_response(step_latents(pipe, step_preview_mode, latents, generator, step, steps)) try: - if init_image is not None: + if image is not None: if mask_image is not None: result = pipe( prompt=prompt, negative_prompt=negative_prompt if use_negative_prompt else None, control_image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, - image=init_image.convert('RGB'), + image=image, mask_image=mask_image, strength=strength, width=rounded_size[0], @@ -145,7 +144,8 @@ def callback(step, timestep, latents): num_inference_steps=steps, guidance_scale=cfg_scale, generator=generator, - callback=callback + callback=callback, + output_type="np" ) else: result = pipe( @@ -153,14 +153,15 @@ def callback(step, timestep, latents): negative_prompt=negative_prompt if use_negative_prompt else None, control_image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, - image=init_image.convert('RGB'), + image=image, strength=strength, width=rounded_size[0], height=rounded_size[1], num_inference_steps=steps, guidance_scale=cfg_scale, generator=generator, - callback=callback + callback=callback, + output_type="np" ) else: result = pipe( @@ -173,16 +174,11 @@ def callback(step, timestep, latents): num_inference_steps=steps, guidance_scale=cfg_scale, generator=generator, - callback=callback + callback=callback, + output_type="np" ) - future.add_response(ImageGenerationResult( - [np.asarray(PIL.ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in result.images], - [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()], - steps, - True - )) + future.add_response(step_images(result.images, generator, steps, steps)) except InterruptedError: pass diff --git a/generator_process/actions/depth_to_image.py b/generator_process/actions/depth_to_image.py index a2b06db8..f883f2ec 100644 --- a/generator_process/actions/depth_to_image.py +++ b/generator_process/actions/depth_to_image.py @@ -2,12 +2,12 @@ import os from contextlib import nullcontext -from numpy.typing import NDArray import numpy as np import random -from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, _configure_model_padding +from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, step_latents, step_images, _configure_model_padding from ...api.models.seamless_axes import SeamlessAxes from ..future import Future +from ...image_utils import image_to_np, ImageOrPath def depth_to_image( self, @@ -18,8 +18,8 @@ def depth_to_image( optimizations: Optimizations, - depth: NDArray | None, - image: NDArray | str | None, + depth: ImageOrPath | None, + image: ImageOrPath | None, strength: float, prompt: str | list[str], steps: int, @@ -44,7 +44,6 @@ def depth_to_image( import diffusers import torch import PIL.Image - import PIL.ImageOps class DreamTexturesDepth2ImgPipeline(diffusers.StableDiffusionInpaintPipeline): def prepare_depth(self, depth, image, dtype, device): @@ -56,7 +55,7 @@ def prepare_depth(self, depth, image, dtype, device): depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-large") depth_estimator = depth_estimator.to(device) - pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values + pixel_values = feature_extractor(images=image, return_tensors="pt", do_rescale=False).pixel_values pixel_values = pixel_values.to(device=device) # The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16. # So we use `torch.autocast` here for half precision inference. @@ -212,9 +211,8 @@ def __call__( # 4. Prepare the depth image depth = self.prepare_depth(depth_image, image, text_embeddings.dtype, device) - - if image is not None and isinstance(image, PIL.Image.Image): - image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess(image) + if image is not None: + image = self.image_processor.preprocess(image) # 5. set timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) @@ -350,15 +348,15 @@ def __call__( int(8 * (width // 8)), int(8 * (height // 8)), ) - depth_image = PIL.ImageOps.flip(PIL.Image.fromarray(np.uint8(depth * 255)).convert('L')).resize(rounded_size) if depth is not None else None - init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).convert('RGB').resize(rounded_size) + depth = image_to_np(depth, mode="L", size=rounded_size, to_color_space=None) + image = image_to_np(image, mode="RGB", size=rounded_size) # Seamless if seamless_axes == SeamlessAxes.AUTO: - init_sa = None if init_image is None else self.detect_seamless(np.array(init_image) / 255) - depth_sa = None if depth_image is None else self.detect_seamless(np.array(depth_image.convert('RGB')) / 255) + init_sa = None if image is None else self.detect_seamless(image) + depth_sa = None if depth is None else self.detect_seamless(depth) if init_sa is not None and depth_sa is not None: - seamless_axes = SeamlessAxes((init_sa.x and depth_sa.x, init_sa.y and depth_sa.y)) + seamless_axes = init_sa & depth_sa elif init_sa is not None: seamless_axes = init_sa elif depth_sa is not None: @@ -371,29 +369,24 @@ def __call__( def callback(step, timestep, latents): if future.check_cancelled(): raise InterruptedError() - future.add_response(ImageGenerationResult.step_preview(self, step_preview_mode, width, height, latents, generator, step)) + future.add_response(step_latents(pipe, step_preview_mode, latents, generator, step, steps)) try: result = pipe( prompt=prompt, negative_prompt=negative_prompt if use_negative_prompt else None, - depth_image=depth_image, - image=init_image, + depth_image=depth, + image=image, strength=strength, width=rounded_size[0], height=rounded_size[1], num_inference_steps=steps, guidance_scale=cfg_scale, generator=generator, - callback=callback + callback=callback, + output_type="np" ) - future.add_response(ImageGenerationResult( - [np.asarray(PIL.ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in result.images], - [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()], - steps, - True - )) + future.add_response(step_images(result.images, generator, steps, steps)) except InterruptedError: pass diff --git a/generator_process/actions/detect_seamless/__init__.py b/generator_process/actions/detect_seamless/__init__.py index 48cba657..a20846ab 100644 --- a/generator_process/actions/detect_seamless/__init__.py +++ b/generator_process/actions/detect_seamless/__init__.py @@ -4,8 +4,9 @@ from numpy.typing import NDArray from ....api.models.seamless_axes import SeamlessAxes +from .... import image_utils -def detect_seamless(self, image: NDArray) -> SeamlessAxes: +def detect_seamless(self, image: image_utils.ImageOrPath) -> SeamlessAxes: import os import torch from torch import nn @@ -58,9 +59,7 @@ def forward(self, x: torch.Tensor): else: device = 'cpu' - if image.shape[2] == 4: - # only trained on RGB channels, not alpha - image = image[:, :, :3] + image = image_utils.image_to_np(image, mode="RGB") # slice 8 pixels off each edge and combine opposing sides where the seam/seamless portion is in the middle # may trim up to 3 pixels off the length of each edge to make them a multiple of 4 diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index 8a60a4f1..43a7837f 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -5,9 +5,10 @@ from numpy.typing import NDArray import numpy as np import random -from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, _configure_model_padding +from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, step_latents, step_images, _configure_model_padding from ...api.models.seamless_axes import SeamlessAxes from ..future import Future +from ...image_utils import image_to_np, size, resize, ImageOrPath def image_to_image( @@ -19,7 +20,7 @@ def image_to_image( optimizations: Optimizations, - image: NDArray, + image: ImageOrPath, fit: bool, strength: float, prompt: str | list[str], @@ -46,8 +47,6 @@ def image_to_image( import diffusers import torch - from PIL import Image, ImageOps - import PIL.Image device = self.choose_device(optimizations) @@ -68,19 +67,17 @@ def image_to_image( generator = generator[0] # Init Image - init_image = Image.fromarray(image).convert('RGB') - + image = image_to_np(image, mode="RGB") if fit: height = height or pipe.unet.config.sample_size * pipe.vae_scale_factor width = width or pipe.unet.config.sample_size * pipe.vae_scale_factor - init_image = init_image.resize((width, height)) + image = resize(image, (width, height)) else: - width = init_image.width - height = init_image.height + width, height = size(image) # Seamless if seamless_axes == SeamlessAxes.AUTO: - seamless_axes = self.detect_seamless(np.array(init_image) / 255) + seamless_axes = self.detect_seamless(image) _configure_model_padding(pipe.unet, seamless_axes) _configure_model_padding(pipe.vae, seamless_axes) @@ -89,25 +86,20 @@ def image_to_image( def callback(step, timestep, latents): if future.check_cancelled(): raise InterruptedError() - future.add_response(ImageGenerationResult.step_preview(self, step_preview_mode, width, height, latents, generator, step)) + future.add_response(step_latents(pipe, step_preview_mode, latents, generator, step, steps)) try: result = pipe( prompt=prompt, negative_prompt=negative_prompt if use_negative_prompt else None, - image=[init_image] * batch_size, + image=[image] * batch_size, strength=strength, num_inference_steps=steps, guidance_scale=cfg_scale, generator=generator, - callback=callback + callback=callback, + output_type="np" ) - future.add_response(ImageGenerationResult( - [np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in result.images], - [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()], - steps, - True - )) + future.add_response(step_images(result.images, generator, steps, steps)) except InterruptedError: pass diff --git a/generator_process/actions/inpaint.py b/generator_process/actions/inpaint.py index b163db58..23ea8645 100644 --- a/generator_process/actions/inpaint.py +++ b/generator_process/actions/inpaint.py @@ -1,12 +1,12 @@ from typing import Union, Generator, Callable, List, Optional import os from contextlib import nullcontext -from numpy.typing import NDArray import numpy as np import random -from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, _configure_model_padding +from .prompt_to_image import Checkpoint, Scheduler, Optimizations, StepPreviewMode, step_latents, step_images, _configure_model_padding from ...api.models.seamless_axes import SeamlessAxes from ..future import Future +from ...image_utils import image_to_np, size, resize, rgb, ImageOrPath def inpaint( self, @@ -17,7 +17,7 @@ def inpaint( optimizations: Optimizations, - image: NDArray, + image: ImageOrPath, fit: bool, strength: float, prompt: str | list[str], @@ -44,14 +44,12 @@ def inpaint( key: str | None = None, **kwargs -) -> Generator[NDArray, None, None]: +) -> Generator[Future, None, None]: future = Future() yield future import diffusers import torch - from PIL import Image, ImageOps - import PIL.Image device = self.choose_device(optimizations) @@ -74,11 +72,17 @@ def inpaint( generator = generator[0] # Init Image - init_image = Image.fromarray(image) + image = image_to_np(image) + if fit: + height = height or pipe.unet.config.sample_size * pipe.vae_scale_factor + width = width or pipe.unet.config.sample_size * pipe.vae_scale_factor + image = resize(image, (width, height)) + else: + width, height = size(image) # Seamless if seamless_axes == SeamlessAxes.AUTO: - seamless_axes = self.detect_seamless(np.array(init_image) / 255) + seamless_axes = self.detect_seamless(image) _configure_model_padding(pipe.unet, seamless_axes) _configure_model_padding(pipe.vae, seamless_axes) @@ -86,42 +90,40 @@ def inpaint( with torch.inference_mode() if device not in ('mps', "dml") else nullcontext(): match inpaint_mask_src: case 'alpha': - mask_image = ImageOps.invert(init_image.getchannel('A')) + mask_image = 1-image[..., -1] + image = rgb(image) case 'prompt': + image = rgb(image) from transformers import AutoProcessor, CLIPSegForImageSegmentation - processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") + processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined", do_rescale=False) clipseg = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") - inputs = processor(text=[text_mask], images=[init_image.convert('RGB')], return_tensors="pt", padding=True) + inputs = processor(text=[text_mask], images=[image], return_tensors="pt", padding=True) outputs = clipseg(**inputs) - mask_image = Image.fromarray(np.uint8((1 - torch.sigmoid(outputs.logits).lt(text_mask_confidence).int().detach().numpy()) * 255), 'L').resize(init_image.size) + mask_image = (torch.sigmoid(outputs.logits) >= text_mask_confidence).detach().numpy().astype(np.float32) + mask_image = resize(mask_image, (width, height)) def callback(step, timestep, latents): if future.check_cancelled(): raise InterruptedError() - future.add_response(ImageGenerationResult.step_preview(self, step_preview_mode, width, height, latents, generator, step)) + future.add_response(step_latents(pipe, step_preview_mode, latents, generator, step, steps)) try: result = pipe( prompt=prompt, negative_prompt=negative_prompt if use_negative_prompt else None, - image=[init_image.convert('RGB')] * batch_size, + image=[image] * batch_size, mask_image=[mask_image] * batch_size, strength=strength, - height=init_image.size[1] if fit else height, - width=init_image.size[0] if fit else width, + height=height, + width=width, num_inference_steps=steps, guidance_scale=cfg_scale, generator=generator, - callback=callback + callback=callback, + output_type="np" ) - future.add_response(ImageGenerationResult( - [np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in result.images], - [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()], - steps, - True - )) + future.add_response(step_images(result.images, generator, steps, steps)) except InterruptedError: pass diff --git a/generator_process/actions/ocio_transform.py b/generator_process/actions/ocio_transform.py deleted file mode 100644 index 35ec8cfc..00000000 --- a/generator_process/actions/ocio_transform.py +++ /dev/null @@ -1,98 +0,0 @@ -import sys -from numpy.typing import NDArray -import math - -def ocio_transform( - self, - input_image: NDArray, - config_path: str, - exposure: float, - gamma: float, - view_transform: str, - display_device: str, - look: str, - inverse: bool -): - import PyOpenColorIO as OCIO - - ocio_config = OCIO.Config.CreateFromFile(config_path) - - # A reimplementation of `OCIOImpl::createDisplayProcessor` from the Blender source. - # https://github.com/dfelinto/blender/blob/87a0770bb969ce37d9a41a04c1658ea09c63933a/intern/opencolorio/ocio_impl.cc#L643 - def create_display_processor( - config, - input_colorspace, - view, - display, - look, - scale, # Exposure - exponent, # Gamma - inverse=False - ): - group = OCIO.GroupTransform() - - # Exposure - if scale != 1: - # Always apply exposure in scene linear. - color_space_transform = OCIO.ColorSpaceTransform() - color_space_transform.setSrc(input_colorspace) - color_space_transform.setDst(OCIO.ROLE_SCENE_LINEAR) - group.appendTransform(color_space_transform) - - # Make further transforms aware of the color space change - input_colorspace = OCIO.ROLE_SCENE_LINEAR - - # Apply scale - matrix_transform = OCIO.MatrixTransform([scale, 0.0, 0.0, 0.0, 0.0, scale, 0.0, 0.0, 0.0, 0.0, scale, 0.0, 0.0, 0.0, 0.0, 1.0]) - group.appendTransform(matrix_transform) - - # Add look transform - use_look = look is not None and len(look) > 0 - if use_look: - look_output = config.getLook(look).getProcessSpace() - if look_output is not None and len(look_output) > 0: - look_transform = OCIO.LookTransform() - look_transform.setSrc(input_colorspace) - look_transform.setDst(look_output) - look_transform.setLooks(look) - group.appendTransform(look_transform) - # Make further transforms aware of the color space change. - input_colorspace = look_output - else: - # For empty looks, no output color space is returned. - use_look = False - - # Add view and display transform - display_view_transform = OCIO.DisplayViewTransform() - display_view_transform.setSrc(input_colorspace) - display_view_transform.setLooksBypass(True) - display_view_transform.setView(view) - display_view_transform.setDisplay(display) - group.appendTransform(display_view_transform) - - # Gamma - if exponent != 1: - exponent_transform = OCIO.ExponentTransform([exponent, exponent, exponent, 1.0]) - group.appendTransform(exponent_transform) - - # Create processor from transform. This is the moment were OCIO validates - # the entire transform, no need to check for the validity of inputs above. - try: - if inverse: - group.setDirection(OCIO.TransformDirection.TRANSFORM_DIR_INVERSE) - processor = config.getProcessor(group) - if processor is not None: - return processor - except Exception as e: - self.send_exception(True, msg=str(e), trace="") - - return None - - # Exposure and gamma transformations derived from Blender source: - # https://github.com/dfelinto/blender/blob/87a0770bb969ce37d9a41a04c1658ea09c63933a/source/blender/imbuf/intern/colormanagement.c#L825 - scale = math.pow(2, exposure) - exponent = 1 if gamma == 1 else (1 / (gamma if gamma > sys.float_info.epsilon else sys.float_info.epsilon)) - processor = create_display_processor(ocio_config, OCIO.ROLE_SCENE_LINEAR, view_transform, display_device, look if look != 'None' else None, scale, exponent, inverse) - - processor.getDefaultCPUProcessor().applyRGBA(input_image) - return input_image \ No newline at end of file diff --git a/generator_process/actions/outpaint.py b/generator_process/actions/outpaint.py index b0f3d95c..faa452e6 100644 --- a/generator_process/actions/outpaint.py +++ b/generator_process/actions/outpaint.py @@ -1,13 +1,13 @@ from typing import Tuple, Generator -from numpy.typing import NDArray import numpy as np -from .prompt_to_image import ImageGenerationResult from ..future import Future +from ...api.models.generation_result import GenerationResult +from ...image_utils import image_to_np, rgba, ImageOrPath def outpaint( self, - image: NDArray, + image: ImageOrPath, width: int | None, height: int | None, @@ -15,61 +15,49 @@ def outpaint( outpaint_origin: Tuple[int, int], **kwargs -) -> Generator[ImageGenerationResult, None, None]: - from PIL import Image, ImageOps +) -> Generator[Future, None, None]: future = Future() yield future - init_image = Image.fromarray(image) width = width or 512 height = height or 512 + image = image_to_np(image) - if outpaint_origin[0] > init_image.size[0] or outpaint_origin[0] < -width: - raise ValueError(f"Outpaint origin X ({outpaint_origin[0]}) must be between {-width} and {init_image.size[0]}") - if outpaint_origin[1] > init_image.size[1] or outpaint_origin[1] < -height: - raise ValueError(f"Outpaint origin Y ({outpaint_origin[1]}) must be between {-height} and {init_image.size[1]}") - - outpaint_bounds = Image.new( - 'RGBA', - ( - max(init_image.size[0], outpaint_origin[0] + width) - min(0, outpaint_origin[0]), - max(init_image.size[1], outpaint_origin[1] + height) - min(0, outpaint_origin[1]), - ), - (0, 0, 0, 0) - ) - outpaint_bounds.paste( - init_image, - ( - 0 if outpaint_origin[0] > 0 else -outpaint_origin[0], - 0 if outpaint_origin[1] > 0 else -outpaint_origin[1], - ) - ) + if outpaint_origin[0] > image.shape[1] or outpaint_origin[0] < -width: + raise ValueError(f"Outpaint origin X ({outpaint_origin[0]}) must be between {-width} and {image.shape[1]}") + if outpaint_origin[1] > image.shape[0] or outpaint_origin[1] < -height: + raise ValueError(f"Outpaint origin Y ({outpaint_origin[1]}) must be between {-height} and {image.shape[0]}") + + outpaint_bounds = np.zeros(( + max(image.shape[0], outpaint_origin[1] + height) - min(0, outpaint_origin[1]), + max(image.shape[1], outpaint_origin[0] + width) - min(0, outpaint_origin[0]), + 4 + ), dtype=np.float32) + + def paste(under, over, offset): + under[offset[0]:offset[0] + over.shape[0], offset[1]:offset[1] + over.shape[1]] = over + return under + + paste(outpaint_bounds, image, ( + 0 if outpaint_origin[1] > 0 else -outpaint_origin[1], + 0 if outpaint_origin[0] > 0 else -outpaint_origin[0] + )) + offset_origin = ( - max(outpaint_origin[0], 0), # left max(outpaint_origin[1], 0), # upper + max(outpaint_origin[0], 0), # left ) # Crop out the area to generate - inpaint_tile = outpaint_bounds.crop( - ( - *offset_origin, - offset_origin[0] + width, # right - offset_origin[1] + height, # lower - ) - ) + inpaint_tile = outpaint_bounds[offset_origin[0]:offset_origin[0]+height, offset_origin[1]:offset_origin[1]+width] - def process(_, step: ImageGenerationResult): - for i, result_image in enumerate(step.images): - image = outpaint_bounds.copy() - image.paste( - ImageOps.flip(Image.fromarray((result_image * 255.).astype(np.uint8))), - offset_origin - ) - step.images[i] = np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + def process(_, step: [GenerationResult]): + for res in step: + res.image = paste(outpaint_bounds.copy(), rgba(res.image), offset_origin) future.add_response(step) inpaint_generator = self.inpaint( - image=np.array(inpaint_tile), + image=inpaint_tile, width=width, height=height, **kwargs diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 94d20213..95f8c97c 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -1,3 +1,4 @@ +import functools from typing import Generator from contextlib import nullcontext @@ -6,7 +7,7 @@ from ...api.models.seamless_axes import SeamlessAxes from ...api.models.step_preview_mode import StepPreviewMode from ..models import Checkpoint, Optimizations, Scheduler -from ..models.image_generation_result import ImageGenerationResult +from ..models.image_generation_result import step_latents, step_images from ..future import Future def prompt_to_image( @@ -46,7 +47,6 @@ def prompt_to_image( import diffusers import torch - from PIL import Image, ImageOps device = self.choose_device(optimizations) @@ -79,11 +79,11 @@ def prompt_to_image( # Inference with torch.inference_mode() if device not in ('mps', "dml") else nullcontext(): is_sdxl = isinstance(pipe, diffusers.StableDiffusionXLPipeline) - output_type = "latent" if is_sdxl and sdxl_refiner_model is not None else "pil" - def callback(step, timestep, latents): + output_type = "latent" if is_sdxl and sdxl_refiner_model is not None else "np" + def callback(pipe, step, timestep, latents): if future.check_cancelled(): raise InterruptedError() - future.add_response(ImageGenerationResult.step_preview(self, step_preview_mode, width, height, latents, generator, step)) + future.add_response(step_latents(pipe, step_preview_mode, latents, generator, step, steps)) try: result = pipe( prompt=prompt, @@ -98,7 +98,7 @@ def callback(step, timestep, latents): latents=None, output_type=output_type, return_dict=True, - callback=callback, + callback=functools.partial(callback, pipe), callback_steps=1, #cfg_end=optimizations.cfg_end ) @@ -111,19 +111,14 @@ def callback(step, timestep, latents): result = refiner( prompt=prompt, negative_prompt=[""], - callback=callback, + callback=functools.partial(callback, refiner), callback_steps=1, num_inference_steps=steps, - image=result.images + image=result.images, + output_type="np" ) - - future.add_response(ImageGenerationResult( - [np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in result.images], - [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()], - steps, - True - )) + + future.add_response(step_images(result.images, generator, steps, steps)) except InterruptedError: pass diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index b70d8ca8..bf66053a 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -3,9 +3,10 @@ from ...api.models.seamless_axes import SeamlessAxes import random from numpy.typing import NDArray -from ..models import Checkpoint, Optimizations, Scheduler, UpscaleTiler, ImageGenerationResult +from ..models import Checkpoint, Optimizations, Scheduler, UpscaleTiler, step_images from ..future import Future from contextlib import nullcontext +from ...image_utils import rgb, rgba def upscale( self, @@ -32,7 +33,6 @@ def upscale( future = Future() yield future - from PIL import Image, ImageOps import torch import diffusers @@ -58,15 +58,13 @@ def upscale( _configure_model_padding(pipe.unet, seamless_axes & ~tiler.seamless_axes) _configure_model_padding(pipe.vae, seamless_axes & ~tiler.seamless_axes) - if image.shape[2] == 4: - image = image[:, :, :3] for i in range(0, len(tiler), optimizations.batch_size): if future.check_cancelled(): future.set_done() return batch_size = min(len(tiler)-i, optimizations.batch_size) ids = list(range(i, i+batch_size)) - low_res_tiles = [Image.fromarray(tiler[id]).convert('RGB') for id in ids] + low_res_tiles = [rgb(tiler[id]) for id in ids] # Inference with torch.inference_mode() if device not in ('mps', "dml") else nullcontext(): high_res_tiles = pipe( @@ -75,28 +73,24 @@ def upscale( num_inference_steps=steps, generator=generator, guidance_scale=cfg_scale, + output_type="np" ).images - for id, tile in zip(ids, high_res_tiles): - tiler[id] = np.array(tile.convert('RGBA')) - step = None + tiler[id] = rgba(tile) + if step_preview_mode != StepPreviewMode.NONE: - step = Image.fromarray(tiler.combined().astype(np.uint8)) - future.add_response(ImageGenerationResult( - [(np.asarray(ImageOps.flip(step).convert('RGBA'), dtype=np.float32) / 255.)], - [seed], + future.add_response(step_images( + [tiler.combined()], + generator, i + batch_size, - (i + batch_size) == len(tiler), - total=len(tiler) + len(tiler), )) if step_preview_mode == StepPreviewMode.NONE: - final = Image.fromarray(tiler.combined().astype(np.uint8)) - future.add_response(ImageGenerationResult( - [np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255.], - [seed], + future.add_response(step_images( + [tiler.combined()], + generator, len(tiler), - True, - total=len(tiler) + len(tiler) )) future.set_done() diff --git a/generator_process/actor.py b/generator_process/actor.py index 4b76e7d7..916952b3 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -22,7 +22,8 @@ def _load_dependencies(): os.add_dll_directory(os.path.dirname(python3_path)) main_thread_rendering = False -if current_process().name == "__actor__": +is_actor_process = current_process().name == "__actor__" +if is_actor_process: _load_dependencies() elif {"-b", "-f", "-a"}.intersection(sys.argv): main_thread_rendering = True diff --git a/generator_process/models/image_generation_result.py b/generator_process/models/image_generation_result.py index bdff8bfc..a62666da 100644 --- a/generator_process/models/image_generation_result.py +++ b/generator_process/models/image_generation_result.py @@ -1,67 +1,74 @@ -from typing import List -import math -from dataclasses import dataclass -from numpy.typing import NDArray -import numpy as np from ...api.models.step_preview_mode import StepPreviewMode +from ...api.models.generation_result import GenerationResult -@dataclass -class ImageGenerationResult: - images: List[NDArray] - seeds: List[int] - step: int - final: bool - total: int | None = None - - @staticmethod - def step_preview(pipe, mode, width, height, latents, generator, iteration): - from PIL import Image, ImageOps - seeds = [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()] - match mode: - case StepPreviewMode.FAST: - return ImageGenerationResult( - [np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents[-1:]))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255.], - seeds[-1:], - iteration, - False +def step_latents(pipe, mode, latents, generator, iteration, steps): + seeds = [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()] + scale = 2 ** (len(pipe.vae.config.block_out_channels) - 1) + match mode: + case StepPreviewMode.FAST: + return [ + GenerationResult( + progress=iteration, + total=steps, + seed=seeds[-1], + image=approximate_decoded_latents(latents[-1:], scale) ) - case StepPreviewMode.FAST_BATCH: - return ImageGenerationResult( - [ - np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents[i:i + 1]))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), - dtype=np.float32) / 255. - for i in range(latents.size(0)) - ], - seeds, - iteration, - False + ] + case StepPreviewMode.FAST_BATCH: + return [ + GenerationResult( + progress=iteration, + total=steps, + seed=seed, + image=approximate_decoded_latents(latent, scale) ) - case StepPreviewMode.ACCURATE: - return ImageGenerationResult( - [np.asarray(ImageOps.flip(pipe.numpy_to_pil(pipe.decode_latents(latents[-1:]))[0]).convert('RGBA'), - dtype=np.float32) / 255.], - seeds[-1:], - iteration, - False + for latent, seed in zip(latents[:, None], seeds) + ] + case StepPreviewMode.ACCURATE: + return [ + GenerationResult( + progress=iteration, + total=steps, + seed=seeds[-1], + image=decode_latents(pipe, latents[-1:]) ) - case StepPreviewMode.ACCURATE_BATCH: - return ImageGenerationResult( - [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in pipe.numpy_to_pil(pipe.decode_latents(latents)) - ], - seeds, - iteration, - False + ] + case StepPreviewMode.ACCURATE_BATCH: + return [ + GenerationResult( + progress=iteration, + total=steps, + seed=seed, + image=decode_latents(pipe, latent) ) - return ImageGenerationResult( - [], - seeds, - iteration, - False + for latent, seed in zip(latents[:, None], seeds) + ] + return [ + GenerationResult( + progress=iteration, + total=steps, + seed=seeds[-1] + ) + ] + +def step_images(images, generator, iteration, steps): + if not isinstance(images, list) and images.ndim == 3: + images = images[None] + seeds = [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()] + return [ + GenerationResult( + progress=iteration, + total=steps, + seed=seed, + image=image ) + for image, seed in zip(images, seeds) + ] -def approximate_decoded_latents(latents): +def decode_latents(pipe, latents): + return pipe.image_processor.postprocess(pipe.vae.decode(latents / pipe.vae.config.scaling_factor).sample, output_type="np") + +def approximate_decoded_latents(latents, scale=None): """ Approximate the decoded latents without using the VAE. """ @@ -79,9 +86,9 @@ def approximate_decoded_latents(latents): ], dtype=latents.dtype, device=latents.device) latent_image = latents[0].permute(1, 2, 0) @ v1_5_latent_rgb_factors - latents_ubyte = (((latent_image + 1) / 2) - .clamp(0, 1) # change scale from -1..1 to 0..1 - .mul(0xFF) # to 0..255 - .byte()).cpu() - - return latents_ubyte.numpy() \ No newline at end of file + if scale is not None: + latent_image = torch.nn.functional.interpolate( + latent_image.permute(2, 0, 1).unsqueeze(0), scale_factor=scale, mode="nearest" + ).squeeze(0).permute(1, 2, 0) + latent_image = ((latent_image + 1) / 2).clamp(0, 1).cpu() + return latent_image.numpy() diff --git a/image_utils.py b/image_utils.py new file mode 100644 index 00000000..852c9f0a --- /dev/null +++ b/image_utils.py @@ -0,0 +1,863 @@ +import importlib.util +import os +import sys +from os import PathLike +from typing import Tuple, Literal, Union, TYPE_CHECKING + +import numpy as np +from numpy.typing import NDArray, DTypeLike + +from .generator_process import RunInSubprocess + + +""" +This module allows for simple handling of image data in numpy ndarrays in some common formats. + +Dimensions: + 2: HW - L + 3: HWC - L/LA/RGB/RGBA + 4: NHWC - batched HWC + +Channels: + 1: L + 2: LA + 3: RGB + 4: RGBA +""" + + +def version_str(version): + return ".".join(str(x) for x in version) + + +# find_spec("bpy") will never return None +has_bpy = sys.modules.get("bpy", None) is not None +has_ocio = importlib.util.find_spec("PyOpenColorIO") is not None +has_oiio = importlib.util.find_spec("OpenImageIO") is not None +has_pil = importlib.util.find_spec("PIL") is not None + +if has_bpy: + # frontend + import bpy + BLENDER_VERSION = bpy.app.version + OCIO_CONFIG = os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio') + # Easier to share via environment variables than to enforce backends with subprocesses to use their own methods of sharing. + os.environ["BLENDER_VERSION"] = version_str(BLENDER_VERSION) + os.environ["BLENDER_OCIO_CONFIG"] = OCIO_CONFIG +else: + # backend + BLENDER_VERSION = tuple(int(x) for x in os.environ["BLENDER_VERSION"].split(".")) + OCIO_CONFIG = os.environ["BLENDER_OCIO_CONFIG"] + +if TYPE_CHECKING: + import bpy + import PIL.Image + + +def _bpy_version_error(required_version, feature, module): + if BLENDER_VERSION >= required_version: + return Exception(f"{module} is unexpectedly missing in Blender {version_str(BLENDER_VERSION)}") + return Exception(f"{feature} requires Blender {version_str(required_version)} or higher, you are using {version_str(BLENDER_VERSION)}") + + +def size(array: NDArray) -> Tuple[int, int]: + if array.ndim == 2: + return array.shape[1], array.shape[0] + if array.ndim in [3, 4]: + return array.shape[-2], array.shape[-3] + raise ValueError(f"Can't determine size from {array.ndim} dimensions") + + +def channels(array: NDArray) -> int: + if array.ndim == 2: + return 1 + if array.ndim in [3, 4]: + return array.shape[-1] + raise ValueError(f"Can't determine channels from {array.ndim} dimensions") + + +def ensure_alpha(array: NDArray, alpha=None) -> NDArray: + """ + Args: + array: Image pixels values. + alpha: Default alpha value if an alpha channel will be made. Will be inferred from `array.dtype` if None. + + Returns: The converted image or the original image if it already had alpha. + """ + c = channels(array) + if c in [2, 4]: + return array + if c not in [1, 3]: + raise ValueError(f"Can't ensure alpha from {c} channels") + + if alpha is None: + alpha = 0 + if np.issubdtype(array.dtype, np.floating): + alpha = 1 + elif np.issubdtype(array.dtype, np.integer): + alpha = np.iinfo(array.dtype).max + array = ensure_channel_dim(array) + return np.pad(array, [*[(0, 0)]*(array.ndim-1), (0, 1)], constant_values=alpha) + + +def ensure_opaque(array: NDArray) -> NDArray: + """ + Removes the alpha channel if it exists. + """ + if channels(array) in [2, 4]: + return array[..., :-1] + return array + + +def ensure_channel_dim(array: NDArray) -> NDArray: + """ + Expands a HW grayscale image to HWC. + """ + if array.ndim == 2: + return array[..., np.newaxis] + return array + + +def rgb(array: NDArray) -> NDArray: + """ + Converts a grayscale image to RGB or removes the alpha channel from an RGBA image. + If the image was already RGB the original array will be returned. + """ + c = channels(array) + match channels(array): + case 1: + return np.concatenate([ensure_channel_dim(array)] * 3, axis=-1) + case 2: + return np.concatenate([array[..., :1]] * 3, axis=-1) + case 3: + return array + case 4: + return array[..., :3] + raise ValueError(f"Can't make {c} channels RGB") + + +def rgba(array: NDArray, alpha=None) -> NDArray: + """ + Args: + array: Image pixels values. + alpha: Default alpha value if an alpha channel will be made. Will be inferred from `array.dtype` if None. + + Returns: The converted image or the original image if it already was RGBA. + """ + c = channels(array) + if c == 4: + return array + if c == 2: + l, a = np.split(array, 2, axis=-1) + return np.concatenate([l, l, l, a], axis=-1) + return ensure_alpha(rgb(array), alpha) + + +def grayscale(array: NDArray) -> NDArray: + """ + Converts `array` into HW or NHWC grayscale. This is intended for converting an + RGB image that is already visibly grayscale, such as a depth map. It will not + make a good approximation of perceived lightness of an otherwise colored image. + """ + if array.ndim == 2: + return array + c = channels(array) + if array.ndim == 3: + if c in [1, 2]: + return array[..., 0] + elif c in [3, 4]: + return np.max(array[..., :3], axis=-1) + raise ValueError(f"Can't make {c} channels grayscale") + elif array.ndim == 4: + if c in [1, 2]: + return array[..., :1] + elif c in [3, 4]: + return np.max(array[..., :3], axis=-1, keepdims=True) + raise ValueError(f"Can't make {c} channels grayscale") + raise ValueError(f"Can't make {array.ndim} dimensions grayscale") + + +def _passthrough_alpha(from_array, to_array): + if channels(from_array) not in [2, 4]: + return to_array + to_array = np.concatenate([ensure_channel_dim(to_array), from_array[..., -1:]], axis=-1) + return to_array + + +def linear_to_srgb(array: NDArray, clamp=True) -> NDArray: + """ + Args: + array: Image to convert from linear to sRGB color space. Will be converted to float32 if it isn't already a float dtype. + clamp: whether to restrict the result between 0..1 + """ + if not np.issubdtype(array.dtype, np.floating): + array = to_dtype(array, np.float32) + srgb = ensure_opaque(array) + srgb = np.where( + srgb <= 0.0031308, + srgb * 12.92, + (np.abs(srgb) ** (1/2.4) * 1.055) - 0.055 + # abs() to suppress `RuntimeWarning: invalid value encountered in power` for negative values + ) + if clamp: + # conversion may produce values outside standard range, usually >1 + srgb = np.clip(srgb, 0, 1) + srgb = _passthrough_alpha(array, srgb) + return srgb + + +def srgb_to_linear(array: NDArray) -> NDArray: + """ + Converts from sRGB to linear color space. Will be converted to float32 if it isn't already a float dtype. + """ + if not np.issubdtype(array.dtype, np.floating): + array = to_dtype(array, np.float32) + linear = ensure_opaque(array) + linear = np.where( + linear <= 0.04045, + linear / 12.92, + ((linear + 0.055) / 1.055) ** 2.4 + ) + linear = _passthrough_alpha(array, linear) + return linear + + +@RunInSubprocess.when_raised +def color_transform(array: NDArray, from_color_space: str, to_color_space: str, *, clamp_srgb=True) -> NDArray: + """ + Args: + array: Pixel values in `from_color_space` + from_color_space: Color space of `array` + to_color_space: Desired color space + clamp_srgb: Restrict values inside the standard range when converting to sRGB. + + Returns: Pixel values in `to_color_space`. The image will be converted to RGB/RGBA float32 for most transforms. + Transforms between linear and sRGB may remain grayscale and keep the original DType if it was floating point. + """ + # Blender handles Raw and Non-Color images as if they were in Linear color space. + if from_color_space in ["Raw", "Non-Color"]: + from_color_space = "Linear" + if to_color_space in ["Raw", "Non-Color"]: + to_color_space = "Linear" + + if from_color_space == to_color_space: + return array + elif from_color_space == "Linear" and to_color_space == "sRGB": + return linear_to_srgb(array, clamp_srgb) + elif from_color_space == "sRGB" and to_color_space == "Linear": + return srgb_to_linear(array) + + if not has_ocio: + raise RunInSubprocess + + import PyOpenColorIO as OCIO + config = OCIO.Config.CreateFromFile(OCIO_CONFIG) + proc = config.getProcessor(from_color_space, to_color_space).getDefaultCPUProcessor() + # OCIO requires RGB/RGBA float32. + # There is a channel agnostic apply(), but I can't seem to get it to work. + # getOptimizedCPUProcessor() can handle different precisions, but I doubt it would have meaningful use. + array = to_dtype(array, np.float32) + c = channels(array) + if c in [1, 3]: + array = rgb(array) + proc.applyRGB(array) + if clamp_srgb and to_color_space == "sRGB": + array = np.clip(array, 0, 1) + return array + elif c in [2, 4]: + array = rgba(array) + proc.applyRGBA(array) + if clamp_srgb and to_color_space == "sRGB": + array = np.clip(array, 0, 1) + return array + raise ValueError(f"Can't color transform {c} channels") + + +# inverse=True is often crashing from EXCEPTION_ACCESS_VIOLATION while on frontend. +# Normally this is caused by not running on the main thread or accessing a deleted +# object, neither seem to be the issue here. Doesn't matter if the backend imports +# its own OCIO or the one packaged with Blender. +# Stack trace: +# OpenColorIO_2_2.dll :0x00007FFDE8961160 OpenColorIO_v2_2::GradingTone::validate +# OpenColorIO_2_2.dll :0x00007FFDE8A2BD40 OpenColorIO_v2_2::Processor::isNoOp +# OpenColorIO_2_2.dll :0x00007FFDE882EA00 OpenColorIO_v2_2::CPUProcessor::apply +# PyOpenColorIO.pyd :0x00007FFDEB0F0E40 pybind11::error_already_set::what +# PyOpenColorIO.pyd :0x00007FFDEB0F0E40 pybind11::error_already_set::what +# PyOpenColorIO.pyd :0x00007FFDEB0F0E40 pybind11::error_already_set::what +# PyOpenColorIO.pyd :0x00007FFDEB0E7510 pybind11::error_already_set::discard_as_unraisable +@RunInSubprocess.when(lambda *_, inverse=False, **__: inverse or not has_ocio) +def render_color_transform( + array: NDArray, + exposure: float, + gamma: float, + view_transform: str, + display_device: str, + look: str, + *, + inverse: bool = False, + color_space: str | None = None, + clamp_srgb: bool = True, +) -> NDArray: + import PyOpenColorIO as OCIO + + ocio_config = OCIO.Config.CreateFromFile(OCIO_CONFIG) + + # A reimplementation of `OCIOImpl::createDisplayProcessor` from the Blender source. + # https://github.com/blender/blender/blob/3816fcd8611bc2836ee8b2a5225b378a02141ce4/intern/opencolorio/ocio_impl.cc#L666 + # Modified to support a final color space transform. + def create_display_processor( + config, + input_colorspace, + view, + display, + look, + scale, # Exposure + exponent, # Gamma + inverse, + color_space + ): + group = OCIO.GroupTransform() + + # Exposure + if scale != 1: + # Always apply exposure in scene linear. + color_space_transform = OCIO.ColorSpaceTransform() + color_space_transform.setSrc(input_colorspace) + color_space_transform.setDst(OCIO.ROLE_SCENE_LINEAR) + group.appendTransform(color_space_transform) + + # Make further transforms aware of the color space change + input_colorspace = OCIO.ROLE_SCENE_LINEAR + + # Apply scale + matrix_transform = OCIO.MatrixTransform( + [scale, 0.0, 0.0, 0.0, 0.0, scale, 0.0, 0.0, 0.0, 0.0, scale, 0.0, 0.0, 0.0, 0.0, 1.0]) + group.appendTransform(matrix_transform) + + # Add look transform + use_look = look is not None and len(look) > 0 + if use_look: + look_output = config.getLook(look).getProcessSpace() + if look_output is not None and len(look_output) > 0: + look_transform = OCIO.LookTransform() + look_transform.setSrc(input_colorspace) + look_transform.setDst(look_output) + look_transform.setLooks(look) + group.appendTransform(look_transform) + # Make further transforms aware of the color space change. + input_colorspace = look_output + else: + # For empty looks, no output color space is returned. + use_look = False + + # Add view and display transform + display_view_transform = OCIO.DisplayViewTransform() + display_view_transform.setSrc(input_colorspace) + display_view_transform.setLooksBypass(use_look) + display_view_transform.setView(view) + display_view_transform.setDisplay(display) + group.appendTransform(display_view_transform) + + if color_space is not None: + group.appendTransform(OCIO.ColorSpaceTransform(input_colorspace if display == "None" else display, color_space)) + + # Gamma + if exponent != 1: + exponent_transform = OCIO.ExponentTransform([exponent, exponent, exponent, 1.0]) + group.appendTransform(exponent_transform) + + if inverse: + group.setDirection(OCIO.TransformDirection.TRANSFORM_DIR_INVERSE) + + # Create processor from transform. This is the moment were OCIO validates + # the entire transform, no need to check for the validity of inputs above. + return config.getProcessor(group) + + # Exposure and gamma transformations derived from Blender source: + # https://github.com/blender/blender/blob/3816fcd8611bc2836ee8b2a5225b378a02141ce4/source/blender/imbuf/intern/colormanagement.cc#L867 + scale = 2 ** exposure + exponent = 1 / max(gamma, np.finfo(np.float32).eps) + processor = create_display_processor(ocio_config, OCIO.ROLE_SCENE_LINEAR, view_transform, display_device, look if look != 'None' else None, scale, exponent, inverse, color_space) + array = to_dtype(array, np.float32) + c = channels(array) + if c in [1, 3]: + array = rgb(array) + processor.getDefaultCPUProcessor().applyRGB(array) + elif c in [2, 4]: + array = rgba(array) + processor.getDefaultCPUProcessor().applyRGBA(array) + else: + raise ValueError(f"Can't color transform {c} channels") + if clamp_srgb and (color_space == "sRGB" or (display_device == "sRGB" and color_space is None)) and not inverse: + array = np.clip(array, 0, 1) + return array + + +def scene_color_transform(array: NDArray, scene: Union["bpy.types.Scene", None] = None, *, inverse: bool = False, color_space: str | None = None, clamp_srgb=True) -> NDArray: + if scene is None: + import bpy + scene = bpy.context.scene + view = scene.view_settings + display = scene.display_settings.display_device + return render_color_transform( + array, + view.exposure, + view.gamma, + view.view_transform, + display, + view.look, + inverse=inverse, + clamp_srgb=clamp_srgb, + color_space=color_space + ) + + +def _unsigned(dtype: DTypeLike) -> DTypeLike: + match bits := np.iinfo(dtype).bits: + case 8: + return np.uint8 + case 16: + return np.uint16 + case 32: + return np.uint32 + case 64: + return np.uint64 + raise ValueError(f"unexpected bit depth {bits} from {repr(dtype)}") + + +def to_dtype(array: NDArray, dtype: DTypeLike) -> NDArray: + """ + Remaps values with respect to ranges rather than simply casting for integer DTypes. + `integer(0)=float(0)`, `integer.MAX=float(1)`, and signed `integer.MIN+1=float(-1)` + """ + dtype = np.dtype(dtype) + from_dtype = array.dtype + if dtype == from_dtype: + return array + from_floating = np.issubdtype(from_dtype, np.floating) + from_integer = np.issubdtype(from_dtype, np.integer) + to_floating = np.issubdtype(dtype, np.floating) + to_integer = np.issubdtype(dtype, np.integer) + if from_floating and to_floating: + array = array.astype(dtype) + if np.finfo(from_dtype).bits > np.finfo(dtype).bits: + # prevent inf when lowering precision + array = np.nan_to_num(array) + elif from_floating and to_integer: + iinfo = np.iinfo(dtype) + array = (array.clip(-1 if iinfo.min < 0 else 0, 1) * iinfo.max).round().astype(dtype) + elif from_integer and to_floating: + iinfo = np.iinfo(from_dtype) + array = (array / iinfo.max).astype(dtype) + elif from_integer and to_integer: + from_signed = np.issubdtype(from_dtype, np.signedinteger) + to_signed = np.issubdtype(dtype, np.signedinteger) + from_bits = np.iinfo(from_dtype).bits + to_bits = np.iinfo(dtype).bits + if from_signed: + from_bits -= 1 + if to_signed: + to_bits -= 1 + bit_diff = to_bits - from_bits + + if from_signed and not to_signed: + # unsigned output does not support negative + array = np.maximum(array, 0) + if from_signed and to_signed: + # simpler to handle bit manipulation in unsigned + sign = np.sign(array) + array = np.abs(array) + + if bit_diff > 0: + # Repeat bits rather than using a single left shift + # so that from_iinfo.max turns into to_iinfo.max + # and all values remain equally spaced. + # Example 8 to 16 bits: + # (incorrect) 0x00FF << 8 = 0xFF00 + # (correct) 0x00FF << 8 | 0x00FF = 0xFFFF + # Implementation uses multiplication instead of potentially multiple left shifts and ors: + # 0x00FF * 0x0101 = 0xFFFF + base = array.astype(_unsigned(dtype)) + m = 0 + for i in range(bit_diff, -1, -from_bits): + m += 2 ** i + array = base * m + remaining_bits = bit_diff % from_bits + if remaining_bits > 0: + # when changing between signed and unsigned bit_diff is not a multiple of from_bits + array |= base >> (from_bits-remaining_bits) + elif bit_diff < 0: + array = array.astype(_unsigned(from_dtype), copy=False) >> -bit_diff + + if from_signed and to_signed: + array = np.multiply(array, sign, dtype=dtype) + array = array.astype(dtype, copy=False) + else: + raise TypeError(f"Unable to convert from {array.dtype} to {dtype}") + return array + + +@RunInSubprocess.when(not has_oiio) +def resize(array: NDArray, size: Tuple[int, int], clamp=True): + no_channels = array.ndim == 2 + if no_channels: + array = array[..., np.newaxis] + no_batch = array.ndim < 4 + if no_batch: + array = array[np.newaxis, ...] + if clamp: + c_min = np.min(array, axis=(1, 2), keepdims=True) + c_max = np.max(array, axis=(1, 2), keepdims=True) + + if has_oiio: + import OpenImageIO as oiio + resized = [] + for unbatched in array: + # OpenImageIO can have batched images, but doesn't support resizing them + image_in = oiio.ImageBuf(unbatched) + image_out = oiio.ImageBufAlgo.resize(image_in, roi=oiio.ROI(0, int(size[0]), 0, int(size[1]))) + if image_out.has_error: + raise Exception(image_out.geterror()) + resized.append(image_out.get_pixels(image_in.spec().format)) + array = np.stack(resized) + else: + original_dtype = array.dtype + if np.issubdtype(original_dtype, np.floating): + if original_dtype == np.float16: + # interpolation not implemented for float16 on CPU + array = to_dtype(array, np.float32) + elif np.issubdtype(original_dtype, np.integer): + # integer interpolation only supported for uint8 nearest, nearest-exact or bilinear + bits = np.iinfo(original_dtype).bits + array = to_dtype(array, np.float64 if bits >= 32 else np.float32) + + import torch + array = torch.from_numpy(np.transpose(array, (0, 3, 1, 2))) + array = torch.nn.functional.interpolate(array, size=(size[1], size[0]), mode="bilinear") + array = np.transpose(array, (0, 2, 3, 1)).numpy() + array = to_dtype(array, original_dtype) + + if clamp: + array = np.clip(array, c_min, c_max) + if no_batch: + array = np.squeeze(array, 0) + if no_channels: + array = np.squeeze(array, -1) + return array + + +def bpy_to_np(image: "bpy.types.Image", *, color_space: str | None = "sRGB", clamp_srgb=True, top_to_bottom=True) -> NDArray: + """ + Args: + image: Image to extract pixels values from. + color_space: The color space to convert to. `None` will apply no color transform. + Keep in mind that Raw/Non-Color images are handled as if they were in Linear color space. + clamp_srgb: Restrict values inside the standard range when converting to sRGB. + top_to_bottom: The y-axis is flipped to a more common standard of `top=0` to `bottom=height-1`. + + Returns: A ndarray copy of `image.pixels` in RGBA float32 format. + """ + if image.type == "RENDER_RESULT": + # can't get pixels automatically without rendering again and freezing Blender until it finishes, or saving to disk + raise ValueError(f"{image.name} image can't be used directly, alternatively use a compositor viewer node") + array = np.empty((image.size[1], image.size[0], image.channels), dtype=np.float32) + # foreach_get/set is extremely fast to read/write an entire image compared to alternatives + # see https://projects.blender.org/blender/blender/commit/9075ec8269e7cb029f4fab6c1289eb2f1ae2858a + image.pixels.foreach_get(array.ravel()) + if color_space is not None: + if image.type == "COMPOSITING": + # Viewer Node + array = scene_color_transform(array, color_space=color_space, clamp_srgb=clamp_srgb) + else: + array = color_transform(array, image.colorspace_settings.name, color_space, clamp_srgb=clamp_srgb) + if top_to_bottom: + array = np.flipud(array) + return rgba(array) + + +def np_to_bpy(array: NDArray, name=None, existing_image=None, float_buffer=None, color_space: str = "sRGB", top_to_bottom=True) -> "bpy.types.Image": + """ + Args: + array: Image pixel values. The y-axis is expected to be ordered `top=0` to `bottom=height-1`. + name: Name of the image data-block. If None it will be `existing_image.name` or "Untitled". + existing_image: Image data-block to overwrite. + float_buffer: + Make Blender keep data in (`True`) 32-bit float values, or (`False`) 8-bit integer values. + `None` won't invalidate `existing_image`, but if a new image is created it will be `False`. + color_space: Color space of `array`. + + Returns: A new Blender image or `existing_image` if it didn't require replacement. + """ + if array.ndim == 4 and array.shape[0] > 1: + raise ValueError(f"Can't convert a batched array of {array.shape[0]} images to a Blender image") + + # create or replace image + import bpy + width, height = size(array) + if name is None: + name = "Untitled" if existing_image is None else existing_image.name + if existing_image is not None and existing_image.type in ["RENDER_RESULT", "COMPOSITING"]: + existing_image = None + elif existing_image is not None and ( + existing_image.size[0] != width + or existing_image.size[1] != height + or (existing_image.channels != channels(array) and existing_image.channels != 4) + or (existing_image.is_float != float_buffer and float_buffer is not None) + ): + bpy.data.images.remove(existing_image) + existing_image = None + if existing_image is None: + image = bpy.data.images.new( + name, + width=width, + height=height, + alpha=channels(array) == 4, + float_buffer=False if float_buffer is None else float_buffer + ) + else: + image = existing_image + image.name = name + image.colorspace_settings.name = color_space + + # adjust array pixels to fit into image + if array.ndim == 4: + array = array[0] + if top_to_bottom: + array = np.flipud(array) + array = to_dtype(array, np.float32) + if image.channels == 4: + array = rgba(array) + elif image.channels == 3: + # I believe image.channels only exists for backwards compatibility and modern versions of Blender + # will always handle images as RGBA. I can't manage to make or import an image and end up with + # anything but 4 channels. Support for images with 3 channels will be kept just in case. + array = rgb(array) + else: + raise NotImplementedError(f"Blender image unexpectedly has {image.channels} channels") + + # apply pixels to image + image.pixels.foreach_set(array.ravel()) + image.pack() + image.update() + return image + + +def render_pass_to_np( + render_pass: "bpy.types.RenderPass", + size: Tuple[int, int], + *, + color_management: bool = False, + color_space: str | None = None, + clamp_srgb: bool = True, + top_to_bottom: bool = True +): + array = np.empty((*reversed(size), render_pass.channels), dtype=np.float32) + render_pass.rect.foreach_get(array.reshape((-1, render_pass.channels))) + if color_management: + array = scene_color_transform(array, color_space=color_space, clamp_srgb=clamp_srgb) + elif color_space is not None: + array = color_transform(array, "Linear", color_space, clamp_srgb=clamp_srgb) + if top_to_bottom: + array = np.flipud(array) + return array + + +def np_to_render_pass( + array: NDArray, + render_pass: "bpy.types.RenderPass", + *, + inverse_color_management: bool = False, + color_space: str | None = None, + dtype: DTypeLike = np.float32, + top_to_bottom: bool = True +): + if inverse_color_management: + array = scene_color_transform(array, inverse=True, color_space=color_space) + elif color_space is not None: + array = color_transform(color_space, "Linear") + if channels(array) != render_pass.channels: + match render_pass.channels: + case 1: + array = grayscale(array) + case 3: + array = rgb(array) + case 4: + array = rgba(array) + case _: + raise NotImplementedError(f"Render pass {render_pass.name} unexpectedly requires {render_pass.channels} channels") + if dtype is not None: + array = to_dtype(array, dtype) + if top_to_bottom: + array = np.flipud(array) + render_pass.rect.foreach_set(array.reshape(-1, render_pass.channels)) + + +def _mode(array, mode): + if mode is None: + return array + elif mode == "RGBA": + return rgba(array) + elif mode == "RGB": + return rgb(array) + elif mode == "L": + return grayscale(array) + elif mode == "LA": + return ensure_alpha(_passthrough_alpha(array, grayscale(array))) + raise ValueError(f"mode expected one of {['RGB', 'RGBA', 'L', 'LA', None]}, got {repr(mode)}") + + +def pil_to_np(image, *, dtype: DTypeLike | None = np.float32, mode: Literal["RGB", "RGBA", "L", "LA"] | None = None) -> NDArray: + # some modes don't require being converted to RGBA for proper handling in other module functions + # see for other modes https://pillow.readthedocs.io/en/stable/handbook/concepts.html#concept-modes + if image.mode not in ["RGB", "RGBA", "L", "LA", "I", "F", "I;16"]: + image = image.convert("RGBA") + array = np.array(image) + if dtype is not None: + array = to_dtype(array, dtype) + array = _mode(array, mode) + return array + + +def np_to_pil(array: NDArray, *, mode: Literal["RGB", "RGBA", "L", "LA"] | None = None): + from PIL import Image + array = to_dtype(array, np.uint8) + if mode is None: + if channels(array) == 1 and array.ndim == 3: + # PIL L mode can't have a channel dimension + array = array[..., 1] + else: + array = _mode(array, mode) + # PIL does support higher precision modes for a single channel, but I don't see a need for supporting them yet. + # uint16="I;16", int32="I", float32="F" + return Image.fromarray(array, mode=mode) + + +def _dtype_to_type_desc(dtype): + import OpenImageIO as oiio + dtype = np.dtype(dtype) + match dtype: + case np.uint8: + return oiio.TypeUInt8 + case np.uint16: + return oiio.TypeUInt16 + case np.uint32: + return oiio.TypeUInt32 + case np.uint64: + return oiio.TypeUInt64 + case np.int8: + return oiio.TypeInt8 + case np.int16: + return oiio.TypeInt16 + case np.int32: + return oiio.TypeInt32 + case np.int64: + return oiio.TypeInt64 + case np.float16: + return oiio.TypeHalf + case np.float32: + return oiio.TypeFloat + case np.float64: + # no oiio.TypeDouble + return oiio.TypeDesc(oiio.BASETYPE.DOUBLE) + raise TypeError(f"can't convert {dtype} to OpenImageIO.TypeDesc") + + +@RunInSubprocess.when(not has_oiio) +def path_to_np( + path: str | PathLike, + *, + dtype: DTypeLike | None = np.float32, + default_color_space: str | None = None, + to_color_space: str | None = "sRGB" +) -> NDArray: + """ + Args: + path: Path to an image file. + dtype: Data type of the returned array. `None` won't change the data type. The data type may still change if a color transform occurs. + default_color_space: The color space that `image_or_path` will be handled as when it can't be determined automatically. + to_color_space: Color space of the returned array. `None` won't apply a color transform. + """ + if has_oiio: + import OpenImageIO as oiio + image = oiio.ImageInput.open(str(path)) + if image is None: + raise IOError(oiio.geterror()) + type_desc = image.spec().format + if dtype is not None: + type_desc = _dtype_to_type_desc(dtype) + array = image.read_image(type_desc) + from_color_space = image.spec().get_string_attribute("oiio:ColorSpace", default_color_space) + image.close() + else: + from PIL import Image + array = pil_to_np(Image.open(path)) + if dtype is not None: + array = to_dtype(array, dtype) + from_color_space = "sRGB" + if from_color_space is not None and to_color_space is not None: + array = color_transform(array, from_color_space, to_color_space) + return array + + +ImageOrPath = Union[NDArray, "PIL.Image.Image", str, PathLike] +"""Backend compatible image types""" + + +def image_to_np( + image_or_path: ImageOrPath | "bpy.types.Image" | None, + *, + dtype: DTypeLike | None = np.float32, + mode: Literal["RGB", "RGBA", "L", "LA"] | None = "RGBA", + default_color_space: str | None = None, + to_color_space: str | None = "sRGB", + size: Tuple[int, int] | None = None, + top_to_bottom: bool = True +) -> NDArray: + """ + Opens an image from disk or takes an image object and converts it to `numpy.ndarray`. + Usable for image argument sanitization when the source can vary in type or format. + + Args: + image_or_path: Either a file path or an instance of `bpy.types.Image`, `PIL.Image.Image`, or `numpy.ndarray`. `None` will return `None`. + dtype: Data type of the returned array. `None` won't change the data type. The data type may still change if a color transform occurs. + mode: Channel mode of the returned array. `None` won't change the mode. The mode may still change if a color transform occurs. + default_color_space: The color space that `image_or_path` will be handled as when it can't be determined automatically. + to_color_space: Color space of the returned array. `None` won't apply a color transform. + size: Resize to specific dimensions. `None` won't change the size. + top_to_bottom: Flips the image like `bpy_to_np(top_to_bottom=True)` does when `True` and `image_or_path` is a Blender image. Other image sources will only be flipped when `False`. + """ + + if image_or_path is None: + return None + + # convert image_or_path to numpy.ndarray + match image_or_path: + case PathLike() | str(): + array = path_to_np(image_or_path, dtype=dtype, default_color_space=default_color_space, to_color_space=to_color_space) + from_color_space = None + case object(__module__="PIL.Image", __class__=type(__name__="Image")): + # abnormal class check because PIL cannot be imported on frontend + array = pil_to_np(image_or_path) + from_color_space = "sRGB" + case object(__module__="bpy.types", __class__=type(__name__="Image")): + # abnormal class check because bpy cannot be imported on backend + array = bpy_to_np(image_or_path, color_space=to_color_space) + from_color_space = None + case np.ndarray(): + array = image_or_path + from_color_space = default_color_space + case _: + raise TypeError(f"not an image or path {repr(type(image_or_path))}") + + # apply image requirements + if not top_to_bottom: + array = np.flipud(array) + if from_color_space is not None and to_color_space is not None: + array = color_transform(array, from_color_space, to_color_space) + if dtype is not None: + array = to_dtype(array, dtype) + array = _mode(array, mode) + if size is not None: + array = resize(array, size) + + return array diff --git a/operators/dream_texture.py b/operators/dream_texture.py index aec1a2a5..a1649482 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -4,27 +4,13 @@ from typing import List, Literal from .notify_result import NotifyResult -from ..pil_to_image import * from ..prompt_engineering import * from ..generator_process import Generator from .. import api +from .. import image_utils import time import math -def bpy_image(name, width, height, pixels, existing_image): - if existing_image is not None and (existing_image.size[0] != width or existing_image.size[1] != height): - bpy.data.images.remove(existing_image) - existing_image = None - if existing_image is None: - image = bpy.data.images.new(name, width=width, height=height) - else: - image = existing_image - image.name = name - image.pixels.foreach_set(pixels) - image.pack() - image.update() - return image - def get_source_image(context, source: Literal['file', 'open_editor']): match source: case 'file': @@ -95,19 +81,15 @@ def execute(self, context): except ValueError: init_image = None if init_image is not None: - init_image = np.flipud( - (np.array(init_image.pixels) * 255) - .astype(np.uint8) - .reshape((init_image.size[1], init_image.size[0], init_image.channels)) - ) + init_image_color_space = "sRGB" + if scene.dream_textures_prompt.use_init_img and scene.dream_textures_prompt.modify_action_source_type in ['depth_map', 'depth']: + init_image_color_space = None + init_image = image_utils.bpy_to_np(init_image, color_space=init_image_color_space) control_images = None if len(prompt.control_nets) > 0: control_images = [ - np.flipud( - np.array(net.control_image.pixels) - .reshape((net.control_image.size[1], net.control_image.size[0], net.control_image.channels)) - ) + image_utils.bpy_to_np(net.control_image, color_space=None) for net in prompt.control_nets ] @@ -125,7 +107,7 @@ def step_callback(progress: List[api.GenerationResult]) -> bool: image = api.GenerationResult.tile_images(progress) if image is None: return CancelGenerator.should_continue - last_data_block = bpy_image(f"Step {progress[-1].progress}/{progress[-1].total}", image.shape[1], image.shape[0], image.ravel(), last_data_block) + last_data_block = image_utils.np_to_bpy(image, f"Step {progress[-1].progress}/{progress[-1].total}", last_data_block) for area in screen.areas: if area.type == 'IMAGE_EDITOR' and not area.spaces.active.use_image_pin: area.spaces.active.image = last_data_block @@ -161,7 +143,7 @@ def callback(results: List[api.GenerationResult] | Exception): seed_str_length = len(str(result.seed)) trim_aware_name = (prompt_string[:54 - seed_str_length] + '..') if len(prompt_string) > 54 else prompt_string name_with_trimmed_prompt = f"{trim_aware_name} ({result.seed})" - image = bpy_image(name_with_trimmed_prompt, result.image.shape[1], result.image.shape[0], result.image.ravel(), last_data_block) + image = image_utils.np_to_bpy(result.image, name_with_trimmed_prompt, last_data_block) last_data_block = None if node_tree is not None: nodes = node_tree.nodes diff --git a/operators/project.py b/operators/project.py index 3a647116..7180763d 100644 --- a/operators/project.py +++ b/operators/project.py @@ -23,6 +23,7 @@ from ..engine.annotations.depth import render_depth_map from .. import api +from .. import image_utils framebuffer_arguments = [ ('depth', 'Depth', 'Only provide the scene depth as input'), @@ -357,7 +358,7 @@ def vert_to_uv(v): context.scene.dream_textures_info = "Rendering viewport depth..." - depth = render_depth_map( + depth = np.flipud(render_depth_map( context.evaluated_depsgraph_get(), collection=None, width=region_width, @@ -365,7 +366,7 @@ def vert_to_uv(v): matrix=context.space_data.region_3d.view_matrix, projection_matrix=context.space_data.region_3d.window_matrix, main_thread=True - ) + )) texture = None @@ -373,11 +374,7 @@ def step_callback(progress: List[api.GenerationResult]) -> bool: nonlocal texture context.scene.dream_textures_progress = progress[-1].progress image = api.GenerationResult.tile_images(progress) - if texture is None: - texture = bpy.data.images.new(name="Step", width=image.shape[1], height=image.shape[0]) - texture.name = f"Step {progress[-1].progress}/{progress[-1].total}" - texture.pixels[:] = image.ravel() - texture.update() + texture = image_utils.np_to_bpy(image, f"Step {progress[-1].progress}/{progress[-1].total}", texture) image_texture_node.image = texture return CancelGenerator.should_continue @@ -399,13 +396,7 @@ def callback(results: List[api.GenerationResult] | Exception): trim_aware_name = (prompt_subject[:54 - seed_str_length] + '..') if len(prompt_subject) > 54 else prompt_subject name_with_trimmed_prompt = f"{trim_aware_name} ({result.seed})" - if texture is None: - texture = bpy.data.images.new(name=name_with_trimmed_prompt, width=result.image.shape[1], height=result.image.shape[0]) - texture.name = name_with_trimmed_prompt - material.name = name_with_trimmed_prompt - texture.pixels[:] = result.image.ravel() - texture.update() - texture.pack() + texture = image_utils.np_to_bpy(result.image, name_with_trimmed_prompt, texture) image_texture_node.image = texture if context.scene.dream_textures_project_bake: for bm, src_uv_layer in target_objects: @@ -430,7 +421,7 @@ def callback(results: List[api.GenerationResult] | Exception): image_data = bpy.data.images.load(init_img_path) if init_img_path is not None else None image = np.asarray(image_data.pixels).reshape((*depth.shape, image_data.channels)) if image_data is not None else None if context.scene.dream_textures_project_use_control_net: - generated_args: api.GenerationArguments = context.scene.dream_textures_project_prompt.generate_args(context, init_image=image, control_images=[np.flipud(depth)]) + generated_args: api.GenerationArguments = context.scene.dream_textures_project_prompt.generate_args(context, init_image=image, control_images=[image_utils.rgba(depth)]) backend.generate(generated_args, step_callback=step_callback, callback=callback) else: generated_args: api.GenerationArguments = context.scene.dream_textures_project_prompt.generate_args(context) diff --git a/operators/upscale.py b/operators/upscale.py index be15f859..29f3621d 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -5,6 +5,7 @@ from ..prompt_engineering import custom_structure from ..generator_process import Generator from .dream_texture import CancelGenerator +from .. import image_utils upscale_options = [ ("2", "2x", "", 2), @@ -12,20 +13,6 @@ ("8", "8x", "", 8), ] -def bpy_image(name, width, height, pixels, existing_image): - if existing_image is not None and (existing_image.size[0] != width or existing_image.size[1] != height): - bpy.data.images.remove(existing_image) - existing_image = None - if existing_image is None: - image = bpy.data.images.new(name, width=width, height=height) - else: - image = existing_image - image.name = name - image.pixels.foreach_set(pixels) - image.pack() - image.update() - return image - def get_source_image(context): node_tree = context.material.node_tree if hasattr(context, 'material') else None active_node = next((node for node in node_tree.nodes if node.select and node.bl_idname == 'ShaderNodeTexImage'), None) if node_tree is not None else None @@ -70,11 +57,7 @@ def step_progress_update(self, context): if input_image is None: self.report({"ERROR"}, "No open image in the Image Editor space, or selected Image Texture node.") return {"FINISHED"} - image_pixels = np.flipud( - (np.array(input_image.pixels) * 255) - .astype(np.uint8) - .reshape((input_image.size[1], input_image.size[0], input_image.channels)) - ) + image_pixels = image_utils.bpy_to_np(input_image) generated_args = context.scene.dream_textures_upscale_prompt.generate_args(context) context.scene.dream_textures_upscale_seamless_result.update_args(generated_args) @@ -97,7 +80,7 @@ def step_callback(progress: List[api.GenerationResult]) -> bool: scene.dream_textures_progress = progress[-1].progress if progress[-1].image is not None: - last_data_block = bpy_image(f"Tile {progress[-1].progress}/{progress[-1].total}", progress[-1].image.shape[1], progress[-1].image.shape[0], progress[-1].image.ravel(), last_data_block) + last_data_block = image_utils.np_to_bpy(progress[-1].image, f"Tile {progress[-1].progress}/{progress[-1].total}", last_data_block) for area in screen.areas: if area.type == 'IMAGE_EDITOR' and not area.spaces.active.use_image_pin: area.spaces.active.image = last_data_block @@ -115,7 +98,7 @@ def callback(results: List[api.GenerationResult] | Exception): last_data_block = None if results[-1].image is None: return - image = bpy_image(f"{input_image.name} (Upscaled)", results[-1].image.shape[1], results[-1].image.shape[0], results[-1].image.ravel(), last_data_block) + image = image_utils.np_to_bpy(results[-1].image, f"{input_image.name} (Upscaled)", last_data_block) for area in screen.areas: if area.type == 'IMAGE_EDITOR' and not area.spaces.active.use_image_pin: area.spaces.active.image = image diff --git a/pil_to_image.py b/pil_to_image.py deleted file mode 100644 index 7e898564..00000000 --- a/pil_to_image.py +++ /dev/null @@ -1,20 +0,0 @@ -import bpy -import numpy as np - -def pil_to_image(pil_image, name): - """ - PIL image pixels is 2D array of byte tuple (when mode is 'RGB', 'RGBA') or byte (when mode is 'L') - bpy image pixels is flat array of normalized values in RGBA order - """ - from PIL import ImageOps - width = pil_image.width - height = pil_image.height - byte_to_normalized = 1.0 / 255.0 - - bpy_image = bpy.data.images.new(name, width=width, height=height) - - # Images are upside down for Blender, so use `ImageOps.flip` to fix it. - bpy_image.pixels[:] = (np.asarray(ImageOps.flip(pil_image).convert('RGBA'),dtype=np.float32) * byte_to_normalized).ravel() - bpy_image.pack() - - return bpy_image \ No newline at end of file diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 36fa33b5..3e2e5c71 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -15,6 +15,7 @@ from functools import reduce from .. import api +from ..image_utils import bpy_to_np, grayscale def scheduler_options(self, context): return [ @@ -236,16 +237,14 @@ def generate_args(self, context, iteration=0, init_image=None, control_images=No ) case 'depth_map': task = api.DepthToImage( - depth=np.array(context.scene.init_depth.pixels) - .astype(np.float32) - .reshape((context.scene.init_depth.size[1], context.scene.init_depth.size[0], context.scene.init_depth.channels)), + depth=None if init_image is None else grayscale(bpy_to_np(context.scene.init_depth, color_space=None)), image=init_image, strength=self.strength ) case 'depth': task = api.DepthToImage( image=None, - depth=np.flipud(init_image.astype(np.float32) / 255.), + depth=None if init_image is None else grayscale(init_image), strength=self.strength ) case 'inpaint': diff --git a/property_groups/seamless_result.py b/property_groups/seamless_result.py index f1ee952f..6425feeb 100644 --- a/property_groups/seamless_result.py +++ b/property_groups/seamless_result.py @@ -5,6 +5,7 @@ from ..generator_process import Generator from ..preferences import StableDiffusionPreferences from ..api.models import GenerationArguments +from .. import image_utils def update(self, context): if hasattr(context.area, "regions"): @@ -44,9 +45,7 @@ def init(): if not can_process: return - pixels = np.empty(image.size[0]*image.size[1]*4, dtype=np.float32) - image.pixels.foreach_get(pixels) - pixels = pixels.reshape(image.size[1], image.size[0], -1) + pixels = image_utils.bpy_to_np(image) def result(future): self.result = future.result().text diff --git a/render_pass.py b/render_pass.py index a089cb9a..6b59a520 100644 --- a/render_pass.py +++ b/render_pass.py @@ -6,6 +6,7 @@ import threading from .generator_process import Generator from . import api +from . import image_utils pass_inputs = [ ('color', 'Color', 'Provide the scene color as input'), @@ -84,70 +85,53 @@ def unregister_render_pass(): # 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 - - match scene.dream_textures_render_properties_pass_inputs: - case 'color': pass - case 'depth' | 'color_depth': - depth = np.empty((size[0] * size[1], 1), dtype=np.float32) - layer.passes["Depth"].rect.foreach_get(depth) - depth = (1 - np.interp(depth, [0, np.ma.masked_equal(depth, depth.max(), copy=False).max()], [0, 1])).reshape((size[1], size[0])) - - combined_pixels = np.empty((size[0] * size[1], 4), dtype=np.float32) - rect.foreach_get(combined_pixels) + def combined(): + self.update_stats("Dream Textures", "Applying color management transforms") + return image_utils.render_pass_to_np(layer.passes["Combined"], size, color_management=True, color_space="sRGB") - 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() + def depth(): + d = image_utils.render_pass_to_np(layer.passes["Depth"], size).squeeze(2) + return (1 - np.interp(d, [0, np.ma.masked_equal(d, d.max(), copy=False).max()], [0, 1])) - self.update_stats("Dream Textures", "Generating...") + self.update_stats("Dream Textures", "Starting") prompt = scene.dream_textures_render_properties_prompt match scene.dream_textures_render_properties_pass_inputs: case 'color': task = api.ImageToImage( - np.flipud(combined_pixels.reshape((size[1], size[0], 4)) * 255).astype(np.uint8), + combined(), prompt.strength, True ) case 'depth': task = api.DepthToImage( - depth, + depth(), None, prompt.strength ) case 'color_depth': task = api.DepthToImage( - depth, - np.flipud(combined_pixels.reshape((size[1], size[0], 4)) * 255).astype(np.uint8), + depth(), + combined(), prompt.strength ) event = threading.Event() + dream_pixels = None def step_callback(progress: List[api.GenerationResult]) -> bool: self.update_progress(progress[-1].progress / progress[-1].total) - render_pass.rect.foreach_set(progress[-1].image.reshape((size[0] * size[1], 4))) + image_utils.np_to_render_pass(progress[-1].image, render_pass) self.update_result(render_result) # This does not seem to have an effect. return True def callback(results: List[api.GenerationResult] | Exception): - nonlocal combined_pixels - combined_pixels = results[-1].image + nonlocal dream_pixels + dream_pixels = results[-1].image event.set() backend: api.Backend = prompt.get_backend() generated_args: api.GenerationArguments = prompt.generate_args(bpy.context) generated_args.task = task generated_args.size = size + self.update_stats("Dream Textures", "Generating...") backend.generate( generated_args, step_callback=step_callback, @@ -158,18 +142,6 @@ def callback(results: List[api.GenerationResult] | Exception): # 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) + image_utils.np_to_render_pass(dream_pixels, render_pass, inverse_color_management=True, color_space="sRGB") self.update_stats("Dream Textures", "Finished") \ No newline at end of file diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index a6c4dea7..e300a182 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -7,7 +7,6 @@ from ...absolute_path import CLIPSEG_WEIGHTS_PATH from ..presets import DREAM_PT_AdvancedPresets -from ...pil_to_image import * from ...prompt_engineering import * from ...operators.dream_texture import DreamTexture, ReleaseGenerator, CancelGenerator, get_source_image from ...operators.open_latest_version import OpenLatestVersion, is_force_show_download, new_version_available @@ -76,7 +75,7 @@ def get_seamless_result(context, prompt): yield create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, actions_panel, get_prompt) def create_panel(space_type, region_type, parent_id, ctor, get_prompt, use_property_decorate=False, **kwargs): - class BasePanel(bpy.types.Panel): + class BasePanel(Panel): bl_category = "Dream" bl_space_type = space_type bl_region_type = region_type diff --git a/ui/panels/history.py b/ui/panels/history.py index a475e36a..2188c71d 100644 --- a/ui/panels/history.py +++ b/ui/panels/history.py @@ -1,6 +1,5 @@ import bpy from bpy.types import Panel -from ...pil_to_image import * from ...prompt_engineering import * from ...operators.dream_texture import DreamTexture, ReleaseGenerator from ...operators.view_history import ExportHistorySelection, ImportPromptFile, RecallHistoryEntry, ClearHistory, RemoveHistorySelection diff --git a/ui/panels/upscaling.py b/ui/panels/upscaling.py index a926a2e8..f5484ceb 100644 --- a/ui/panels/upscaling.py +++ b/ui/panels/upscaling.py @@ -1,5 +1,4 @@ from bpy.types import Panel -from ...pil_to_image import * from ...prompt_engineering import * from ...operators.upscale import Upscale, get_source_image from ...operators.dream_texture import CancelGenerator, ReleaseGenerator