From 9313f2573a772c79ac35340c191ddf0a710e4590 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 20 Nov 2022 16:10:02 -0500 Subject: [PATCH 01/36] Initial diffusers backend scaffolding --- .gitignore | 2 - .gitmodules | 4 - .python_dependencies/.gitignore | 1 - .python_dependencies/README | 2 - __init__.py | 187 +++++----- absolute_path.py | 2 +- classes.py | 7 +- generator_process/__init__.py | 344 +------------------ generator_process/__main__.py | 7 - generator_process/actions/huggingface_hub.py | 80 +++++ generator_process/actions/prompt_to_image.py | 109 ++++++ generator_process/actions/upscale.py | 5 + generator_process/actor.py | 152 ++++++++ operators/dream_texture.py | 14 +- operators/upscale.py | 2 +- preferences.py | 141 ++++++-- property_groups/dream_prompt.py | 8 +- render_pass.py | 2 +- requirements-dreamstudio.txt | 3 +- requirements-lin-AMD.txt | 10 - requirements-lin-win-colab-CUDA.txt | 37 -- requirements-mac-MPS-CPU.txt | 12 - requirements.txt | 6 + stable_diffusion | 1 - weights/clipseg/.gitignore | 2 - weights/config.yml | 6 - weights/realesrgan/.gitignore | 2 - weights/stable-diffusion-v1.4/.gitignore | 2 - 28 files changed, 583 insertions(+), 567 deletions(-) delete mode 100644 .gitmodules delete mode 100644 .python_dependencies/README delete mode 100644 generator_process/__main__.py create mode 100644 generator_process/actions/huggingface_hub.py create mode 100644 generator_process/actions/prompt_to_image.py create mode 100644 generator_process/actions/upscale.py create mode 100644 generator_process/actor.py delete mode 100644 requirements-lin-AMD.txt delete mode 100644 requirements-lin-win-colab-CUDA.txt delete mode 100644 requirements-mac-MPS-CPU.txt create mode 100644 requirements.txt delete mode 160000 stable_diffusion delete mode 100644 weights/clipseg/.gitignore delete mode 100644 weights/config.yml delete mode 100644 weights/realesrgan/.gitignore delete mode 100644 weights/stable-diffusion-v1.4/.gitignore diff --git a/.gitignore b/.gitignore index 37f81dcf..98fd9f00 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,5 @@ .DS_Store -python-devel.tgz - # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/.gitmodules b/.gitmodules deleted file mode 100644 index 84229e39..00000000 --- a/.gitmodules +++ /dev/null @@ -1,4 +0,0 @@ -[submodule "stable-diffusion"] - path = stable_diffusion - url = https://github.com/lstein/stable-diffusion - branch = development diff --git a/.python_dependencies/.gitignore b/.python_dependencies/.gitignore index 4a16a23d..c96a04f0 100644 --- a/.python_dependencies/.gitignore +++ b/.python_dependencies/.gitignore @@ -1,3 +1,2 @@ * -!README !.gitignore \ No newline at end of file diff --git a/.python_dependencies/README b/.python_dependencies/README deleted file mode 100644 index d8fabe15..00000000 --- a/.python_dependencies/README +++ /dev/null @@ -1,2 +0,0 @@ -# Python Dependencies -When releasing the addon, the dependencies are put into this directory. It is then added to `sys.path` for module lookup. \ No newline at end of file diff --git a/__init__.py b/__init__.py index edac8757..d6426544 100644 --- a/__init__.py +++ b/__init__.py @@ -21,94 +21,99 @@ "category": "Paint" } -import bpy -from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty -import sys -import os - -module_name = os.path.basename(os.path.dirname(__file__)) -def clear_modules(): - for name in list(sys.modules.keys()): - if name.startswith(module_name) and name != module_name: - del sys.modules[name] -clear_modules() # keep before all addon imports - -from .render_pass import register_render_pass, unregister_render_pass -from .prompt_engineering import * -from .operators.open_latest_version import check_for_updates -from .classes import CLASSES, PREFERENCE_CLASSES -from .tools import TOOLS -from .operators.dream_texture import DreamTexture, kill_generator -from .property_groups.dream_prompt import DreamPrompt -from .operators.upscale import upscale_options -from .preferences import StableDiffusionPreferences -from .ui.presets import register_default_presets - -requirements_path_items = ( - # Use the old version of requirements-win.txt to fix installation issues with Blender + PyTorch 1.12.1 - ('requirements-lin-win-colab-CUDA.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'), - ('requirements-mac-MPS-CPU.txt', 'Apple Silicon', 'Apple M1/M2'), - ('requirements-lin-AMD.txt', 'Linux (AMD)', 'Linux with AMD GPU'), -) - -def register(): - dt_op = bpy.ops - for name in DreamTexture.bl_idname.split("."): - dt_op = getattr(dt_op, name) - if hasattr(bpy.types, dt_op.idname()): # objects under bpy.ops are created on the fly, have to check that it actually exists a little differently - raise RuntimeError("Another instance of Dream Textures is already running.") - - bpy.types.Scene.dream_textures_requirements_path = EnumProperty(name="Platform", items=requirements_path_items, description="Specifies which set of dependencies to install", default='requirements-mac-MPS-CPU.txt' if sys.platform == 'darwin' else 'requirements-lin-win-colab-CUDA.txt') - - for cls in PREFERENCE_CLASSES: - bpy.utils.register_class(cls) - - check_for_updates() - - bpy.types.Scene.dream_textures_prompt = PointerProperty(type=DreamPrompt) - bpy.types.Scene.dream_textures_prompt_file = PointerProperty(type=bpy.types.Text) - bpy.types.Scene.init_img = PointerProperty(name="Init Image", type=bpy.types.Image) - bpy.types.Scene.init_mask = PointerProperty(name="Init Mask", type=bpy.types.Image) - def get_selection_preview(self): - history = bpy.context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.history - if self.dream_textures_history_selection > 0 and self.dream_textures_history_selection < len(history): - return history[self.dream_textures_history_selection].generate_prompt() - return "" - bpy.types.Scene.dream_textures_history_selection = IntProperty(default=1) - bpy.types.Scene.dream_textures_history_selection_preview = bpy.props.StringProperty(name="", default="", get=get_selection_preview, set=lambda _, __: None) - bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=0, min=0, max=0) - bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info") - - bpy.types.Scene.dream_textures_viewport_enabled = BoolProperty(name="Viewport Enabled", default=False) - bpy.types.Scene.dream_textures_render_properties_enabled = BoolProperty(default=False) - bpy.types.Scene.dream_textures_render_properties_prompt = PointerProperty(type=DreamPrompt) - bpy.types.Scene.dream_textures_upscale_outscale = bpy.props.EnumProperty(name="Target Size", items=upscale_options) - bpy.types.Scene.dream_textures_upscale_full_precision = bpy.props.BoolProperty(name="Full Precision", default=True) - bpy.types.Scene.dream_textures_upscale_seamless = bpy.props.BoolProperty(name="Seamless", default=False) - - for cls in CLASSES: - bpy.utils.register_class(cls) - - for tool in TOOLS: - bpy.utils.register_tool(tool) - - # Monkey patch cycles render passes - register_render_pass() - - register_default_presets() - -def unregister(): - for cls in PREFERENCE_CLASSES: - bpy.utils.unregister_class(cls) - - for cls in CLASSES: - bpy.utils.unregister_class(cls) - for tool in TOOLS: - bpy.utils.unregister_tool(tool) - - unregister_render_pass() - - kill_generator() - -if __name__ == "__main__": - register() \ No newline at end of file +from multiprocessing import current_process + +if current_process().name != "__actor__": + import bpy + from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty + import sys + import os + + module_name = os.path.basename(os.path.dirname(__file__)) + def clear_modules(): + for name in list(sys.modules.keys()): + if name.startswith(module_name) and name != module_name: + del sys.modules[name] + clear_modules() # keep before all addon imports + + from .render_pass import register_render_pass, unregister_render_pass + from .prompt_engineering import * + from .operators.open_latest_version import check_for_updates + from .classes import CLASSES, PREFERENCE_CLASSES + from .tools import TOOLS + from .operators.dream_texture import DreamTexture, kill_generator + from .property_groups.dream_prompt import DreamPrompt + from .operators.upscale import upscale_options + from .preferences import StableDiffusionPreferences + from .ui.presets import register_default_presets + + requirements_path_items = ( + # Use the old version of requirements-win.txt to fix installation issues with Blender + PyTorch 1.12.1 + ('requirements-lin-win-colab-CUDA.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'), + ('requirements-mac-MPS-CPU.txt', 'Apple Silicon', 'Apple M1/M2'), + ('requirements-lin-AMD.txt', 'Linux (AMD)', 'Linux with AMD GPU'), + ) + + def register(): + dt_op = bpy.ops + for name in DreamTexture.bl_idname.split("."): + dt_op = getattr(dt_op, name) + if hasattr(bpy.types, dt_op.idname()): # objects under bpy.ops are created on the fly, have to check that it actually exists a little differently + raise RuntimeError("Another instance of Dream Textures is already running.") + + bpy.types.Scene.dream_textures_requirements_path = EnumProperty(name="Platform", items=requirements_path_items, description="Specifies which set of dependencies to install", default='requirements-mac-MPS-CPU.txt' if sys.platform == 'darwin' else 'requirements-lin-win-colab-CUDA.txt') + + for cls in PREFERENCE_CLASSES: + bpy.utils.register_class(cls) + + check_for_updates() + + bpy.types.Scene.dream_textures_prompt = PointerProperty(type=DreamPrompt) + bpy.types.Scene.dream_textures_prompt_file = PointerProperty(type=bpy.types.Text) + bpy.types.Scene.init_img = PointerProperty(name="Init Image", type=bpy.types.Image) + bpy.types.Scene.init_mask = PointerProperty(name="Init Mask", type=bpy.types.Image) + def get_selection_preview(self): + history = bpy.context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.history + if self.dream_textures_history_selection > 0 and self.dream_textures_history_selection < len(history): + return history[self.dream_textures_history_selection].generate_prompt() + return "" + bpy.types.Scene.dream_textures_history_selection = IntProperty(default=1) + bpy.types.Scene.dream_textures_history_selection_preview = bpy.props.StringProperty(name="", default="", get=get_selection_preview, set=lambda _, __: None) + bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=0, min=0, max=0) + bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info") + + bpy.types.Scene.dream_textures_viewport_enabled = BoolProperty(name="Viewport Enabled", default=False) + bpy.types.Scene.dream_textures_render_properties_enabled = BoolProperty(default=False) + bpy.types.Scene.dream_textures_render_properties_prompt = PointerProperty(type=DreamPrompt) + bpy.types.Scene.dream_textures_upscale_outscale = bpy.props.EnumProperty(name="Target Size", items=upscale_options) + bpy.types.Scene.dream_textures_upscale_full_precision = bpy.props.BoolProperty(name="Full Precision", default=True) + bpy.types.Scene.dream_textures_upscale_seamless = bpy.props.BoolProperty(name="Seamless", default=False) + + for cls in CLASSES: + bpy.utils.register_class(cls) + + for tool in TOOLS: + bpy.utils.register_tool(tool) + + # Monkey patch cycles render passes + register_render_pass() + + register_default_presets() + + def unregister(): + for cls in PREFERENCE_CLASSES: + bpy.utils.unregister_class(cls) + + for cls in CLASSES: + bpy.utils.unregister_class(cls) + for tool in TOOLS: + bpy.utils.unregister_tool(tool) + + # unregister_render_pass() + + kill_generator() + + from .generator_process import Generator + from .generator_process.actor import ActorContext + gen = Generator(ActorContext.FRONTEND) + gen.start() \ No newline at end of file diff --git a/absolute_path.py b/absolute_path.py index 142df72a..16e64d2a 100644 --- a/absolute_path.py +++ b/absolute_path.py @@ -1,6 +1,6 @@ import os -def absolute_path(component): +def absolute_path(component: str): """ Returns the absolute path to a file in the addon directory. diff --git a/classes.py b/classes.py index 8c3be177..d1ff8c8d 100644 --- a/classes.py +++ b/classes.py @@ -6,7 +6,7 @@ from .operators.upscale import Upscale from .property_groups.dream_prompt import DreamPrompt from .ui.panels import dream_texture, history, upscaling, render_properties -from .preferences import OpenHuggingFace, OpenContributors, StableDiffusionPreferences, OpenDreamStudio, ImportWeights, WeightsFile, DeleteSelectedWeights +from .preferences import OpenHuggingFace, OpenContributors, StableDiffusionPreferences, OpenDreamStudio, ImportWeights, Model, DeleteSelectedWeights, ModelSearch, InstallModel, PREFERENCES_UL_ModelList from .ui.presets import DREAM_PT_AdvancedPresets, DREAM_MT_AdvancedPresets, AddAdvancedPreset, RestoreDefaultPresets @@ -44,8 +44,11 @@ ) PREFERENCE_CLASSES = ( + PREFERENCES_UL_ModelList, + ModelSearch, + InstallModel, DeleteSelectedWeights, - WeightsFile, + Model, DreamPrompt, InstallDependencies, OpenHuggingFace, diff --git a/generator_process/__init__.py b/generator_process/__init__.py index 36f746f7..3d9fb4b2 100644 --- a/generator_process/__init__.py +++ b/generator_process/__init__.py @@ -1,340 +1,10 @@ -from enum import IntEnum -import json -import subprocess -import sys -import os -import threading -import site -import traceback -import numpy as np -from multiprocessing.shared_memory import SharedMemory +from .actor import Actor -from .action import ACTION_BYTE_LENGTH, Action -from .intent import INTENT_BYTE_LENGTH, Intent - -from .registrar import BackendTarget, registrar -from .intents.apply_ocio_transforms import * -from .intents.prompt_to_image import * -from .intents.send_stop import * -from .intents.upscale import * - -MISSING_DEPENDENCIES_ERROR = "Python dependencies are missing. Click Download Latest Release to fix." - -_shared_instance = None -class GeneratorProcess(): - def __init__(self, backend: BackendTarget = BackendTarget.LOCAL): - import bpy - env = os.environ.copy() - env.pop('PYTHONPATH', None) # in case if --python-use-system-env - self.backend = backend - self.process = subprocess.Popen( - [sys.executable,'-s','generator_process', bpy.app.binary_path, '--backend', backend.name], - cwd=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), - stdin=subprocess.PIPE, stdout=subprocess.PIPE, env=env - ) - self.reader = self.process.stdout - self.queue = [] - self.in_use = False - self.killed = False - self.thread = threading.Thread(target=self._run,daemon=True,name="BackgroundReader") - self.thread.start() - - for intent in registrar._generator_intents: - # Bind self with __get__ - setattr(self, intent.name, intent.func.__get__(self, GeneratorProcess)) - - @classmethod - def shared(self, backend: BackendTarget | None = None, create=True): - global _shared_instance - if create: - if _shared_instance is None or (backend is not None and _shared_instance.backend != backend): - GeneratorProcess.kill_shared() - _shared_instance = GeneratorProcess(backend=backend if backend is not None else BackendTarget.STABILITY_SDK) - - return _shared_instance - - @classmethod - def kill_shared(self): - global _shared_instance - if _shared_instance is None: - return - _shared_instance.kill() - _shared_instance = None - - @classmethod - def can_use(self): - self = self.shared(create=False) - return not (self and self.in_use) - - def kill(self): - self.killed = True - self.process.kill() - - def send_intent(self, intent, *, payload = None, **kwargs): - """Sends intent messages to backend. - - Arguments: - * intent -- Intent enum or int - * payload -- Bytes-like value that is not suitable for json, it is recommended to use a shared memory approach instead if the payload is large - * **kwargs -- json serializable key-value pairs used for subprocess function arguments - """ - if Intent(intent) == Intent.UNKNOWN: - raise ValueError(f"Internal error, invalid Intent: {intent}") - kwargs_len = payload_len = b'\x00'*8 - if kwargs: - kwargs = bytes(json.dumps(kwargs), encoding='utf-8') - kwargs_len = len(kwargs).to_bytes(len(kwargs_len), sys.byteorder, signed=False) - if payload is not None: - payload = bytes(payload) - payload_len = len(payload).to_bytes(len(payload_len), sys.byteorder, signed=False) - # keep all checks before writing so ipc doesn't get broken intents - - stdin = self.process.stdin - stdin.write(intent.to_bytes(INTENT_BYTE_LENGTH, sys.byteorder, signed=False)) - stdin.write(kwargs_len) - if kwargs: - stdin.write(kwargs) - stdin.write(payload_len) - if payload: - stdin.write(payload) - stdin.flush() - - def _run(self): - reader = self.reader - def readUInt(length): - return int.from_bytes(reader.read(length),sys.byteorder,signed=False) - - queue = self.queue - def queue_exception_msg(msg): - queue.append((Action.EXCEPTION, {'fatal': True, 'msg': msg, 'trace': None})) - - while not self.killed: - action = readUInt(ACTION_BYTE_LENGTH) - if action == Action.CLOSED: - if not self.killed: - queue_exception_msg("Process closed unexpectedly") - return - kwargs_len = readUInt(8) - kwargs = {} if kwargs_len == 0 else json.loads(reader.read(kwargs_len)) - payload_len = readUInt(8) - if payload_len > 0: - kwargs['payload'] = reader.read(payload_len) - - if action in [Action.INFO, Action.STEP_NO_SHOW, Action.IMAGE, Action.STEP_IMAGE, Action.STOPPED]: - queue.append((action, kwargs)) - elif action == Action.EXCEPTION: - queue.append((action, kwargs)) - if kwargs['fatal']: - return - else: - queue_exception_msg(f"Internal error, unexpected action id: {action}") - return - -BYTE_TO_NORMALIZED = 1.0 / 255.0 -class Backend(): - def __init__(self, backend: BackendTarget): - self.backend_target = backend - self.intent_backends = {} - self.stdin = sys.stdin.buffer - self.stdout = sys.stdout.buffer - # stdin and stdout are piped for ipc, they should not be accessed through usual methods - # stdout can be redirected to stderr so it's still visible within the parent process's console - sys.stdout = sys.stderr - self.stderr = sys.stderr - self.shared_memory = None - self.stop_requested = False - self.intent = Intent.UNKNOWN - self.stopped_was_sent = False - self.queue = [] - self.queue_appended = threading.Event() - self.thread = threading.Thread(target=self._run,daemon=True,name="BackgroundReader") - - def check_stop(self): - if self.stop_requested: - self.stop_requested = False - raise KeyboardInterrupt - - def send_action(self, action, *, payload = None, **kwargs): - """Sends action messages to frontend. - - Arguments: - * action -- Action enum or int - * payload -- Bytes-like value that is not suitable for json, it is recommended to use a shared memory approach instead if the payload is large - * **kwargs -- json serializable key-value pairs used for callback function arguments - """ - if Action(action) == Action.UNKNOWN: - raise ValueError(f"Internal error, invalid Action: {action}") - kwargs_len = payload_len = b'\x00'*8 - if kwargs: - kwargs = bytes(json.dumps(kwargs), encoding='utf-8') - kwargs_len = len(kwargs).to_bytes(len(kwargs_len), sys.byteorder, signed=False) - if payload is not None: - payload = bytes(payload) - payload_len = len(payload).to_bytes(len(payload_len), sys.byteorder, signed=False) - # keep all checks before writing so ipc doesn't get broken actions - - self.stdout.write(action.to_bytes(ACTION_BYTE_LENGTH, sys.byteorder, signed=False)) - self.stdout.write(kwargs_len) - if kwargs: - self.stdout.write(kwargs) - self.stdout.write(payload_len) - if payload: - self.stdout.write(payload) - self.stdout.flush() - if action in [Action.EXCEPTION, Action.STOPPED]: - self.stopped_was_sent = True - - def send_info(self, msg): - """Sends information to be shown to the user before generation begins.""" - self.send_action(Action.INFO, msg=msg) - - def send_exception(self, fatal = True, msg: str = None, trace: str = None): - """Send exception information to frontend. When called within an except block arguments can be inferred. - - Arguments: - * fatal -- whether the subprocess should be killed - * msg -- user notified prompt - * trace -- traceback string - """ - exc = sys.exc_info() - if msg is None: - msg = repr(exc[1]) if exc[1] is not None else "Internal error, see system console for details" - if trace is None and exc[2] is not None: - trace = traceback.format_exc() - if msg is None and trace is None: - raise TypeError("msg and trace cannot be None outside of an except block") - self.send_action(Action.EXCEPTION, fatal=fatal, msg=msg, trace=trace) - if fatal: - sys.exit(1) - - def share_image_memory(self, image): - from PIL import ImageOps - image_bytes = (np.asarray(ImageOps.flip(image).convert('RGBA'),dtype=np.float32) * BYTE_TO_NORMALIZED).tobytes() - image_bytes_len = len(image_bytes) - shared_memory = self.shared_memory - if shared_memory is None or shared_memory.size != image_bytes_len: - if shared_memory is not None: - shared_memory.close() - self.shared_memory = shared_memory = SharedMemory(create=True, size=image_bytes_len) - shared_memory.buf[:] = image_bytes - return shared_memory.name - - """ intent generator function format - def intent(self): - args = yield - # imports and prior setup - while True: - try: # try...except is only needed if you call self.check_stop() within - ... # execute intent - except KeyboardInterrupt: - pass - args = yield +class Generator(Actor): + """ + The actor used for all background processes. """ - def _run(self): - reader = self.stdin - def readUInt(length): - return int.from_bytes(reader.read(length),sys.byteorder,signed=False) - - while True: - intent = readUInt(INTENT_BYTE_LENGTH) - json_len = readUInt(8) - if json_len == 0: - return # stdin closed - args = json.loads(reader.read(json_len)) - payload_len = readUInt(8) - if payload_len > 0: - args['payload'] = reader.read(payload_len) - if intent == Intent.STOP: - if 'stop_intent' in args and self.intent == args['stop_intent']: - self.stop_requested = True - else: - self.queue.append((intent, args)) - self.queue_appended.set() - - def main_loop(self): - intents = {} - for intent in filter(lambda x: x.backend == None or x.backend == self.backend_target, registrar._intent_backends): - print(f"Using {intent.intent} for {intent.backend}") - intents[intent.intent] = intent.func(self) - for fn in intents.values(): - next(fn) - - while True: - if len(self.queue) == 0: - self.queue_appended.clear() - self.queue_appended.wait() - (intent, args) = self.queue.pop(0) - if intent in intents: - self.intent = intent - self.stopped_was_sent = False - intents[intent].send(args) - self.stop_requested = False - self.intent = Intent.UNKNOWN - if not self.stopped_was_sent: - self.send_action(Action.STOPPED) - else: - self.send_exception(True, f"Unknown intent {intent} sent to process. Expected one of {Intent._member_names_}.") - -def main(): - import argparse - parser = argparse.ArgumentParser() - parser.add_argument('binary_path') - parser.add_argument("--backend", dest="backend", type=lambda x: BackendTarget[x], choices=list(BackendTarget)) - args = parser.parse_args() - - back = Backend(backend=args.backend) - - try: - if sys.platform == 'win32': - from ctypes import WinDLL - WinDLL(os.path.join(os.path.dirname(args.binary_path),"python3.dll")) # fix for ImportError: DLL load failed while importing cv2: The specified module could not be found. - - from absolute_path import absolute_path, CLIPSEG_WEIGHTS_PATH - # Support Apple Silicon GPUs as much as possible. - os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" - - # Move Python runtime paths to end. (prioritize addon modules) - paths = sys.path[1:] - sys.path[:] = sys.path[0:1] - - if args.backend == BackendTarget.LOCAL: - sys.path.append(absolute_path("stable_diffusion/")) - sys.path.append(absolute_path("stable_diffusion/src/clip")) - sys.path.append(absolute_path("stable_diffusion/src/k-diffusion")) - sys.path.append(absolute_path("stable_diffusion/src/taming-transformers")) - sys.path.append(absolute_path("stable_diffusion/src/clipseg")) - site.addsitedir(absolute_path(".python_dependencies")) - sys.path.extend(paths) - - import pkg_resources - pkg_resources._initialize_master_working_set() - - # Import specific modules that cause subprocess to hang if first imported after the background thread is started - if sys.platform == 'win32': - # I'm not sure if scipy will be an issue since we're not using a version that has - # libbanded5x.Q3V52YHHGVBP5BKVHJ5RHQVFWHHSLVWO.gfortran-win_amd64.dll anymore. - # Importing skimage is indirectly loading scipy anyway. Keeping note for future reference. - - # Couldn't track down exactly where it was hanging but it's good enough. - if os.path.exists(absolute_path(".python_dependencies/skimage")): - from skimage import transform - - if args.backend == BackendTarget.LOCAL: - from ldm.invoke import txt2mask - txt2mask.CLIPSEG_WEIGHTS = CLIPSEG_WEIGHTS_PATH - - try: - import certifi - os.environ["SSL_CERT_FILE"] = certifi.where() - using_certifi = True - except ModuleNotFoundError: - using_certifi = False - print(f"Using certifi: {using_certifi}") - - back.thread.start() - back.main_loop() - except SystemExit: - pass - except: - back.send_exception() \ No newline at end of file + from .actions.prompt_to_image import prompt_to_image + from .actions.upscale import upscale + from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models \ No newline at end of file diff --git a/generator_process/__main__.py b/generator_process/__main__.py deleted file mode 100644 index cf405d33..00000000 --- a/generator_process/__main__.py +++ /dev/null @@ -1,7 +0,0 @@ -import sys -import os -sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) - -from generator_process import main - -main() \ No newline at end of file diff --git a/generator_process/actions/huggingface_hub.py b/generator_process/actions/huggingface_hub.py new file mode 100644 index 00000000..e11ff1eb --- /dev/null +++ b/generator_process/actions/huggingface_hub.py @@ -0,0 +1,80 @@ +from dataclasses import dataclass +import os + +@dataclass +class Model: + id: str + author: str + tags: list[str] + likes: int + downloads: int + +def hf_list_models( + self, + query: str +) -> list[Model]: + from huggingface_hub import HfApi, ModelFilter + + if hasattr(self, "huggingface_hub_api"): + api: HfApi = self.huggingface_hub_api + else: + api = HfApi() + setattr(self, "huggingface_hub_api", api) + + filter = ModelFilter(tags="diffusers", task="text-to-image") + models = api.list_models( + filter=filter, + search=query + ) + + return list(map(lambda m: Model(m.modelId, m.author, m.tags, m.likes, getattr(m, "downloads", -1)), models)) + +def hf_list_installed_models(self) -> list[Model]: + from diffusers.utils import DIFFUSERS_CACHE + return list( + filter( + lambda x: os.path.isdir(x.id), + map(lambda x: Model(os.path.join(DIFFUSERS_CACHE, x), "", [], -1, -1), os.listdir(DIFFUSERS_CACHE)) + ) + ) + +def hf_snapshot_download( + self, + model: str, + token: str +) -> None: + from huggingface_hub import snapshot_download + from diffusers import StableDiffusionPipeline + from diffusers.utils import DIFFUSERS_CACHE, WEIGHTS_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME + from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME + from diffusers.hub_utils import http_user_agent + config_dict = StableDiffusionPipeline.get_config_dict( + model, + cache_dir=DIFFUSERS_CACHE, + resume_download=True, + force_download=False, + use_auth_token=token + ) + # make sure we only download sub-folders and `diffusers` filenames + folder_names = [k for k in config_dict.keys() if not k.startswith("_")] + allow_patterns = [os.path.join(k, "*") for k in folder_names] + allow_patterns += [WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME, StableDiffusionPipeline.config_name] + + # make sure we don't download flax weights + ignore_patterns = "*.msgpack" + + requested_pipeline_class = config_dict.get("_class_name", StableDiffusionPipeline.__name__) + user_agent = {"pipeline_class": requested_pipeline_class} + user_agent = http_user_agent(user_agent) + + # download all allow_patterns + cached_folder = snapshot_download( + model, + cache_dir=DIFFUSERS_CACHE, + resume_download=True, + use_auth_token=token, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + user_agent=user_agent, + ) + return \ No newline at end of file diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py new file mode 100644 index 00000000..e9931d35 --- /dev/null +++ b/generator_process/actions/prompt_to_image.py @@ -0,0 +1,109 @@ +from typing import Optional, Annotated, Union, _AnnotatedAlias +import enum +import os +from dataclasses import dataclass + +class Pipeline(enum.IntEnum): + STABLE_DIFFUSION = 0 + + STABILITY_SDK = 1 + +@dataclass(eq=True) +class Optimizations: + attention_slicing = True + attention_slice_size: Union[str, int] = "auto" + cudnn_benchmark: Annotated[bool, "cuda"] = False + tf32: Annotated[bool, "cuda"] = False + amp: Annotated[bool, "cuda"] = False + half_precision: Annotated[bool, "cuda"] = True + sequential_cpu_offload = False + channels_last_memory_format = False + xformers_attention = False + + def can_use(self, property, device) -> bool: + if not getattr(self, property): + return False + if isinstance(getattr(self.__annotations__, property, None), _AnnotatedAlias): + annotation: _AnnotatedAlias = self.__annotations__[property] + return annotation.__metadata__ != device + return True + +def _choose_device(): + """ + Automatically select which PyTorch device to use. + """ + import torch + if torch.cuda.is_available(): + return "cuda" + elif torch.backends.mps.is_available(): + return "mps" + else: + return "cpu" + +def prompt_to_image( + self, + pipeline: Pipeline, + + model: str, + + prompt: str, + steps: int, + width: int, + height: int, + + optimizations: Optimizations, + + **kwargs +) -> Optional[bytes]: + match pipeline: + case Pipeline.STABLE_DIFFUSION: + import diffusers + import torch + from ...absolute_path import WEIGHTS_PATH + + device = _choose_device() + + if optimizations.can_use("cudnn_benchmark", device): + torch.backends.cudnn.benchmark = True + + if optimizations.can_use("tf32", device): + torch.backends.cuda.matmul.allow_tf32 = True + + if hasattr(self, "_cached_pipe") and self._cached_pipe[1] == optimizations: + pipe = self._cached_pipe[0] + else: + pipe = diffusers.StableDiffusionPipeline.from_pretrained( + os.path.join(WEIGHTS_PATH, model), + torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32, + ) + pipe = pipe.to(device) + + if optimizations.can_use("attention_slicing", device): + pipe.enable_attention_slicing(optimizations.attention_slice_size) + + if optimizations.can_use("sequential_cpu_offload", device): + pipe.enable_sequential_cpu_offload() + + if optimizations.can_use("channels_last_memory_format", device): + pipe.unet.to(memory_format=torch.channels_last) + + if optimizations.can_use("xformers_attention", device): + pipe.enable_xformers_memory_efficient_attention() + + setattr(self, "_cached_pipe", (pipe, optimizations)) + + if device == "mps": + # First-time "warmup" pass (necessary on MPS as of diffusers 0.7.2) + _ = pipe(prompt, num_inference_steps=1) + + with torch.inference_mode(mode=False), torch.autocast(device, enabled=optimizations.can_use("amp", device)): + return pipe( + prompt, + num_inference_steps=steps, + width=width, + height=height + ).images[0] + case Pipeline.STABILITY_SDK: + import stability_sdk + case _: + raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py new file mode 100644 index 00000000..1cf08fee --- /dev/null +++ b/generator_process/actions/upscale.py @@ -0,0 +1,5 @@ +def upscale( + self, + input: bytes +) -> bytes: + return input \ No newline at end of file diff --git a/generator_process/actor.py b/generator_process/actor.py new file mode 100644 index 00000000..a196f9bc --- /dev/null +++ b/generator_process/actor.py @@ -0,0 +1,152 @@ +from multiprocessing import Queue, Process +import subprocess +import enum +import threading +import functools +from typing import Type, TypeVar, Optional +from concurrent.futures import Future +import site +import os +import sys + +class ActorContext(enum.IntEnum): + """ + The context of an `Actor` object. + + One `Actor` instance is the `FRONTEND`, while the other instance is the backend, which runs in a separate process. + The `FRONTEND` sends messages to the `BACKEND`, which does work and returns a result. + """ + FRONTEND = 0 + BACKEND = 1 + +class Message(): + """ + Represents a function signature with a method name, positonal arguments, and keyword arguments. + + Note: All arguments must be picklable. + """ + + def __init__(self, method_name, args, kwargs): + self.method_name = method_name + self.args = args + self.kwargs = kwargs + + +def _start_backend(cls, message_queue, response_queue): + cls( + ActorContext.BACKEND, + message_queue=message_queue, + response_queue=response_queue + ).start() + +T = TypeVar('T', bound='Actor') + +class Actor(): + """ + Base class for specialized actors. + + Uses queues to serialize actions from different threads, and automatically dispatches methods to a separate process. + """ + + _message_queue: Queue + _response_queue: Queue + + _shared_instance = None + + # Methods that are not used for message passing, and should not be overridden in `_setup`. + _protected_methods = { + "start", + "close", + "is_alive", + "can_use", + "shared" + } + + def __init__(self, context: ActorContext, message_queue: Queue = Queue(), response_queue: Queue = Queue()): + self.context = context + self._message_queue = message_queue + self._response_queue = response_queue + self._setup() + self.__class__._shared_instance = self + + def _setup(self): + """ + Setup the Actor after initialization. + """ + match self.context: + case ActorContext.FRONTEND: + for name in filter(lambda name: callable(getattr(self, name)) and not name.startswith("_") and name not in self._protected_methods, dir(self)): + setattr(self, name, self._send(name)) + case ActorContext.BACKEND: + pass + + @classmethod + def shared(cls: Type[T]) -> T: + return cls._shared_instance or cls(ActorContext.FRONTEND) + + def start(self): + """ + Start the actor process. + """ + match self.context: + case ActorContext.FRONTEND: + self.process = Process(target=_start_backend, args=(self.__class__, self._message_queue, self._response_queue), name="__actor__") + self.process.start() + case ActorContext.BACKEND: + self._backend_loop() + + def close(self): + """ + Stop the actor process. + """ + match self.context: + case ActorContext.FRONTEND: + self.process.terminate() + self._message_queue.close() + self._response_queue.close() + case ActorContext.BACKEND: + pass + + def is_alive(self): + match self.context: + case ActorContext.FRONTEND: + return self.process.is_alive() + case ActorContext.BACKEND: + return True + + def can_use(self): + return self._message_queue.empty() and self._response_queue.empty() + + def _load_dependencies(self): + from ..absolute_path import absolute_path + site.addsitedir(absolute_path(".python_dependencies")) + + def _backend_loop(self): + self._load_dependencies() + while True: + self._receive(self._message_queue.get()) + + def _receive(self, message: Message): + try: + response = getattr(self, message.method_name)(*message.args, **message.kwargs) + except Exception as e: + response = e + self._response_queue.put(response) + + def _send(self, name): + def _send(*args, **kwargs): + future = Future() + def wait_for_response(future: Future): + response = self._response_queue.get() + if isinstance(response, Exception): + future.set_exception(response) + else: + future.set_result(response) + thread = threading.Thread(target=functools.partial(wait_for_response, future)) + thread.start() + self._message_queue.put(Message(name, args, kwargs)) + return future + return _send + + def __del__(self): + self.close() \ No newline at end of file diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 156b3ca1..b2049fd4 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -12,7 +12,8 @@ from ..pil_to_image import * from ..prompt_engineering import * from ..absolute_path import WEIGHTS_PATH, CLIPSEG_WEIGHTS_PATH -from ..generator_process import MISSING_DEPENDENCIES_ERROR, GeneratorProcess, Intent +from ..generator_process import Generator +from ..generator_process.actions.prompt_to_image import Pipeline, Optimizations import tempfile @@ -34,7 +35,7 @@ class DreamTexture(bpy.types.Operator): @classmethod def poll(cls, context): - return GeneratorProcess.can_use() + return Generator.shared().can_use() def invoke(self, context, event): if weights_are_installed(self.report): @@ -115,7 +116,12 @@ def view_step(step, width=None, height=None, shared_memory_name=None): bpy.data.images.remove(last_data_block) last_data_block = step_image return # Only perform this on the first image editor found. - dream_texture(context.scene.dream_textures_prompt, view_step, image_writer) + # dream_texture(context.scene.dream_textures_prompt, view_step, image_writer) + Generator.shared().prompt_to_image( + Pipeline.STABLE_DIFFUSION, + optimizations=Optimizations(), + **scene.dream_textures_prompt.generate_args(), + ) return {"FINISHED"} headless_prompt = None @@ -293,7 +299,7 @@ def modal_stopped(context): last_data_block = None def kill_generator(context=bpy.context): - GeneratorProcess.kill_shared() + Generator.shared().close() modal_stopped(context) class ReleaseGenerator(bpy.types.Operator): diff --git a/operators/upscale.py b/operators/upscale.py index cce7afff..abb0dc2c 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -3,7 +3,7 @@ from multiprocessing.shared_memory import SharedMemory import numpy as np import sys -from ..generator_process import GeneratorProcess +# from ..generator_process import GeneratorProcess upscale_options = [ ("2", "2x", "", 2), diff --git a/preferences.py b/preferences.py index cfc15268..7a957d4d 100644 --- a/preferences.py +++ b/preferences.py @@ -4,21 +4,26 @@ import os import webbrowser import shutil +from concurrent.futures import Future +import functools from .absolute_path import WEIGHTS_PATH, absolute_path from .operators.install_dependencies import InstallDependencies from .operators.open_latest_version import OpenLatestVersion from .property_groups.dream_prompt import DreamPrompt from .ui.presets import RestoreDefaultPresets, default_presets_missing +from .generator_process import Generator +from .generator_process.actions.huggingface_hub import Model +from typing import List class OpenHuggingFace(bpy.types.Operator): bl_idname = "dream_textures.open_hugging_face" - bl_label = "Download Weights from Hugging Face" - bl_description = ("Opens huggingface.co to the download page for the model weights.") + bl_label = "Get Access Token" + bl_description = ("Opens huggingface.co to the tokens page") bl_options = {"REGISTER", "INTERNAL"} def execute(self, context): - webbrowser.open("https://huggingface.co/CompVis/stable-diffusion-v-1-4-original") + webbrowser.open("https://huggingface.co/settings/tokens") return {"FINISHED"} class ImportWeights(bpy.types.Operator, ImportHelper): @@ -73,59 +78,121 @@ def execute(self, context): webbrowser.open("https://beta.dreamstudio.ai/membership?tab=apiKeys") return {"FINISHED"} -class WeightsFile(bpy.types.PropertyGroup): - bl_label = "Weights File" - bl_idname = "dream_textures.WeightsFile" +class Model(bpy.types.PropertyGroup): + bl_label = "Model" + bl_idname = "dream_textures.Model" - name: bpy.props.StringProperty(name="Path") + model: bpy.props.StringProperty(name="Model") + downloads: bpy.props.IntProperty(name="Downloads") + likes: bpy.props.IntProperty(name="Likes") -class PREFERENCES_UL_WeightsFileList(bpy.types.UIList): +class PREFERENCES_UL_ModelList(bpy.types.UIList): def draw_item(self, context, layout, data, item, icon, active_data, active_propname): - layout.label(text=item.name) + model_name = item.model + is_installed = False + if os.path.exists(item.model): + model_name = os.path.basename(item.model).replace('models--', '').replace('--', '/') + is_installed = True + split = layout.split(factor=0.75) + split.label(text=model_name) + if item.downloads != -1: + split.label(text=str(item.downloads), icon="IMPORT") + if item.downloads != -1: + split.label(text=str(item.likes), icon="HEART") + layout.operator(InstallModel.bl_idname, text="", icon="FILE_FOLDER" if is_installed else "IMPORT").model = item.model + +model_installing = False + +@staticmethod +def set_model_list(model_list: str, future: Future): + getattr(bpy.context.preferences.addons[__package__].preferences, model_list).clear() + for model in future.result(): + m = getattr(bpy.context.preferences.addons[__package__].preferences, model_list).add() + m.model = model.id + m.downloads = model.downloads + m.likes = model.likes + +class ModelSearch(bpy.types.Operator): + bl_idname = "dream_textures.model_search" + bl_label = "Search" + bl_description = ("Searches HuggingFace Hub for models") + bl_options = {"REGISTER", "INTERNAL"} + + def execute(self, context): + + return {"FINISHED"} + +class InstallModel(bpy.types.Operator): + bl_idname = "dream_textures.install_model" + bl_label = "Install or Open" + bl_description = ("Install or open a model from the cache") + bl_options = {"REGISTER", "INTERNAL"} + + model: StringProperty(name="Model ID") + + def execute(self, context): + if os.path.exists(self.model): + webbrowser.open(f"file://{self.model}") + else: + global model_installing + model_installing = True + def done_installing(_): + global model_installing + model_installing = False + Generator.shared().hf_list_installed_models().add_done_callback(functools.partial(set_model_list, 'installed_models')) + Generator.shared().hf_snapshot_download(self.model, bpy.context.preferences.addons[__package__].preferences.hf_token).add_done_callback(done_installing) + return {"FINISHED"} + +def _model_search(self, context): + Generator.shared().hf_list_models(self.model_query).add_done_callback(functools.partial(set_model_list, 'model_results')) class StableDiffusionPreferences(bpy.types.AddonPreferences): bl_idname = __package__ history: CollectionProperty(type=DreamPrompt) - - weights: CollectionProperty(type=WeightsFile) - active_weights: bpy.props.IntProperty(name="Active Weights", default=0) dream_studio_key: StringProperty(name="DreamStudio Key") + model_query: StringProperty(name="Search", update=_model_search) + model_results: CollectionProperty(type=Model) + active_model_result: bpy.props.IntProperty(name="Active Model", default=0) + hf_token: StringProperty(name="HuggingFace Token") + + installed_models: CollectionProperty(type=Model) + active_installed_model: bpy.props.IntProperty(name="Active Model", default=0) + + @staticmethod + def register(): + Generator.shared().hf_list_installed_models().add_done_callback(functools.partial(set_model_list, 'installed_models')) + def draw(self, context): layout = self.layout - self.weights.clear() - for path in filter(lambda f: f.endswith('.ckpt'), os.listdir(WEIGHTS_PATH)): - weights_file = self.weights.add() - weights_file.name = path - weights_installed = len(self.weights) > 0 + weights_installed = len(self.installed_models) > 0 if not weights_installed: layout.label(text="Complete the following steps to finish setting up the addon:") - + has_dependencies = len(os.listdir(absolute_path(".python_dependencies"))) > 2 if has_dependencies: - has_local = len(os.listdir(absolute_path("stable_diffusion"))) > 0 - if has_local: - dependencies_box = layout.box() - dependencies_box.label(text="Dependencies Located", icon="CHECKMARK") - dependencies_box.label(text="All dependencies (except for model weights) are included in the release.") - - model_weights_box = layout.box() - model_weights_box.label(text="Setup Model Weights", icon="SETTINGS") - if weights_installed: - model_weights_box.label(text="Model weights setup successfully.", icon="CHECKMARK") - else: - model_weights_box.label(text="The model weights are not distributed with the addon.") - model_weights_box.label(text="Follow the steps below to download and install them.") - model_weights_box.label(text="1. Download the file 'sd-v1-4.ckpt'") - model_weights_box.operator(OpenHuggingFace.bl_idname, icon="URL") - model_weights_box.label(text="2. Select the downloaded weights to install.") - model_weights_box.operator(ImportWeights.bl_idname, text="Import Model Weights", icon="IMPORT") - model_weights_box.template_list("UI_UL_list", "dream_textures_weights", self, "weights", self, "active_weights") - model_weights_box.operator(DeleteSelectedWeights.bl_idname, text="Delete Selected Weights", icon="X") + has_local = os.path.exists(absolute_path(".python_dependencies/diffusers")) > 0 + + if has_local: + search_box = layout.box() + search_box.label(text="Find Models", icon="SETTINGS") + search_box.label(text="Search HuggingFace Hub for compatible models.") + + auth_row = search_box.row() + auth_row.prop(self, "hf_token", text="Token") + auth_row.operator(OpenHuggingFace.bl_idname, text="Get Your Token", icon="KEYINGSET") + + search_box.prop(self, "model_query", text="", icon="VIEWZOOM") + + if len(self.model_results) > 0: + search_box.enabled = not model_installing + search_box.template_list(PREFERENCES_UL_ModelList.__name__, "dream_textures_model_results", self, "model_results", self, "active_model_result") + + layout.template_list(PREFERENCES_UL_ModelList.__name__, "dream_textures_installed_models", self, "installed_models", self, "active_installed_model") dream_studio_box = layout.box() dream_studio_box.label(text=f"DreamStudio{' (Optional)' if has_local else ''}", icon="HIDE_OFF") diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 7bec51ef..fea3ef92 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -53,12 +53,12 @@ def inpaint_mask_sources_filtered(self, context): ('xy', 'Both', '', 3), ] -def weights_options(self, context): +def model_options(self, context): return [(f, f, '', i) for i, f in enumerate(filter(lambda f: f.endswith('.ckpt'), os.listdir(WEIGHTS_PATH)))] def backend_options(self, context): def options(): - if len(os.listdir(absolute_path("stable_diffusion"))) > 0: + if os.path.exists(absolute_path(".python_dependencies/diffusers")): yield (BackendTarget.LOCAL.name, 'Local', 'Run on your own hardware', 1) if len(context.preferences.addons[__package__.split('.')[0]].preferences.dream_studio_key) > 0: yield (BackendTarget.STABILITY_SDK.name, 'DreamStudio', 'Run in the cloud with DreamStudio', 2) @@ -74,8 +74,8 @@ def seed_clamp(self, ctx): pass # will get hashed once generated attributes = { - "backend": EnumProperty(name="Backend", items=backend_options, default=1 if len(os.listdir(absolute_path("stable_diffusion"))) > 0 else 2, description="Fill in a few simple options to create interesting images quickly"), - "model": EnumProperty(name="Model", items=weights_options, description="Specify which weights file to use for inference"), + "backend": EnumProperty(name="Backend", items=backend_options, default=1 if os.path.exists(absolute_path(".python_dependencies/diffusers")) else 2, description="Fill in a few simple options to create interesting images quickly"), + "model": EnumProperty(name="Model", items=model_options, description="Specify which model to use for inference"), # Prompt "prompt_structure": EnumProperty(name="Preset", items=prompt_structures_items, description="Fill in a few simple options to create interesting images quickly"), diff --git a/render_pass.py b/render_pass.py index 9546131c..1f1d8678 100644 --- a/render_pass.py +++ b/render_pass.py @@ -7,7 +7,7 @@ import os from multiprocessing.shared_memory import SharedMemory -from .generator_process import GeneratorProcess +from .generator_process import Generator from .operators.dream_texture import dream_texture, weights_are_installed diff --git a/requirements-dreamstudio.txt b/requirements-dreamstudio.txt index c4a4957b..6f8f619d 100644 --- a/requirements-dreamstudio.txt +++ b/requirements-dreamstudio.txt @@ -1 +1,2 @@ -stability-sdk==0.2.6 \ No newline at end of file +stability-sdk==0.2.6 +opencolorio \ No newline at end of file diff --git a/requirements-lin-AMD.txt b/requirements-lin-AMD.txt deleted file mode 100644 index bfc290b0..00000000 --- a/requirements-lin-AMD.txt +++ /dev/null @@ -1,10 +0,0 @@ --r requirements.txt - -# Get hardware-appropriate torch/torchvision ---extra-index-url https://download.pytorch.org/whl/rocm5.1.1 --trusted-host https://download.pytorch.org -torch -torchvision -realesrgan==0.2.5.0 -opencolorio -stability-sdk==0.2.6 --e . diff --git a/requirements-lin-win-colab-CUDA.txt b/requirements-lin-win-colab-CUDA.txt deleted file mode 100644 index 48696758..00000000 --- a/requirements-lin-win-colab-CUDA.txt +++ /dev/null @@ -1,37 +0,0 @@ -albumentations==0.4.3 -einops==0.3.0 -huggingface-hub==0.8.1 -imageio-ffmpeg==0.4.2 -imageio==2.9.0 -kornia==0.6.0 -# pip will resolve the version which matches torch -numpy -omegaconf==2.1.1 -opencv-python==4.6.0.66 -pillow==9.2.0 -pip>=22 -pudb==2019.2 -pytorch-lightning==1.4.2 -streamlit==1.12.0 -# "CompVis/taming-transformers" doesn't work -# ldm\models\autoencoder.py", line 6, in -# from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer -# ModuleNotFoundError -taming-transformers-rom1504==0.0.6 -test-tube>=0.7.5 -torch-fidelity==0.3.0 -torchmetrics==0.6.0 -transformers==4.19.2 -git+https://github.com/openai/CLIP.git@main#egg=clip -git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion -git+https://github.com/lstein/GFPGAN@fix-dark-cast-images#egg=gfpgan --e git+https://github.com/invoke-ai/clipseg.git@models-rename#egg=clipseg -# No CUDA in PyPi builds ---extra-index-url https://download.pytorch.org/whl/cu113 --trusted-host https://download.pytorch.org -torch==1.11.0 -# Same as numpy - let pip do its thing -torchvision -realesrgan==0.2.5.0 -opencolorio -stability-sdk==0.2.6 --e . diff --git a/requirements-mac-MPS-CPU.txt b/requirements-mac-MPS-CPU.txt deleted file mode 100644 index b279d3a1..00000000 --- a/requirements-mac-MPS-CPU.txt +++ /dev/null @@ -1,12 +0,0 @@ --r stable_diffusion/requirements.txt - ---pre ---extra-index-url https://download.pytorch.org/whl/nightly/cpu --trusted-host https://download.pytorch.org - -protobuf==3.19.4 -torch==1.12.1 -torchvision -realesrgan==0.2.5.0 -opencolorio -stability-sdk==0.2.6 --e . diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..6c76a587 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +diffusers +transformers +huggingface_hub +torch>=1.13 +stability-sdk==0.2.6 +opencolorio \ No newline at end of file diff --git a/stable_diffusion b/stable_diffusion deleted file mode 160000 index 3081b6b7..00000000 --- a/stable_diffusion +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 3081b6b7dd4c2fb1156e7a99dc461012c4ecda35 diff --git a/weights/clipseg/.gitignore b/weights/clipseg/.gitignore deleted file mode 100644 index c96a04f0..00000000 --- a/weights/clipseg/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -* -!.gitignore \ No newline at end of file diff --git a/weights/config.yml b/weights/config.yml deleted file mode 100644 index b8863379..00000000 --- a/weights/config.yml +++ /dev/null @@ -1,6 +0,0 @@ -stable-diffusion-1.4: - config: stable_diffusion/configs/stable-diffusion/v1-inference.yaml - weights: weights/stable-diffusion-v1.4/model.ckpt - description: Stable Diffusion inference model version 1.4 - width: 512 - height: 512 \ No newline at end of file diff --git a/weights/realesrgan/.gitignore b/weights/realesrgan/.gitignore deleted file mode 100644 index c96a04f0..00000000 --- a/weights/realesrgan/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -* -!.gitignore \ No newline at end of file diff --git a/weights/stable-diffusion-v1.4/.gitignore b/weights/stable-diffusion-v1.4/.gitignore deleted file mode 100644 index c96a04f0..00000000 --- a/weights/stable-diffusion-v1.4/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -* -!.gitignore \ No newline at end of file From 3946ee43e605214cb88d50040f5537705ef680d5 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 24 Nov 2022 11:42:14 -0500 Subject: [PATCH 02/36] Support scheduler selection --- __init__.py | 7 +- generator_process/__init__.py | 2 +- generator_process/actions/prompt_to_image.py | 113 +++++++++++++++---- generator_process/actor.py | 41 +++++-- generator_process/intents/prompt_to_image.py | 2 +- operators/dream_texture.py | 23 ++-- operators/view_history.py | 4 +- property_groups/dream_prompt.py | 58 +++++++--- render_pass.py | 4 +- requirements.txt | 1 + ui/panels/dream_texture.py | 54 ++++++++- ui/panels/render_properties.py | 2 +- ui/presets.py | 2 +- 13 files changed, 230 insertions(+), 83 deletions(-) diff --git a/__init__.py b/__init__.py index d6426544..d765da7e 100644 --- a/__init__.py +++ b/__init__.py @@ -111,9 +111,4 @@ def unregister(): # unregister_render_pass() - kill_generator() - - from .generator_process import Generator - from .generator_process.actor import ActorContext - gen = Generator(ActorContext.FRONTEND) - gen.start() \ No newline at end of file + kill_generator() \ No newline at end of file diff --git a/generator_process/__init__.py b/generator_process/__init__.py index 3d9fb4b2..1ba3fcef 100644 --- a/generator_process/__init__.py +++ b/generator_process/__init__.py @@ -5,6 +5,6 @@ class Generator(Actor): The actor used for all background processes. """ - from .actions.prompt_to_image import prompt_to_image + from .actions.prompt_to_image import prompt_to_image, choose_device from .actions.upscale import upscale from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models \ No newline at end of file diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index e9931d35..bed33e96 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -2,33 +2,64 @@ import enum import os from dataclasses import dataclass +from contextlib import nullcontext +import numpy as np +from numpy.typing import NDArray class Pipeline(enum.IntEnum): STABLE_DIFFUSION = 0 STABILITY_SDK = 1 +class Scheduler(enum.Enum): + DDIM = "DDIM" + PNDM = "PNDM" + LMS_DISCRETE = "LMS Discrete" + EULER_DISCRETE = "Euler Discrete" + EULER_ANCESTRAL_DISCRETE = "Euler Ancestral Discrete" + + def create(self, pipeline, pretrained): + import diffusers + def scheduler_class(): + match self: + case Scheduler.DDIM: + return diffusers.DDIMScheduler + case Scheduler.PNDM: + return diffusers.PNDMScheduler + case Scheduler.LMS_DISCRETE: + return diffusers.LMSDiscreteScheduler + case Scheduler.EULER_DISCRETE: + return diffusers.EulerDiscreteScheduler + case Scheduler.EULER_ANCESTRAL_DISCRETE: + return diffusers.EulerAncestralDiscreteScheduler + if pretrained is not None: + return scheduler_class().from_pretrained(pretrained.model_path, subfolder=pretrained.subfolder) + else: + return scheduler_class().from_config(pipeline.scheduler.config) + + @dataclass(eq=True) class Optimizations: - attention_slicing = True + attention_slicing: bool = True attention_slice_size: Union[str, int] = "auto" + inference_mode: Annotated[bool, "cuda"] = True cudnn_benchmark: Annotated[bool, "cuda"] = False tf32: Annotated[bool, "cuda"] = False amp: Annotated[bool, "cuda"] = False half_precision: Annotated[bool, "cuda"] = True - sequential_cpu_offload = False - channels_last_memory_format = False - xformers_attention = False + sequential_cpu_offload: bool = False + channels_last_memory_format: bool = False + # xformers_attention: bool = False # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. def can_use(self, property, device) -> bool: if not getattr(self, property): return False - if isinstance(getattr(self.__annotations__, property, None), _AnnotatedAlias): + if isinstance(self.__annotations__.get(property, None), _AnnotatedAlias): annotation: _AnnotatedAlias = self.__annotations__[property] - return annotation.__metadata__ != device + return annotation.__metadata__ == device return True -def _choose_device(): +def choose_device(self) -> str: """ Automatically select which PyTorch device to use. """ @@ -46,22 +77,27 @@ def prompt_to_image( model: str, + scheduler: Scheduler, + + optimizations: Optimizations, + prompt: str, steps: int, width: int, height: int, - - optimizations: Optimizations, + seed: int, **kwargs -) -> Optional[bytes]: +) -> Optional[NDArray]: match pipeline: case Pipeline.STABLE_DIFFUSION: import diffusers import torch + from PIL import ImageOps from ...absolute_path import WEIGHTS_PATH - device = _choose_device() + device = self.choose_device() + optimizations.can_use("amp", device) if optimizations.can_use("cudnn_benchmark", device): torch.backends.cudnn.benchmark = True @@ -69,13 +105,24 @@ def prompt_to_image( if optimizations.can_use("tf32", device): torch.backends.cuda.matmul.allow_tf32 = True - if hasattr(self, "_cached_pipe") and self._cached_pipe[1] == optimizations: + if hasattr(self, "_cached_pipe") and self._cached_pipe[1] == (model, scheduler, optimizations): pipe = self._cached_pipe[0] else: + storage_folder = os.path.join(WEIGHTS_PATH, model) + revision = "main" + ref_path = os.path.join(storage_folder, "refs", revision) + with open(ref_path) as f: + commit_hash = f.read() + + snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) pipe = diffusers.StableDiffusionPipeline.from_pretrained( - os.path.join(WEIGHTS_PATH, model), + snapshot_folder, torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32, ) + pipe.scheduler = scheduler.create(pipe, { + 'model_path': snapshot_folder, + 'subfolder': 'scheduler', + } if ('stable-diffusion-2' in snapshot_folder) else None) pipe = pipe.to(device) if optimizations.can_use("attention_slicing", device): @@ -87,23 +134,39 @@ def prompt_to_image( if optimizations.can_use("channels_last_memory_format", device): pipe.unet.to(memory_format=torch.channels_last) - if optimizations.can_use("xformers_attention", device): - pipe.enable_xformers_memory_efficient_attention() + # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. + # if optimizations.can_use("xformers_attention", device): + # pipe.enable_xformers_memory_efficient_attention() - setattr(self, "_cached_pipe", (pipe, optimizations)) + setattr(self, "_cached_pipe", (pipe, (model, scheduler, optimizations))) - if device == "mps": # First-time "warmup" pass (necessary on MPS as of diffusers 0.7.2) - _ = pipe(prompt, num_inference_steps=1) + if device == "mps": + _ = pipe(prompt, num_inference_steps=1) - with torch.inference_mode(mode=False), torch.autocast(device, enabled=optimizations.can_use("amp", device)): - return pipe( - prompt, - num_inference_steps=steps, - width=width, - height=height - ).images[0] + generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + + with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + i = pipe.__call__( + prompt=prompt, + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=7.5, + negative_prompt=None, + num_images_per_prompt=1, + eta=0.0, + generator=generator, + latents=None, + output_type="pil", + return_dict=True, + callback=None, + callback_steps=1 + ).images[0] + return np.asarray(ImageOps.flip(i).convert('RGBA'), dtype=np.float32) * 1/255. case Pipeline.STABILITY_SDK: import stability_sdk + raise NotImplementedError() case _: raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actor.py b/generator_process/actor.py index a196f9bc..3a625dbb 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -1,13 +1,11 @@ from multiprocessing import Queue, Process -import subprocess +import queue import enum +import traceback import threading -import functools -from typing import Type, TypeVar, Optional +from typing import Type, TypeVar from concurrent.futures import Future import site -import os -import sys class ActorContext(enum.IntEnum): """ @@ -31,7 +29,6 @@ def __init__(self, method_name, args, kwargs): self.args = args self.kwargs = kwargs - def _start_backend(cls, message_queue, response_queue): cls( ActorContext.BACKEND, @@ -39,6 +36,11 @@ def _start_backend(cls, message_queue, response_queue): response_queue=response_queue ).start() +class TracedError(BaseException): + def __init__(self, base: BaseException, trace: str): + self.base = base + self.trace = trace + T = TypeVar('T', bound='Actor') class Actor(): @@ -82,9 +84,9 @@ def _setup(self): @classmethod def shared(cls: Type[T]) -> T: - return cls._shared_instance or cls(ActorContext.FRONTEND) + return cls._shared_instance or cls(ActorContext.FRONTEND).start() - def start(self): + def start(self: T) -> T: """ Start the actor process. """ @@ -94,6 +96,7 @@ def start(self): self.process.start() case ActorContext.BACKEND: self._backend_loop() + return self def close(self): """ @@ -107,6 +110,11 @@ def close(self): case ActorContext.BACKEND: pass + @classmethod + def shared_close(cls: Type[T]): + cls._shared_instance.close() + cls._shared_instance = None + def is_alive(self): match self.context: case ActorContext.FRONTEND: @@ -130,19 +138,28 @@ def _receive(self, message: Message): try: response = getattr(self, message.method_name)(*message.args, **message.kwargs) except Exception as e: - response = e + trace = traceback.format_exc() + response = TracedError(e, trace) self._response_queue.put(response) def _send(self, name): def _send(*args, **kwargs): future = Future() def wait_for_response(future: Future): - response = self._response_queue.get() - if isinstance(response, Exception): + response = None + while response is None: + try: + response = self._response_queue.get(block=False) + except queue.Empty: + continue + if isinstance(response, TracedError): + response.base.__cause__ = Exception(response.trace) + future.set_exception(response.base) + elif isinstance(response, Exception): future.set_exception(response) else: future.set_result(response) - thread = threading.Thread(target=functools.partial(wait_for_response, future)) + thread = threading.Thread(target=wait_for_response, args=(future,), daemon=True) thread.start() self._message_queue.put(Message(name, args, kwargs)) return future diff --git a/generator_process/intents/prompt_to_image.py b/generator_process/intents/prompt_to_image.py index 35391c4f..3bb61dfb 100644 --- a/generator_process/intents/prompt_to_image.py +++ b/generator_process/intents/prompt_to_image.py @@ -254,7 +254,7 @@ def realtime_viewport_init_image(): start_schedule=1.0 * args['strength'], end_schedule=0.01, cfg_scale=args['cfg_scale'], - sampler=algorithms[args['sampler_name']], + sampler=algorithms[args['scheduler']], steps=args['steps'], seed=seed, samples=args['iterations'], diff --git a/operators/dream_texture.py b/operators/dream_texture.py index b2049fd4..8c74b6df 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -2,6 +2,7 @@ import bpy import os import numpy as np +from numpy.typing import NDArray from multiprocessing.shared_memory import SharedMemory from ..property_groups.dream_prompt import backend_options @@ -21,12 +22,6 @@ last_data_block = None timer = None -def weights_are_installed(report = None): - weights_installed = os.path.exists(WEIGHTS_PATH) - if report and (not weights_installed): - report({'ERROR'}, "Please complete setup in the preferences window.") - return weights_installed - class DreamTexture(bpy.types.Operator): bl_idname = "shade.dream_texture" bl_label = "Dream Texture" @@ -37,11 +32,6 @@ class DreamTexture(bpy.types.Operator): def poll(cls, context): return Generator.shared().can_use() - def invoke(self, context, event): - if weights_are_installed(self.report): - return self.execute(context) - return {'CANCELLED'} - def execute(self, context): history_entries = [] is_file_batch = context.scene.dream_textures_prompt.prompt_structure == file_batch_structure.id @@ -117,11 +107,16 @@ def view_step(step, width=None, height=None, shared_memory_name=None): last_data_block = step_image return # Only perform this on the first image editor found. # dream_texture(context.scene.dream_textures_prompt, view_step, image_writer) + def image_done(future): + image: NDArray = future.result() + image = bpy_image("diffusers-image", image.shape[0], image.shape[1], image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = image Generator.shared().prompt_to_image( Pipeline.STABLE_DIFFUSION, - optimizations=Optimizations(), **scene.dream_textures_prompt.generate_args(), - ) + ).add_done_callback(image_done) return {"FINISHED"} headless_prompt = None @@ -299,7 +294,7 @@ def modal_stopped(context): last_data_block = None def kill_generator(context=bpy.context): - Generator.shared().close() + Generator.shared_close() modal_stopped(context) class ReleaseGenerator(bpy.types.Operator): diff --git a/operators/view_history.py b/operators/view_history.py index 81a27963..854518c3 100644 --- a/operators/view_history.py +++ b/operators/view_history.py @@ -2,7 +2,7 @@ from bpy_extras.io_utils import ImportHelper, ExportHelper import json import os -from ..property_groups.dream_prompt import sampler_options +from ..property_groups.dream_prompt import scheduler_options from ..preferences import StableDiffusionPreferences class SCENE_UL_HistoryList(bpy.types.UIList): @@ -19,7 +19,7 @@ def draw_item(self, context, layout, data, item, icon, active_data, active_propn layout.label(text=f"{item.seed}", translate=False) layout.label(text=f"{item.width}x{item.height}", translate=False) layout.label(text=f"{item.steps} steps", translate=False) - layout.label(text=next(x for x in sampler_options if x[0] == item.sampler_name)[1], translate=False) + layout.label(text=next(x for x in scheduler_options if x[0] == item.scheduler)[1], translate=False) elif self.layout_type == 'GRID': layout.alignment = 'CENTER' layout.label(text="", icon_value=icon) diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index fea3ef92..ae74b38a 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -1,20 +1,15 @@ import bpy from bpy.props import FloatProperty, IntProperty, EnumProperty, BoolProperty, StringProperty import os +import sys +from typing import _AnnotatedAlias from ..absolute_path import absolute_path, WEIGHTS_PATH from ..generator_process.registrar import BackendTarget +from ..generator_process.actions.prompt_to_image import Optimizations, Scheduler +from ..generator_process import Generator from ..prompt_engineering import * -sampler_options = [ - ("ddim", "DDIM", "", 1), - ("plms", "PLMS", "", 2), - ("k_lms", "KLMS", "", 3), - ("k_dpm_2", "KDPM_2", "", 4), - ("k_dpm_2_a", "KDPM_2A", "", 5), - ("k_euler", "KEULER", "", 6), - ("k_euler_a", "KEULER_A", "", 7), - ("k_heun", "KHEUN", "", 8), -] +scheduler_options = [(scheduler.value, scheduler.value, '') for scheduler in Scheduler] precision_options = [ ('auto', 'Automatic', "", 1), @@ -53,8 +48,13 @@ def inpaint_mask_sources_filtered(self, context): ('xy', 'Both', '', 3), ] +_model_options = [] +def _on_model_options(future): + global _model_options + _model_options = future.result() +Generator.shared().hf_list_installed_models().add_done_callback(_on_model_options) def model_options(self, context): - return [(f, f, '', i) for i, f in enumerate(filter(lambda f: f.endswith('.ckpt'), os.listdir(WEIGHTS_PATH)))] + return [(m.id, os.path.basename(m.id).replace('models--', '').replace('--', '/'), '', i) for i, m in enumerate(_model_options)] def backend_options(self, context): def options(): @@ -94,11 +94,10 @@ def seed_clamp(self, ctx): "show_advanced": BoolProperty(name="", default=False), "random_seed": BoolProperty(name="Random Seed", default=True, description="Randomly pick a seed"), "seed": StringProperty(name="Seed", default="0", description="Manually pick a seed", update=seed_clamp), - "precision": EnumProperty(name="Precision", items=precision_options, default='auto', description="Whether to use full precision or half precision floats. Full precision is slower, but required by some GPUs"), "iterations": IntProperty(name="Iterations", default=1, min=1, description="How many images to generate"), "steps": IntProperty(name="Steps", default=25, min=1), "cfg_scale": FloatProperty(name="CFG Scale", default=7.5, min=1, soft_min=1.01, description="How strongly the prompt influences the image"), - "sampler_name": EnumProperty(name="Sampler", items=sampler_options, default=3), + "scheduler": EnumProperty(name="Scheduler", items=scheduler_options, default=0), "show_steps": BoolProperty(name="Show Steps", description="Displays intermediate steps in the Image Viewer. Disabling can speed up generation", default=False), # Init Image @@ -123,6 +122,26 @@ def seed_clamp(self, ctx): "outpaint_blend": IntProperty(name="Blend", description="Gaussian blur amount to apply to the extended area", default=16, min=0), } +default_optimizations = Optimizations() +for optim in dir(Optimizations): + if optim.startswith('_'): + continue + if hasattr(Optimizations.__annotations__, optim): + annotation = Optimizations.__annotations__[optim] + if annotation != bool or (annotation is _AnnotatedAlias and annotation.__origin__ != bool): + continue + default = getattr(default_optimizations, optim, None) + if default is not None and not isinstance(getattr(default_optimizations, optim), bool): + continue + setattr(default_optimizations, optim, True) + if default_optimizations.can_use(optim, "mps" if sys.platform == "darwin" else "cuda"): + attributes[f"optimizations_{optim}"] = BoolProperty(name=optim.replace('_', ' ').title(), default=default) +attributes["optimizations_attention_slice_size_src"] = EnumProperty(name="Attention Slice Size", items=( + ("auto", "Automatic", "", 1), + ("manual", "Manual", "", 2), +), default=1) +attributes["optimizations_attention_slice_size"] = IntProperty(name="Attention Slice Size", default=1) + def map_structure_token_items(value): return (value[0], value[1], '') for structure in prompt_structures: @@ -177,13 +196,26 @@ def get_seed(self): h = ~h return (h & 0xFFFFFFFF) ^ (h >> 32) # 64 bit hash down to 32 bits +def get_optimizations(self: DreamPrompt): + optimizations = Optimizations() + for prop in dir(self): + split_name = prop.replace('optimizations_', '') + if prop.startswith('optimizations_') and hasattr(optimizations, split_name): + setattr(optimizations, split_name, getattr(self, prop)) + if self.optimizations_attention_slice_size_src == 'auto': + optimizations.attention_slice_size = 'auto' + return optimizations + def generate_args(self): args = { key: getattr(self, key) for key in DreamPrompt.__annotations__ } args['prompt'] = self.generate_prompt() args['seed'] = self.get_seed() + args['optimizations'] = self.get_optimizations() + args['scheduler'] = Scheduler(args['scheduler']) return args DreamPrompt.generate_prompt = generate_prompt DreamPrompt.get_prompt_subject = get_prompt_subject DreamPrompt.get_seed = get_seed +DreamPrompt.get_optimizations = get_optimizations DreamPrompt.generate_args = generate_args diff --git a/render_pass.py b/render_pass.py index 1f1d8678..a1963ee3 100644 --- a/render_pass.py +++ b/render_pass.py @@ -9,7 +9,7 @@ from .generator_process import Generator -from .operators.dream_texture import dream_texture, weights_are_installed +from .operators.dream_texture import dream_texture update_render_passes_original = cycles.CyclesRender.update_render_passes render_original = cycles.CyclesRender.render @@ -38,8 +38,6 @@ def render(self, depsgraph): if size_x % 64 != 0 or size_y % 64 != 0: self.report({"ERROR"}, f"Image dimensions must be multiples of 64 (e.x. 512x512, 512x768, ...) closest is {round(size_x/64)*64}x{round(size_y/64)*64}") return result - if not weights_are_installed(self.report): - return result render_result = self.begin_result(0, 0, size_x, size_y) for original_layer in original_result.layers: layer = None diff --git a/requirements.txt b/requirements.txt index 6c76a587..72eeb0af 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ diffusers transformers +accelerate huggingface_hub torch>=1.13 stability-sdk==0.2.6 diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 5832d925..60842a2c 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -17,6 +17,7 @@ from ..space_types import SPACE_TYPES from ...property_groups.dream_prompt import DreamPrompt, backend_options from ...generator_process.registrar import BackendTarget +from ...generator_process.actions.prompt_to_image import Optimizations def dream_texture_panels(): for space_type in SPACE_TYPES: @@ -63,7 +64,7 @@ def get_prompt(context): yield from create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, prompt_panel, get_prompt) yield create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, size_panel, get_prompt) yield from create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, init_image_panels, get_prompt) - yield create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, advanced_panel, get_prompt) + yield from create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, advanced_panel, get_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): @@ -255,16 +256,61 @@ def draw(self, context): layout = self.layout layout.use_property_split = True - layout.prop(get_prompt(context), "precision") layout.prop(get_prompt(context), "random_seed") if not get_prompt(context).random_seed: layout.prop(get_prompt(context), "seed") # advanced_box.prop(self, "iterations") # Disabled until supported by the addon. layout.prop(get_prompt(context), "steps") layout.prop(get_prompt(context), "cfg_scale") - layout.prop(get_prompt(context), "sampler_name") + layout.prop(get_prompt(context), "scheduler") layout.prop(get_prompt(context), "show_steps") - return AdvancedPanel + + yield AdvancedPanel + + class SpeedOptimizationPanel(sub_panel): + """Create a subpanel for speed optimizations""" + bl_idname = f"DREAM_PT_dream_panel_speed_optimizations_{space_type}" + bl_label = "Speed Optimizations" + bl_parent_id = AdvancedPanel.bl_idname + + def draw(self, context): + super().draw(context) + layout = self.layout + layout.use_property_split = True + prompt = get_prompt(context) + + def optimization(prop): + if hasattr(prompt, f"optimizations_{prop}"): + layout.prop(prompt, f"optimizations_{prop}") + + optimization("inference_mode") + optimization("cudnn_benchmark") + optimization("tf32") + optimization("amp") + optimization("half_precision") + yield SpeedOptimizationPanel + + class MemoryOptimizationPanel(sub_panel): + """Create a subpanel for memory optimizations""" + bl_idname = f"DREAM_PT_dream_panel_memory_optimizations_{space_type}" + bl_label = "Memory Optimizations" + bl_parent_id = AdvancedPanel.bl_idname + + def draw(self, context): + super().draw(context) + layout = self.layout + layout.use_property_split = True + prompt = get_prompt(context) + + layout.prop(prompt, "optimizations_attention_slicing") + slice_size_row = layout.row() + slice_size_row.prop(prompt, "optimizations_attention_slice_size_src") + if prompt.optimizations_attention_slice_size_src == 'manual': + slice_size_row.prop(prompt, "optimizations_attention_slice_size", text="Size") + layout.prop(prompt, "optimizations_sequential_cpu_offload") + layout.prop(prompt, "optimizations_channels_last_memory_format") + # layout.prop(prompt, "optimizations_xformers_attention") # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. + yield MemoryOptimizationPanel def actions_panel(sub_panel, space_type, get_prompt): class ActionsPanel(sub_panel): diff --git a/ui/panels/render_properties.py b/ui/panels/render_properties.py index 998198c5..b18201c8 100644 --- a/ui/panels/render_properties.py +++ b/ui/panels/render_properties.py @@ -39,7 +39,7 @@ def get_prompt(context): region_type = RenderPropertiesPanel.bl_region_type panels = [ *create_panel(space_type, region_type, RenderPropertiesPanel.bl_idname, prompt_panel, get_prompt, True), - create_panel(space_type, region_type, RenderPropertiesPanel.bl_idname, advanced_panel, get_prompt, True), + *create_panel(space_type, region_type, RenderPropertiesPanel.bl_idname, advanced_panel, get_prompt, True), ] for panel in panels: def draw_decorator(original): diff --git a/ui/presets.py b/ui/presets.py index 2f7001c0..8aac830a 100644 --- a/ui/presets.py +++ b/ui/presets.py @@ -41,7 +41,7 @@ class AddAdvancedPreset(AddPresetBase, Operator): "prompt.seed", "prompt.steps", "prompt.cfg_scale", - "prompt.sampler_name", + "prompt.scheduler", "prompt.show_steps", ] From eff0fd5f9e0be643b8519bb100d54079e8abaf12 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 24 Nov 2022 13:53:29 -0500 Subject: [PATCH 03/36] Add more schedulers --- generator_process/actions/prompt_to_image.py | 27 +++++++++++++++----- requirements.txt | 11 +++++--- ui/panels/dream_texture.py | 13 +++++++--- 3 files changed, 37 insertions(+), 14 deletions(-) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index bed33e96..b002148d 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -12,9 +12,11 @@ class Pipeline(enum.IntEnum): STABILITY_SDK = 1 class Scheduler(enum.Enum): + LMS_DISCRETE = "LMS Discrete" DDIM = "DDIM" PNDM = "PNDM" - LMS_DISCRETE = "LMS Discrete" + DDPM = "DDPM" + DPM_SOLVER_MULTISTEP = "DPM Solver Multistep" EULER_DISCRETE = "Euler Discrete" EULER_ANCESTRAL_DISCRETE = "Euler Ancestral Discrete" @@ -22,18 +24,22 @@ def create(self, pipeline, pretrained): import diffusers def scheduler_class(): match self: + case Scheduler.LMS_DISCRETE: + return diffusers.LMSDiscreteScheduler case Scheduler.DDIM: return diffusers.DDIMScheduler case Scheduler.PNDM: return diffusers.PNDMScheduler - case Scheduler.LMS_DISCRETE: - return diffusers.LMSDiscreteScheduler + case Scheduler.DDPM: + return diffusers.DDPMScheduler + case Scheduler.DPM_SOLVER_MULTISTEP: + return diffusers.DPMSolverMultistepScheduler case Scheduler.EULER_DISCRETE: return diffusers.EulerDiscreteScheduler case Scheduler.EULER_ANCESTRAL_DISCRETE: return diffusers.EulerAncestralDiscreteScheduler if pretrained is not None: - return scheduler_class().from_pretrained(pretrained.model_path, subfolder=pretrained.subfolder) + return scheduler_class().from_pretrained(pretrained['model_path'], subfolder=pretrained['subfolder']) else: return scheduler_class().from_config(pipeline.scheduler.config) @@ -51,6 +57,8 @@ class Optimizations: channels_last_memory_format: bool = False # xformers_attention: bool = False # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. + cpu_only: bool = False + def can_use(self, property, device) -> bool: if not getattr(self, property): return False @@ -96,7 +104,10 @@ def prompt_to_image( from PIL import ImageOps from ...absolute_path import WEIGHTS_PATH - device = self.choose_device() + if optimizations.cpu_only: + device = "cpu" + else: + device = self.choose_device() optimizations.can_use("amp", device) if optimizations.can_use("cudnn_benchmark", device): @@ -119,10 +130,12 @@ def prompt_to_image( snapshot_folder, torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32, ) - pipe.scheduler = scheduler.create(pipe, { + is_stable_diffusion_2 = 'stabilityai--stable-diffusion-2' in snapshot_folder + pipe.scheduler = (Scheduler.EULER_DISCRETE if is_stable_diffusion_2 else scheduler).create(pipe, { 'model_path': snapshot_folder, 'subfolder': 'scheduler', - } if ('stable-diffusion-2' in snapshot_folder) else None) + } if is_stable_diffusion_2 else None) + pipe = pipe.to(device) if optimizations.can_use("attention_slicing", device): diff --git a/requirements.txt b/requirements.txt index 72eeb0af..ab3914df 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,12 @@ -diffusers +git+https://github.com/huggingface/diffusers@fix-sd2-attention-slice#egg=diffusers # SD2 attention slicing fix transformers accelerate huggingface_hub + torch>=1.13 -stability-sdk==0.2.6 -opencolorio \ No newline at end of file + +scipy # LMSDiscreteScheduler + +stability-sdk==0.2.6 # DreamStudio + +opencolorio # color management \ No newline at end of file diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 60842a2c..cc675d99 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -302,14 +302,19 @@ def draw(self, context): layout.use_property_split = True prompt = get_prompt(context) - layout.prop(prompt, "optimizations_attention_slicing") + def optimization(prop): + if hasattr(prompt, f"optimizations_{prop}"): + layout.prop(prompt, f"optimizations_{prop}") + + optimization("attention_slicing") slice_size_row = layout.row() slice_size_row.prop(prompt, "optimizations_attention_slice_size_src") if prompt.optimizations_attention_slice_size_src == 'manual': slice_size_row.prop(prompt, "optimizations_attention_slice_size", text="Size") - layout.prop(prompt, "optimizations_sequential_cpu_offload") - layout.prop(prompt, "optimizations_channels_last_memory_format") - # layout.prop(prompt, "optimizations_xformers_attention") # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. + optimization("sequential_cpu_offload") + optimization("channels_last_memory_format") + optimization("cpu_only") + # optimization("xformers_attention") # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. yield MemoryOptimizationPanel def actions_panel(sub_panel, space_type, get_prompt): From 3e6e92d956fa8ce23fa8d0bf14d09b05a68c9252 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 24 Nov 2022 15:01:33 -0500 Subject: [PATCH 04/36] Add ocio_transform action --- generator_process/__init__.py | 3 +- generator_process/actions/ocio_transform.py | 98 ++++++++++++++ .../intents/apply_ocio_transforms.py | 120 ------------------ preferences.py | 24 ++-- property_groups/dream_prompt.py | 2 +- render_pass.py | 64 ++-------- 6 files changed, 125 insertions(+), 186 deletions(-) create mode 100644 generator_process/actions/ocio_transform.py delete mode 100644 generator_process/intents/apply_ocio_transforms.py diff --git a/generator_process/__init__.py b/generator_process/__init__.py index 1ba3fcef..835c7997 100644 --- a/generator_process/__init__.py +++ b/generator_process/__init__.py @@ -7,4 +7,5 @@ class Generator(Actor): from .actions.prompt_to_image import prompt_to_image, choose_device from .actions.upscale import upscale - from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models \ No newline at end of file + from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models + from .actions.ocio_transform import ocio_transform \ No newline at end of file diff --git a/generator_process/actions/ocio_transform.py b/generator_process/actions/ocio_transform.py new file mode 100644 index 00000000..35ec8cfc --- /dev/null +++ b/generator_process/actions/ocio_transform.py @@ -0,0 +1,98 @@ +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/intents/apply_ocio_transforms.py b/generator_process/intents/apply_ocio_transforms.py deleted file mode 100644 index e4c10bcd..00000000 --- a/generator_process/intents/apply_ocio_transforms.py +++ /dev/null @@ -1,120 +0,0 @@ -import math -import sys -import numpy as np -from ..registrar import registrar -from ..intent import Intent -from ..action import Action - -@registrar.generator_intent -def apply_ocio_transforms(self, args, image_callback, exception_callback): - self.send_intent(Intent.APPLY_OCIO_TRANSFORMS, **args) - - queue = self.queue - callbacks = { - Action.IMAGE: image_callback, - Action.EXCEPTION: exception_callback, - Action.STOPPED: lambda: None - } - - while True: - while len(queue) == 0: - yield - tup = queue.pop(0) - action = tup[0] - callbacks[action](**tup[1]) - if action in [Action.STOPPED, Action.EXCEPTION]: - return - -@registrar.intent_backend(Intent.APPLY_OCIO_TRANSFORMS) -def _apply_ocio_transforms(self): - args = yield - import PyOpenColorIO as OCIO - from multiprocessing.shared_memory import SharedMemory - while True: - ocio_config = OCIO.Config.CreateFromFile(args['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, args['exposure']) - exponent = 1 if args['gamma'] == 1 else (1 / (args['gamma'] if args['gamma'] > sys.float_info.epsilon else sys.float_info.epsilon)) - processor = create_display_processor(ocio_config, OCIO.ROLE_SCENE_LINEAR, args['view_transform'], args['display_device'], args['look'] if args['look'] != 'None' else None, scale, exponent, args['inverse']) - - shared_memory = SharedMemory(args['name']) - input_image = np.frombuffer(shared_memory.buf, dtype=np.float32) - processor.getDefaultCPUProcessor().applyRGBA(input_image) - shared_memory.buf[:] = input_image.tobytes() - self.send_action(Action.IMAGE, shared_memory_name=args['name'], seed=args['name'], width=0, height=0) - del input_image - shared_memory.close() - args = yield \ No newline at end of file diff --git a/preferences.py b/preferences.py index 7a957d4d..874827d8 100644 --- a/preferences.py +++ b/preferences.py @@ -101,12 +101,10 @@ def draw_item(self, context, layout, data, item, icon, active_data, active_propn split.label(text=str(item.likes), icon="HEART") layout.operator(InstallModel.bl_idname, text="", icon="FILE_FOLDER" if is_installed else "IMPORT").model = item.model -model_installing = False - @staticmethod -def set_model_list(model_list: str, future: Future): +def set_model_list(model_list: str, models: list): getattr(bpy.context.preferences.addons[__package__].preferences, model_list).clear() - for model in future.result(): + for model in models: m = getattr(bpy.context.preferences.addons[__package__].preferences, model_list).add() m.model = model.id m.downloads = model.downloads @@ -134,17 +132,14 @@ def execute(self, context): if os.path.exists(self.model): webbrowser.open(f"file://{self.model}") else: - global model_installing - model_installing = True - def done_installing(_): - global model_installing - model_installing = False - Generator.shared().hf_list_installed_models().add_done_callback(functools.partial(set_model_list, 'installed_models')) - Generator.shared().hf_snapshot_download(self.model, bpy.context.preferences.addons[__package__].preferences.hf_token).add_done_callback(done_installing) + _ = Generator.shared().hf_snapshot_download(self.model, bpy.context.preferences.addons[__package__].preferences.hf_token).result() + set_model_list('installed_models', Generator.shared().hf_list_installed_models().result()) return {"FINISHED"} def _model_search(self, context): - Generator.shared().hf_list_models(self.model_query).add_done_callback(functools.partial(set_model_list, 'model_results')) + def on_done(future): + set_model_list('model_results', future.result()) + Generator.shared().hf_list_models(self.model_query).add_done_callback(on_done) class StableDiffusionPreferences(bpy.types.AddonPreferences): bl_idname = __package__ @@ -163,7 +158,9 @@ class StableDiffusionPreferences(bpy.types.AddonPreferences): @staticmethod def register(): - Generator.shared().hf_list_installed_models().add_done_callback(functools.partial(set_model_list, 'installed_models')) + print("register") + set_model_list('installed_models', Generator.shared().hf_list_installed_models().result()) + print("done") def draw(self, context): layout = self.layout @@ -189,7 +186,6 @@ def draw(self, context): search_box.prop(self, "model_query", text="", icon="VIEWZOOM") if len(self.model_results) > 0: - search_box.enabled = not model_installing search_box.template_list(PREFERENCES_UL_ModelList.__name__, "dream_textures_model_results", self, "model_results", self, "active_model_result") layout.template_list(PREFERENCES_UL_ModelList.__name__, "dream_textures_installed_models", self, "installed_models", self, "active_installed_model") diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index ae74b38a..264f26d0 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -3,7 +3,7 @@ import os import sys from typing import _AnnotatedAlias -from ..absolute_path import absolute_path, WEIGHTS_PATH +from ..absolute_path import absolute_path from ..generator_process.registrar import BackendTarget from ..generator_process.actions.prompt_to_image import Optimizations, Scheduler from ..generator_process import Generator diff --git a/render_pass.py b/render_pass.py index a1963ee3..77d85eb5 100644 --- a/render_pass.py +++ b/render_pass.py @@ -49,34 +49,6 @@ def render(self, depsgraph): for pass_i in layer.passes: if pass_i.name == original_render_pass.name: render_pass = pass_i - def do_ocio_transform(event, target_pixels, target_pixels_memory, inverse): - ocio_config_path = os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio') - args = { - 'config_path': ocio_config_path, - 'name': target_pixels_memory.name, - - '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': inverse - } - def image_callback(event, shared_memory_name, seed, width, height, upscaled=False): - nonlocal target_pixels - nonlocal target_pixels_memory - target_pixels[:] = np.frombuffer(target_pixels_memory.buf, dtype=np.float32).copy().reshape((size_x * size_y, 4)) - def exception_callback(fatal, msg, trace): - print(fatal, msg, trace) - generator_advance = GeneratorProcess.shared().apply_ocio_transforms(args, functools.partial(image_callback, event), exception_callback) - def timer(): - try: - next(generator_advance) - return 0.01 - except StopIteration: - event.set() - bpy.app.timers.register(timer) if render_pass.name == "Dream Textures": self.update_stats("Dream Textures", "Starting") def image_callback(set_pixels, shared_memory_name, seed, width, height, upscaled=False): @@ -101,32 +73,24 @@ def step_callback(step, width=None, height=None, shared_memory_name=None): combined_pixels = np.empty((size_x * size_y, 4), dtype=np.float32) rect.foreach_get(combined_pixels) - event = threading.Event() - - buf = combined_pixels.tobytes() - combined_pixels_memory = SharedMemory(create=True, size=len(buf)) - combined_pixels_memory.buf[:] = buf - bpy.app.timers.register(functools.partial(do_ocio_transform, event, combined_pixels, combined_pixels_memory, False)) - event.wait() + combined_pixels = Generator.shared().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 + ).get() combined_pass_image.pixels[:] = combined_pixels.ravel() self.update_stats("Dream Textures", "Starting...") - event = threading.Event() - pixels = None - def set_pixels(npbuf): - nonlocal pixels - pixels = npbuf - def do_dream_texture_pass(): - dream_texture(scene.dream_textures_render_properties_prompt, step_callback, functools.partial(image_callback, set_pixels), combined_pass_image, width=size_x, height=size_y, show_steps=False, use_init_img_color=False) - gen = GeneratorProcess.shared(None, False) - def waiter(): - if gen.in_use: - return 0.01 - event.set() - bpy.app.timers.register(waiter) - bpy.app.timers.register(do_dream_texture_pass) - event.wait() + + pixels = Generator.shared().prompt_to_image( + **scene.dream_textures_render_properties_prompt.generate_args() + ) # Perform an inverse transform so when Blender applies its transform everything looks correct. event = threading.Event() From 2b8e728f52318b669d8da66dbb9ac720b1b0ce10 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 24 Nov 2022 15:17:18 -0500 Subject: [PATCH 05/36] Support seamless generation --- generator_process/actions/prompt_to_image.py | 45 ++++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index b002148d..5e8a37f1 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -95,6 +95,9 @@ def prompt_to_image( height: int, seed: int, + seamless: bool, + seamless_axes: list[str], + **kwargs ) -> Optional[NDArray]: match pipeline: @@ -156,12 +159,15 @@ def prompt_to_image( # First-time "warmup" pass (necessary on MPS as of diffusers 0.7.2) if device == "mps": _ = pipe(prompt, num_inference_steps=1) + + _configure_model_padding(pipe.unet, seamless, seamless_axes) + _configure_model_padding(pipe.vae, seamless, seamless_axes) generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): - i = pipe.__call__( + i = pipe( prompt=prompt, height=height, width=width, @@ -177,9 +183,42 @@ def prompt_to_image( callback=None, callback_steps=1 ).images[0] - return np.asarray(ImageOps.flip(i).convert('RGBA'), dtype=np.float32) * 1/255. + return np.asarray(ImageOps.flip(i).convert('RGBA'), dtype=np.float32) / 255. case Pipeline.STABILITY_SDK: import stability_sdk raise NotImplementedError() case _: - raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file + raise Exception(f"Unsupported pipeline {pipeline}.") + +def _conv_forward_asymmetric(self, input, weight, bias): + import torch.nn as nn + """ + Patch for Conv2d._conv_forward that supports asymmetric padding + """ + working = nn.functional.pad(input, self.asymmetric_padding[0], mode=self.asymmetric_padding_mode[0]) + working = nn.functional.pad(working, self.asymmetric_padding[1], mode=self.asymmetric_padding_mode[1]) + return nn.functional.conv2d(working, weight, bias, self.stride, nn.modules.utils._pair(0), self.dilation, self.groups) + +def _configure_model_padding(model, seamless, seamless_axes): + import torch.nn as nn + """ + Modifies the 2D convolution layers to use a circular padding mode based on the `seamless` and `seamless_axes` options. + """ + for m in model.modules(): + if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): + if seamless: + m.asymmetric_padding_mode = ( + 'circular' if ('x' in seamless_axes) else 'constant', + 'circular' if ('y' in seamless_axes) else 'constant' + ) + m.asymmetric_padding = ( + (m._reversed_padding_repeated_twice[0], m._reversed_padding_repeated_twice[1], 0, 0), + (0, 0, m._reversed_padding_repeated_twice[2], m._reversed_padding_repeated_twice[3]) + ) + m._conv_forward = _conv_forward_asymmetric.__get__(m, nn.Conv2d) + else: + m._conv_forward = nn.Conv2d._conv_forward.__get__(m, nn.Conv2d) + if hasattr(m, 'asymmetric_padding_mode'): + del m.asymmetric_padding_mode + if hasattr(m, 'asymmetric_padding'): + del m.asymmetric_padding \ No newline at end of file From ee41241a9de28317f69b337686f145471c66e0af Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 24 Nov 2022 15:48:09 -0500 Subject: [PATCH 06/36] Support cancellation --- generator_process/actions/prompt_to_image.py | 73 ++++++++++++++------ generator_process/actor.py | 13 ++++ operators/dream_texture.py | 19 +++-- 3 files changed, 76 insertions(+), 29 deletions(-) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 5e8a37f1..45b69edf 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -5,6 +5,8 @@ from contextlib import nullcontext import numpy as np from numpy.typing import NDArray +from concurrent.futures import Future +import threading class Pipeline(enum.IntEnum): STABLE_DIFFUSION = 0 @@ -79,6 +81,17 @@ def choose_device(self) -> str: else: return "cpu" +class Cancellable: + def __init__(self, iterable, future: Future): + self.iterable = iterable + self.future = future + + def __iter__(self): + for i in self.iterable: + yield i + if self.future.cancelled(): + return + def prompt_to_image( self, pipeline: Pipeline, @@ -99,7 +112,9 @@ def prompt_to_image( seamless_axes: list[str], **kwargs -) -> Optional[NDArray]: +) -> Future: + future = Future() + match pipeline: case Pipeline.STABLE_DIFFUSION: import diffusers @@ -107,6 +122,7 @@ def prompt_to_image( from PIL import ImageOps from ...absolute_path import WEIGHTS_PATH + # Top level configuration if optimizations.cpu_only: device = "cpu" else: @@ -119,6 +135,7 @@ def prompt_to_image( if optimizations.can_use("tf32", device): torch.backends.cuda.matmul.allow_tf32 = True + # StableDiffusionPipeline w/ caching if hasattr(self, "_cached_pipe") and self._cached_pipe[1] == (model, scheduler, optimizations): pipe = self._cached_pipe[0] else: @@ -160,35 +177,47 @@ def prompt_to_image( if device == "mps": _ = pipe(prompt, num_inference_steps=1) + # RNG + generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + + # Seamless _configure_model_padding(pipe.unet, seamless, seamless_axes) _configure_model_padding(pipe.vae, seamless, seamless_axes) - - generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) - with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ - (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): - i = pipe( - prompt=prompt, - height=height, - width=width, - num_inference_steps=steps, - guidance_scale=7.5, - negative_prompt=None, - num_images_per_prompt=1, - eta=0.0, - generator=generator, - latents=None, - output_type="pil", - return_dict=True, - callback=None, - callback_steps=1 - ).images[0] - return np.asarray(ImageOps.flip(i).convert('RGBA'), dtype=np.float32) / 255. + # Cancellation + pipe.progress_bar = lambda iterable: Cancellable(iterable, future) + + # Inference + def inference(): + with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + i = pipe( + prompt=prompt, + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=7.5, + negative_prompt=None, + num_images_per_prompt=1, + eta=0.0, + generator=generator, + latents=None, + output_type="pil", + return_dict=True, + callback=None, + callback_steps=1 + ).images[0] + if not future.cancelled(): + future.set_result(np.asarray(ImageOps.flip(i).convert('RGBA'), dtype=np.float32) / 255.) + t = threading.Thread(target=inference, daemon=True) + t.start() case Pipeline.STABILITY_SDK: import stability_sdk raise NotImplementedError() case _: raise Exception(f"Unsupported pipeline {pipeline}.") + + return future def _conv_forward_asymmetric(self, input, weight, bias): import torch.nn as nn diff --git a/generator_process/actor.py b/generator_process/actor.py index 3a625dbb..b8f17f8b 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -28,6 +28,8 @@ def __init__(self, method_name, args, kwargs): self.method_name = method_name self.args = args self.kwargs = kwargs + + CANCEL_MESSAGE = "__cancel__" def _start_backend(cls, message_queue, response_queue): cls( @@ -112,6 +114,8 @@ def close(self): @classmethod def shared_close(cls: Type[T]): + if cls._shared_instance is None: + return cls._shared_instance.close() cls._shared_instance = None @@ -135,8 +139,15 @@ def _backend_loop(self): self._receive(self._message_queue.get()) def _receive(self, message: Message): + if message.method_name == Message.CANCEL_MESSAGE and self._active_future is not None: + self._active_future.cancel() + return try: response = getattr(self, message.method_name)(*message.args, **message.kwargs) + if isinstance(response, Future): + self._active_future = response + response.add_done_callback(lambda future: self._response_queue.put(None if future.cancelled() else (future.exception() or future.result()))) + return except Exception as e: trace = traceback.format_exc() response = TracedError(e, trace) @@ -149,6 +160,8 @@ def wait_for_response(future: Future): response = None while response is None: try: + if future.cancelled(): + self._message_queue.put(Message(Message.CANCEL_MESSAGE, (), {})) response = self._response_queue.get(block=False) except queue.Empty: continue diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 8c74b6df..811ef6f4 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -108,15 +108,21 @@ def view_step(step, width=None, height=None, shared_memory_name=None): return # Only perform this on the first image editor found. # dream_texture(context.scene.dream_textures_prompt, view_step, image_writer) def image_done(future): + if future.cancelled(): + del gen._active_generation_future + return image: NDArray = future.result() image = bpy_image("diffusers-image", image.shape[0], image.shape[1], image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = image - Generator.shared().prompt_to_image( + gen = Generator.shared() + f = gen.prompt_to_image( Pipeline.STABLE_DIFFUSION, **scene.dream_textures_prompt.generate_args(), - ).add_done_callback(image_done) + ) + gen._active_generation_future = f + f.add_done_callback(image_done) return {"FINISHED"} headless_prompt = None @@ -315,11 +321,10 @@ class CancelGenerator(bpy.types.Operator): @classmethod def poll(self, context): - global timer - return timer is not None + gen = Generator.shared() + return hasattr(gen, "_active_generation_future") and gen._active_generation_future is not None def execute(self, context): - gen = GeneratorProcess.shared(create=False) - if gen: - gen.send_stop(Intent.PROMPT_TO_IMAGE) + gen = Generator.shared() + gen._active_generation_future.cancel() return {'FINISHED'} From b906bb229fe0e4e457a519e6eee2611ad74d7fb9 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 26 Nov 2022 12:49:47 -0500 Subject: [PATCH 07/36] Remove old intents, upscaling, actor improvements --- __init__.py | 3 +- classes.py | 3 - generator_process/actions/prompt_to_image.py | 254 +++++++++++------ generator_process/actions/upscale.py | 27 +- generator_process/actor.py | 198 ++++++++++--- generator_process/intents/prompt_to_image.py | 279 ------------------- generator_process/intents/send_stop.py | 8 - generator_process/intents/upscale.py | 62 ----- operators/dream_texture.py | 13 +- operators/upscale.py | 67 ++--- preferences.py | 6 - requirements.txt | 2 +- ui/panels/dream_texture.py | 2 +- ui/panels/upscaling.py | 48 +--- 14 files changed, 397 insertions(+), 575 deletions(-) delete mode 100644 generator_process/intents/prompt_to_image.py delete mode 100644 generator_process/intents/send_stop.py delete mode 100644 generator_process/intents/upscale.py diff --git a/__init__.py b/__init__.py index d765da7e..682bade4 100644 --- a/__init__.py +++ b/__init__.py @@ -43,7 +43,6 @@ def clear_modules(): from .tools import TOOLS from .operators.dream_texture import DreamTexture, kill_generator from .property_groups.dream_prompt import DreamPrompt - from .operators.upscale import upscale_options from .preferences import StableDiffusionPreferences from .ui.presets import register_default_presets @@ -85,7 +84,7 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_viewport_enabled = BoolProperty(name="Viewport Enabled", default=False) bpy.types.Scene.dream_textures_render_properties_enabled = BoolProperty(default=False) bpy.types.Scene.dream_textures_render_properties_prompt = PointerProperty(type=DreamPrompt) - bpy.types.Scene.dream_textures_upscale_outscale = bpy.props.EnumProperty(name="Target Size", items=upscale_options) + bpy.types.Scene.dream_textures_upscale_full_precision = bpy.props.BoolProperty(name="Full Precision", default=True) bpy.types.Scene.dream_textures_upscale_seamless = bpy.props.BoolProperty(name="Seamless", default=False) diff --git a/classes.py b/classes.py index d1ff8c8d..30c7327c 100644 --- a/classes.py +++ b/classes.py @@ -36,9 +36,6 @@ *upscaling.upscaling_panels(), *history.history_panels(), - upscaling.OpenRealESRGANDownload, - upscaling.OpenRealESRGANWeightsDirectory, - dream_texture.OpenClipSegDownload, dream_texture.OpenClipSegWeightsDirectory, ) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 45b69edf..429320bb 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -1,12 +1,10 @@ -from typing import Optional, Annotated, Union, _AnnotatedAlias +from typing import Annotated, Union, _AnnotatedAlias, Generator, Callable, List, Optional import enum import os from dataclasses import dataclass from contextlib import nullcontext -import numpy as np from numpy.typing import NDArray -from concurrent.futures import Future -import threading +import numpy as np class Pipeline(enum.IntEnum): STABLE_DIFFUSION = 0 @@ -57,7 +55,7 @@ class Optimizations: half_precision: Annotated[bool, "cuda"] = True sequential_cpu_offload: bool = False channels_last_memory_format: bool = False - # xformers_attention: bool = False # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. + # xformers_attention: bool = False # FIXME: xFormers is not yet available. cpu_only: bool = False @@ -81,17 +79,6 @@ def choose_device(self) -> str: else: return "cpu" -class Cancellable: - def __init__(self, iterable, future: Future): - self.iterable = iterable - self.future = future - - def __iter__(self): - for i in self.iterable: - yield i - if self.future.cancelled(): - return - def prompt_to_image( self, pipeline: Pipeline, @@ -112,31 +99,142 @@ def prompt_to_image( seamless_axes: list[str], **kwargs -) -> Future: - future = Future() - +) -> Generator[NDArray, None, None]: match pipeline: case Pipeline.STABLE_DIFFUSION: import diffusers import torch - from PIL import ImageOps + from PIL import Image, ImageOps from ...absolute_path import WEIGHTS_PATH - # Top level configuration + # Mostly copied from `diffusers.StableDiffusionPipeline`, with slight modifications to yield the latents at each step. + class GeneratorPipeline(diffusers.StableDiffusionPipeline): + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[torch.Generator] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + **kwargs, + ): + # 0. Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # 2. Define call parameters + batch_size = 1 if isinstance(prompt, str) else len(prompt) + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_embeddings = self._encode_prompt( + prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt + ) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Denoising loop + for i, t in enumerate(self.progress_bar(timesteps)): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # NOTE: Modified to yield the latents instead of calling a callback. + yield np.asarray(ImageOps.flip(Image.fromarray(self._approximate_decoded_latents(latents))).convert('RGBA'), dtype=np.float32) / 255. + + # 8. Post-processing + image = self.decode_latents(latents) + + # TODO: Add UI to enable this. + # 9. Run safety checker + # image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype) + + # NOTE: Modified to yield the decoded image as a numpy array. + yield from [ + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for image in self.numpy_to_pil(image) + ] + + + def _approximate_decoded_latents(self, latents): + """ + Approximate the decoded latents without using the VAE. + """ + # origingally adapted from code by @erucipe and @keturn here: + # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7 + + # these updated numbers for v1.5 are from @torridgristle + v1_5_latent_rgb_factors = torch.tensor([ + # R G B + [ 0.3444, 0.1385, 0.0670], # L1 + [ 0.1247, 0.4027, 0.1494], # L2 + [-0.3192, 0.2513, 0.2103], # L3 + [-0.1307, -0.1874, -0.7445] # L4 + ], 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() + if optimizations.cpu_only: device = "cpu" else: device = self.choose_device() - optimizations.can_use("amp", device) - - if optimizations.can_use("cudnn_benchmark", device): - torch.backends.cudnn.benchmark = True - - if optimizations.can_use("tf32", device): - torch.backends.cuda.matmul.allow_tf32 = True + + use_cpu_offload = optimizations.can_use("sequential_cpu_offload", device) # StableDiffusionPipeline w/ caching - if hasattr(self, "_cached_pipe") and self._cached_pipe[1] == (model, scheduler, optimizations): + if hasattr(self, "_cached_pipe") and self._cached_pipe[1] == model and use_cpu_offload == self._cached_pipe[2]: pipe = self._cached_pipe[0] else: storage_folder = os.path.join(WEIGHTS_PATH, model) @@ -146,36 +244,42 @@ def prompt_to_image( commit_hash = f.read() snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) - pipe = diffusers.StableDiffusionPipeline.from_pretrained( + pipe = GeneratorPipeline.from_pretrained( snapshot_folder, + revision="fp16" if optimizations.can_use("half_precision", device) else None, torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32, ) - is_stable_diffusion_2 = 'stabilityai--stable-diffusion-2' in snapshot_folder - pipe.scheduler = (Scheduler.EULER_DISCRETE if is_stable_diffusion_2 else scheduler).create(pipe, { - 'model_path': snapshot_folder, - 'subfolder': 'scheduler', - } if is_stable_diffusion_2 else None) - pipe = pipe.to(device) + setattr(self, "_cached_pipe", (pipe, model, use_cpu_offload, snapshot_folder)) - if optimizations.can_use("attention_slicing", device): - pipe.enable_attention_slicing(optimizations.attention_slice_size) - - if optimizations.can_use("sequential_cpu_offload", device): - pipe.enable_sequential_cpu_offload() - - if optimizations.can_use("channels_last_memory_format", device): - pipe.unet.to(memory_format=torch.channels_last) + # Scheduler + is_stable_diffusion_2 = 'stabilityai--stable-diffusion-2' in model + pipe.scheduler = scheduler.create(pipe, { + 'model_path': self._cached_pipe[3], + 'subfolder': 'scheduler', + } if is_stable_diffusion_2 else None) - # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. - # if optimizations.can_use("xformers_attention", device): - # pipe.enable_xformers_memory_efficient_attention() + # Optimizations + + torch.backends.cudnn.benchmark = optimizations.can_use("cudnn_benchmark", device) + torch.backends.cuda.matmul.allow_tf32 = optimizations.can_use("tf32", device) - setattr(self, "_cached_pipe", (pipe, (model, scheduler, optimizations))) + if optimizations.can_use("attention_slicing", device): + pipe.enable_attention_slicing(optimizations.attention_slice_size) + else: + pipe.disable_attention_slicing() + + if use_cpu_offload: + pipe.enable_sequential_cpu_offload() + + if optimizations.can_use("channels_last_memory_format", device): + pipe.unet.to(memory_format=torch.channels_last) + else: + pipe.unet.to(memory_format=torch.contiguous_format) - # First-time "warmup" pass (necessary on MPS as of diffusers 0.7.2) - if device == "mps": - _ = pipe(prompt, num_inference_steps=1) + # FIXME: xFormers wheels are not yet available (https://github.com/facebookresearch/xformers/issues/533) + # if optimizations.can_use("xformers_attention", device): + # pipe.enable_xformers_memory_efficient_attention() # RNG generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) @@ -184,40 +288,30 @@ def prompt_to_image( _configure_model_padding(pipe.unet, seamless, seamless_axes) _configure_model_padding(pipe.vae, seamless, seamless_axes) - # Cancellation - pipe.progress_bar = lambda iterable: Cancellable(iterable, future) - # Inference - def inference(): - with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ - (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): - i = pipe( - prompt=prompt, - height=height, - width=width, - num_inference_steps=steps, - guidance_scale=7.5, - negative_prompt=None, - num_images_per_prompt=1, - eta=0.0, - generator=generator, - latents=None, - output_type="pil", - return_dict=True, - callback=None, - callback_steps=1 - ).images[0] - if not future.cancelled(): - future.set_result(np.asarray(ImageOps.flip(i).convert('RGBA'), dtype=np.float32) / 255.) - t = threading.Thread(target=inference, daemon=True) - t.start() + with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + yield from pipe( + prompt=prompt, + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=7.5, + negative_prompt=None, + num_images_per_prompt=1, + eta=0.0, + generator=generator, + latents=None, + output_type="pil", + return_dict=True, + callback=None, + callback_steps=1 + ) case Pipeline.STABILITY_SDK: import stability_sdk raise NotImplementedError() case _: raise Exception(f"Unsupported pipeline {pipeline}.") - - return future def _conv_forward_asymmetric(self, input, weight, bias): import torch.nn as nn @@ -250,4 +344,4 @@ def _configure_model_padding(model, seamless, seamless_axes): if hasattr(m, 'asymmetric_padding_mode'): del m.asymmetric_padding_mode if hasattr(m, 'asymmetric_padding'): - del m.asymmetric_padding \ No newline at end of file + del m.asymmetric_padding diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index 1cf08fee..4574c66f 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -1,5 +1,26 @@ +from numpy.typing import NDArray +import numpy as np + def upscale( self, - input: bytes -) -> bytes: - return input \ No newline at end of file + image: str, + prompt: str, + + half_precision: bool +) -> NDArray: + import torch + import diffusers + from PIL import Image, ImageOps + + model_id = "stabilityai/stable-diffusion-x4-upscaler" + pipe = diffusers.StableDiffusionUpscalePipeline.from_pretrained( + model_id, + revision="fp16" if half_precision else None, + torch_dtype=torch.float16 if half_precision else torch.float32 + ) + pipe = pipe.to(self.choose_device()) + + pipe.enable_attention_slicing() + + result = pipe(prompt=prompt, image=Image.open(image).convert('RGB').resize((128, 128))).images[0] + return np.asarray(ImageOps.flip(result).convert('RGBA'), dtype=np.float32) / 255. \ No newline at end of file diff --git a/generator_process/actor.py b/generator_process/actor.py index b8f17f8b..10e82c11 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -1,12 +1,121 @@ -from multiprocessing import Queue, Process -import queue +from multiprocessing import Queue, Process, Lock +import multiprocessing.synchronize import enum import traceback import threading -from typing import Type, TypeVar -from concurrent.futures import Future +from typing import Type, TypeVar, Callable, Any, MutableSet, Generator +# from concurrent.futures import Future import site +class Future: + """ + Object that represents a value that has not completed processing, but will in the future. + + Add callbacks to be notified when values become available, or use `.result()` and `.exception()` to wait for the value. + """ + _response_callbacks: MutableSet[Callable[['Future', Any], None]] = set() + _exception_callbacks: MutableSet[Callable[['Future', BaseException], None]] = set() + _done_callbacks: MutableSet[Callable[['Future'], None]] = set() + _responses: list = [] + _exception: BaseException | None = None + done: bool = False + cancelled: bool = False + + def __init__(self): + self._response_callbacks = set() + self._exception_callbacks = set() + self._done_callbacks = set() + self._responses = [] + self._exception = None + self.done = False + self.cancelled = False + + def result(self): + """ + Get the result value (blocking). + """ + def _response(): + match len(self._responses): + case 0: + return None + case 1: + return self._responses[0] + case _: + return self._responses + if self._exception is not None: + raise self._exception + if self.done: + return _response() + else: + event = threading.Event() + def _done(_): + event.set() + self.add_done_callback(_done) + event.wait() + if self._exception is not None: + raise self._exception + return _response() + + def exception(self): + if self.done: + return self._exception + else: + event = threading.Event() + def _done(_): + event.set() + self.add_done_callback(_done) + event.wait() + return self._exception + + def cancel(self): + self.done = True + self._cancelled = True + self.set_done() + + def add_response(self, response): + """ + Add a response value and notify all consumers. + """ + self._responses.append(response) + for response_callback in self._response_callbacks: + response_callback(self, response) + + def set_exception(self, exception: BaseException): + """ + Set the exception. + """ + self._exception = exception + + def set_done(self): + """ + Mark the future as done. + """ + assert not self.done + self.done = True + for done_callback in self._done_callbacks: + done_callback(self) + + def add_response_callback(self, callback: Callable[['Future', Any], None]): + """ + Add a callback to run whenever a response is received. + Will be called multiple times by generator functions. + """ + self._response_callbacks.add(callback) + + def add_exception_callback(self, callback: Callable[['Future', BaseException], None]): + """ + Add a callback to run when the future errors. + Will only be called once at the first exception. + """ + self._exception_callbacks.add(callback) + + def add_done_callback(self, callback: Callable[['Future'], None]): + """ + Add a callback to run when the future is marked as done. + Will only be called once. + """ + self._done_callbacks.add(callback) + class ActorContext(enum.IntEnum): """ The context of an `Actor` object. @@ -17,7 +126,7 @@ class ActorContext(enum.IntEnum): FRONTEND = 0 BACKEND = 1 -class Message(): +class Message: """ Represents a function signature with a method name, positonal arguments, and keyword arguments. @@ -29,7 +138,8 @@ def __init__(self, method_name, args, kwargs): self.args = args self.kwargs = kwargs - CANCEL_MESSAGE = "__cancel__" + CANCEL = "__cancel__" + END = "__end__" def _start_backend(cls, message_queue, response_queue): cls( @@ -45,15 +155,19 @@ def __init__(self, base: BaseException, trace: str): T = TypeVar('T', bound='Actor') -class Actor(): +class Actor: """ Base class for specialized actors. - Uses queues to serialize actions from different threads, and automatically dispatches methods to a separate process. + Uses queues to send actions to a background process and receive a response. + Calls to any method declared by the frontend are automatically dispatched to the backend. + + All function arguments must be picklable. """ _message_queue: Queue _response_queue: Queue + _lock: multiprocessing.synchronize.Lock _shared_instance = None @@ -66,7 +180,7 @@ class Actor(): "shared" } - def __init__(self, context: ActorContext, message_queue: Queue = Queue(), response_queue: Queue = Queue()): + def __init__(self, context: ActorContext, message_queue: Queue = Queue(maxsize=1), response_queue: Queue = Queue(maxsize=1)): self.context = context self._message_queue = message_queue self._response_queue = response_queue @@ -79,6 +193,7 @@ def _setup(self): """ match self.context: case ActorContext.FRONTEND: + self._lock = Lock() for name in filter(lambda name: callable(getattr(self, name)) and not name.startswith("_") and name not in self._protected_methods, dir(self)): setattr(self, name, self._send(name)) case ActorContext.BACKEND: @@ -127,7 +242,9 @@ def is_alive(self): return True def can_use(self): - return self._message_queue.empty() and self._response_queue.empty() + if result := self._lock.acquire(block=False): + self._lock.release() + return result def _load_dependencies(self): from ..absolute_path import absolute_path @@ -139,42 +256,49 @@ def _backend_loop(self): self._receive(self._message_queue.get()) def _receive(self, message: Message): - if message.method_name == Message.CANCEL_MESSAGE and self._active_future is not None: - self._active_future.cancel() - return try: response = getattr(self, message.method_name)(*message.args, **message.kwargs) - if isinstance(response, Future): - self._active_future = response - response.add_done_callback(lambda future: self._response_queue.put(None if future.cancelled() else (future.exception() or future.result()))) - return + if isinstance(response, Generator): + for res in iter(response): + extra_message = None + try: + self._message_queue.get(block=False) + except: + pass + if extra_message == Message.CANCEL: + break + self._response_queue.put(res) + else: + self._response_queue.put(response) except Exception as e: trace = traceback.format_exc() - response = TracedError(e, trace) - self._response_queue.put(response) + self._response_queue.put(TracedError(e, trace)) + self._response_queue.put(Message.END) def _send(self, name): def _send(*args, **kwargs): future = Future() - def wait_for_response(future: Future): - response = None - while response is None: - try: - if future.cancelled(): - self._message_queue.put(Message(Message.CANCEL_MESSAGE, (), {})) - response = self._response_queue.get(block=False) - except queue.Empty: - continue - if isinstance(response, TracedError): - response.base.__cause__ = Exception(response.trace) - future.set_exception(response.base) - elif isinstance(response, Exception): - future.set_exception(response) - else: - future.set_result(response) - thread = threading.Thread(target=wait_for_response, args=(future,), daemon=True) + def _send_thread(future: Future): + self._lock.acquire() + self._message_queue.put(Message(name, args, kwargs)) + + while not future.done: + if future.cancelled: + self._message_queue.put(Message.CANCEL) + response = self._response_queue.get() + if response == Message.END: + future.set_done() + elif isinstance(response, TracedError): + response.base.__cause__ = Exception(response.trace) + future.set_exception(response.base) + elif isinstance(response, Exception): + future.set_exception(response) + else: + future.add_response(response) + + self._lock.release() + thread = threading.Thread(target=_send_thread, args=(future,), daemon=True) thread.start() - self._message_queue.put(Message(name, args, kwargs)) return future return _send diff --git a/generator_process/intents/prompt_to_image.py b/generator_process/intents/prompt_to_image.py deleted file mode 100644 index 3bb61dfb..00000000 --- a/generator_process/intents/prompt_to_image.py +++ /dev/null @@ -1,279 +0,0 @@ -from ..block_in_use import block_in_use -from ..action import Action -from ..intent import Intent -from ..registrar import BackendTarget, registrar -import sys -import os - -@registrar.generator_intent -@block_in_use -def prompt2image(self, args, step_callback, image_callback, info_callback, exception_callback): - self.send_intent(Intent.PROMPT_TO_IMAGE, **args) - - queue = self.queue - callbacks = { - Action.INFO: info_callback, - Action.IMAGE: image_callback, - Action.STEP_IMAGE: step_callback, - Action.STEP_NO_SHOW: step_callback, - Action.EXCEPTION: exception_callback, - Action.STOPPED: lambda: None - } - - while True: - while len(queue) == 0: - yield # nothing in queue, let blender resume - tup = queue.pop(0) - action = tup[0] - callbacks[action](**tup[1]) - if action in [Action.STOPPED, Action.EXCEPTION]: - return - -@registrar.intent_backend(Intent.PROMPT_TO_IMAGE, BackendTarget.LOCAL) -def prompt_to_image(self): - args = yield - self.send_info("Importing Dependencies") - from absolute_path import absolute_path - from stable_diffusion.ldm.generate import Generate - from stable_diffusion.ldm.invoke.devices import choose_precision - from io import StringIO - - models_config = absolute_path('weights/config.yml') - generator: Generate = None - - def image_writer(image, seed, upscaled=False, first_seed=None): - # Only use the non-upscaled texture, as upscaling is a separate step in this addon. - if not upscaled: - self.send_action(Action.IMAGE, shared_memory_name=self.share_image_memory(image), seed=seed, width=image.width, height=image.height) - - step = 0 - def view_step(samples, i): - self.check_stop() - nonlocal step - step = i - if args['show_steps']: - image = generator.sample_to_image(samples) - self.send_action(Action.STEP_IMAGE, shared_memory_name=self.share_image_memory(image), step=step, width=image.width, height=image.height) - else: - self.send_action(Action.STEP_NO_SHOW, step=step) - - def preload_models(): - tqdm = None - try: - from huggingface_hub.utils.tqdm import tqdm as hfh_tqdm - tqdm = hfh_tqdm - except: - try: - from tqdm.auto import tqdm as auto_tqdm - tqdm = auto_tqdm - except: - return - - current_model_name = "" - def start_preloading(model_name): - nonlocal current_model_name - current_model_name = model_name - self.send_info(f"Loading {model_name}") - - def update_decorator(original): - def update(tqdm_self, n=1): - result = original(tqdm_self, n) - nonlocal current_model_name - frac = tqdm_self.n / tqdm_self.total - percentage = int(frac * 100) - if tqdm_self.n - tqdm_self.last_print_n >= tqdm_self.miniters: - self.send_info(f"Downloading {current_model_name} ({percentage}%)") - return result - return update - old_update = tqdm.update - tqdm.update = update_decorator(tqdm.update) - - import warnings - import transformers - transformers.logging.set_verbosity_error() - - start_preloading("BERT tokenizer") - transformers.BertTokenizerFast.from_pretrained('bert-base-uncased') - - self.send_info("Preloading `kornia` requirements") - with warnings.catch_warnings(): - warnings.filterwarnings('ignore', category=DeprecationWarning) - import kornia - - start_preloading("CLIP") - clip_version = 'openai/clip-vit-large-patch14' - transformers.CLIPTokenizer.from_pretrained(clip_version) - transformers.CLIPTextModel.from_pretrained(clip_version) - - start_preloading("CLIP Segmentation") - from absolute_path import CLIPSEG_WEIGHTS_PATH - from clipseg_models.clipseg import CLIPDensePredT - CLIPDensePredT(version='ViT-B/16', reduce_dim=64) - - tqdm.update = old_update - - from transformers.utils.hub import TRANSFORMERS_CACHE - model_paths = {'bert-base-uncased', 'openai--clip-vit-large-patch14'} - if any(not os.path.isdir(os.path.join(TRANSFORMERS_CACHE, f'models--{path}')) for path in model_paths) or not os.path.exists(os.path.join(os.path.expanduser("~/.cache/clip"), 'ViT-B-16.pt')): - preload_models() - - while True: - try: - self.check_stop() - # Reset the step count - step = 0 - - if generator is None or generator.precision != (choose_precision(generator.device) if args['precision'] == 'auto' else args['precision']) or generator.model_name != args['model']: - self.send_info("Loading Model") - from omegaconf import OmegaConf - def omegaconf_load(func): - def load(path): - if path == models_config: - return OmegaConf.create({ - args['model']: { - "config": "stable_diffusion/configs/stable-diffusion/v1-inference.yaml", - "weights": f"weights/stable-diffusion-v1.4/{args['model']}", - "description": "Stable Diffusion inference model version 1.4", - "width": 512, - "height": 512 - } - }) - else: - return func(path) - load.omegaconf_decorated = True - return load - if not getattr(OmegaConf.load, 'omegaconf_decorated', False): - OmegaConf.load = omegaconf_load(OmegaConf.load) - generator = Generate( - conf=models_config, - model=args['model'], - precision=args['precision'] - ) - generator.free_gpu_mem = False # Not sure what this is for, and why it isn't a flag but read from Args()? - generator.load_model() - self.check_stop() - self.send_info("Starting") - - tmp_stderr = sys.stderr = StringIO() # prompt2image writes exceptions straight to stderr, intercepting - prompt_list = args['prompt'] if isinstance(args['prompt'], list) else [args['prompt']] - for prompt in prompt_list: - generator_args = args.copy() - generator_args['prompt'] = prompt - generator_args['seamless_axes'] = list(generator_args['seamless_axes']) - if args['init_img_action'] == 'inpaint' and args['inpaint_mask_src'] == 'prompt': - generator_args['text_mask'] = (generator_args['text_mask'], generator_args['text_mask_confidence']) - else: - generator_args['text_mask'] = None - if args['use_init_img'] and args['init_img_action'] == 'outpaint': - args['fit'] = False - # Extend the image in the specified directions - from PIL import Image, ImageFilter - init_img = Image.open(args['init_img']) - extended_size = (init_img.size[0] + args['outpaint_left'] + args['outpaint_right'], init_img.size[1] + args['outpaint_top'] + args['outpaint_bottom']) - extended_img = Image.new('RGBA', extended_size, (0, 0, 0, 0)) - blurred_fill = init_img.resize(extended_size).filter(filter=ImageFilter.GaussianBlur(radius=args['outpaint_blend'])) - blurred_fill.putalpha(0) - extended_img.paste(blurred_fill, (0, 0)) - extended_img.paste(init_img, (args['outpaint_left'], args['outpaint_top'])) - extended_img.save(generator_args['init_img'], 'png') - generator_args['width'] = extended_size[0] - generator_args['height'] = extended_size[1] - generator.prompt2image( - # a function or method that will be called each step - step_callback=view_step, - # a function or method that will be called each time an image is generated - image_callback=image_writer, - **generator_args - ) - if tmp_stderr.tell() > 0: - tmp_stderr.seek(0) - s = tmp_stderr.read() - i = s.find("Traceback") # progress also gets printed to stderr so check for an actual exception - if i != -1: - s = s[i:] - import re - low_ram = re.search(r"(Not enough memory, use lower resolution)( \(max approx. \d+x\d+\))",s,re.IGNORECASE) - if low_ram: - self.send_exception(False, f"{low_ram[1]}{' or disable full precision' if args['precision'] == 'float32' else ''}{low_ram[2]}", s) - elif s.find("CUDA out of memory. Tried to allocate") != -1: - self.send_exception(False, f"Not enough memory, use lower resolution{' or disable full precision' if args['precision'] == 'float32' else ''}", s) - else: - self.send_exception(True, msg=None, trace=s) # consider all unknown exceptions to be fatal so the generator process is fully restarted next time - except KeyboardInterrupt: - pass - finally: - sys.stderr = self.stderr - args = yield - -@registrar.intent_backend(Intent.PROMPT_TO_IMAGE, BackendTarget.STABILITY_SDK) -def prompt_to_image_stability_sdk(self): - args = yield - self.send_info("Importing Dependencies") - - from stability_sdk import client, interfaces - from PIL import Image - import io - import random - from multiprocessing.shared_memory import SharedMemory - - # Some of these names are abbreviated. - algorithms = client.algorithms.copy() - algorithms['k_euler_a'] = algorithms['k_euler_ancestral'] - algorithms['k_dpm_2_a'] = algorithms['k_dpm_2_ancestral'] - - stability_inference = client.StabilityInference(key=args['dream_studio_key']) - - def image_writer(image, seed, upscaled=False, first_seed=None): - # Only use the non-upscaled texture, as upscaling is a separate step in this addon. - if not upscaled: - self.send_action(Action.IMAGE, shared_memory_name=self.share_image_memory(image), seed=seed, width=image.width, height=image.height) - - while True: - self.check_stop() - - self.send_info("Generating...") - - seed = random.randrange(0, 4294967295) if args['seed'] is None else args['seed'] - def realtime_viewport_init_image(): - return None # Realtime viewport is not currently available. - if args['init_img_shared_memory'] is not None: - init_img_memory = SharedMemory(args['init_img_shared_memory']) - shared_init_img = Image.frombytes('RGBA', (args['init_img_shared_memory_width'], args['init_img_shared_memory_height']), init_img_memory.buf.tobytes()) - shared_init_img = shared_init_img.resize((512, round(((shared_init_img.height / shared_init_img.width) * 512) / 64)*64)) - init_img_memory.close() - return shared_init_img - return None - shared_init_img = realtime_viewport_init_image() - init_img = Image.open(args['init_img']) if args['init_img'] is not None and args['use_init_img'] else None - answers = stability_inference.generate( - prompt=args['prompt'], - init_image=shared_init_img if shared_init_img is not None and args['use_init_img'] else init_img, - mask_image=init_img.split()[-1] if init_img is not None and args['use_init_img'] and args['init_img_action'] == 'inpaint' else None, - width=shared_init_img.width if shared_init_img is not None else args['width'], - height=shared_init_img.height if shared_init_img is not None else args['height'], - start_schedule=1.0 * args['strength'], - end_schedule=0.01, - cfg_scale=args['cfg_scale'], - sampler=algorithms[args['scheduler']], - steps=args['steps'], - seed=seed, - samples=args['iterations'], - # safety: bool = True, - # classifiers: Optional[generation.ClassifierParameters] = None, - # guidance_preset: generation.GuidancePreset = generation.GUIDANCE_PRESET_NONE, - # guidance_cuts: int = 0, - # guidance_strength: Optional[float] = None, - # guidance_prompt: Union[str, generation.Prompt] = None, - # guidance_models: List[str] = None, - ) - - for answer in answers: - for artifact in answer.artifacts: - if artifact.finish_reason == interfaces.gooseai.generation.generation_pb2.FILTER: - self.send_exception(False, "Your request activated DreamStudio's safety filter. Please modify your prompt and try again.") - if artifact.type == interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE: - response = Image.open(io.BytesIO(artifact.binary)) - image_writer(response, artifact.seed) - - self.send_info("Done") - args = yield \ No newline at end of file diff --git a/generator_process/intents/send_stop.py b/generator_process/intents/send_stop.py deleted file mode 100644 index 6a8270c6..00000000 --- a/generator_process/intents/send_stop.py +++ /dev/null @@ -1,8 +0,0 @@ -from ..intent import Intent -from ..registrar import registrar - -@registrar.generator_intent -def send_stop(self, stop_intent): - self.send_intent(Intent.STOP, stop_intent=stop_intent) - -# The Intent backend is implemented as a special case. \ No newline at end of file diff --git a/generator_process/intents/upscale.py b/generator_process/intents/upscale.py deleted file mode 100644 index 4589d994..00000000 --- a/generator_process/intents/upscale.py +++ /dev/null @@ -1,62 +0,0 @@ -from ..block_in_use import block_in_use -from ..action import Action -from ..intent import Intent -from ..registrar import registrar - -@registrar.generator_intent -@block_in_use -def upscale(self, args, image_callback, info_callback, exception_callback): - self.send_intent(Intent.UPSCALE, **args) - - queue = self.queue - callbacks = { - Action.INFO: info_callback, - Action.IMAGE: image_callback, - Action.EXCEPTION: exception_callback, - Action.STOPPED: lambda: None - } - - while True: - while len(queue) == 0: - yield - tup = queue.pop(0) - action = tup[0] - callbacks[action](**tup[1]) - if action in [Action.STOPPED, Action.EXCEPTION]: - return - -@registrar.intent_backend(Intent.UPSCALE) -def upscale(self): - args = yield - self.send_info("Starting") - from absolute_path import REAL_ESRGAN_WEIGHTS_PATH - import cv2 - from PIL import Image - from realesrgan import RealESRGANer - from realesrgan.archs.srvgg_arch import SRVGGNetCompact - from torch import nn - while True: - image = cv2.imread(args['input'], cv2.IMREAD_UNCHANGED) - real_esrgan_model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu') - if args['seamless']: - for m in real_esrgan_model.body: - if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): - m.padding_mode = 'circular' - netscale = 4 - self.send_info("Loading Upsampler") - upsampler = RealESRGANer( - scale=netscale, - model_path=REAL_ESRGAN_WEIGHTS_PATH, - model=real_esrgan_model, - tile=0, - tile_pad=10, - pre_pad=0, - half=not args['full_precision'] - ) - self.send_info("Enhancing Input") - output, _ = upsampler.enhance(image, outscale=args['outscale']) - self.send_info("Converting Result") - output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) - output = Image.fromarray(output) - self.send_action(Action.IMAGE, shared_memory_name=self.share_image_memory(output), seed=args['name'], width=output.width, height=output.height) - args = yield \ No newline at end of file diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 811ef6f4..c7421f8a 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -106,12 +106,14 @@ def view_step(step, width=None, height=None, shared_memory_name=None): bpy.data.images.remove(last_data_block) last_data_block = step_image return # Only perform this on the first image editor found. - # dream_texture(context.scene.dream_textures_prompt, view_step, image_writer) + def step_callback(future, step_image): + image = bpy_image("diffusers-image", step_image.shape[0], step_image.shape[1], step_image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = image def image_done(future): - if future.cancelled(): - del gen._active_generation_future - return - image: NDArray = future.result() + del gen._active_generation_future + image = future.result()[-1] image = bpy_image("diffusers-image", image.shape[0], image.shape[1], image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': @@ -122,6 +124,7 @@ def image_done(future): **scene.dream_textures_prompt.generate_args(), ) gen._active_generation_future = f + f.add_response_callback(step_callback) f.add_done_callback(image_done) return {"FINISHED"} diff --git a/operators/upscale.py b/operators/upscale.py index abb0dc2c..1448a3c6 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -3,21 +3,7 @@ from multiprocessing.shared_memory import SharedMemory import numpy as np import sys -# from ..generator_process import GeneratorProcess - -upscale_options = [ - ("2", "2x", "", 2), - ("4", "4x", "", 4), -] - -generator_advance = None -timer = None - -def remove_timer(context): - global timer - if timer: - context.window_manager.event_timer_remove(timer) - timer = None +from ..generator_process import Generator class Upscale(bpy.types.Operator): bl_idname = "shade.dream_textures_upscale" @@ -27,20 +13,7 @@ class Upscale(bpy.types.Operator): @classmethod def poll(cls, context): - return GeneratorProcess.can_use() - - def modal(self, context, event): - if event.type != 'TIMER': - return {'PASS_THROUGH'} - try: - next(generator_advance) - except StopIteration: - remove_timer(context) - return {'FINISHED'} - except Exception as e: - remove_timer(context) - raise e - return {'RUNNING_MODAL'} + return Generator.shared().can_use() def execute(self, context): scene = context.scene @@ -103,7 +76,7 @@ def image_callback(shared_memory_name, seed, width, height): scene.dream_textures_info = "" shared_memory = SharedMemory(shared_memory_name) image = bpy_image(seed + ' (Upscaled)', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) - for area in context.screen.areas: + for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = image if active_node is not None: @@ -118,18 +91,22 @@ def exception_callback(fatal, msg, trace): if trace: print(trace, file=sys.stderr) - generator = GeneratorProcess.shared() - - args = { - 'input': input_image_path, - 'name': input_image.name, - 'outscale': int(context.scene.dream_textures_upscale_outscale), - 'full_precision': context.scene.dream_textures_upscale_full_precision, - 'seamless': context.scene.dream_textures_upscale_seamless - } - global generator_advance - generator_advance = generator.upscale(args, image_callback, info_callback, exception_callback) - context.window_manager.modal_handler_add(self) - self.timer = context.window_manager.event_timer_add(1 / 15, window=context.window) - - return {"RUNNING_MODAL"} \ No newline at end of file + # args = { + # 'input': input_image_path, + # 'name': input_image.name, + # 'outscale': int(context.scene.dream_textures_upscale_outscale), + # 'full_precision': context.scene.dream_textures_upscale_full_precision, + # 'seamless': context.scene.dream_textures_upscale_seamless + # } + + def image_done(future): + image = future.result() + image = bpy_image("diffusers-upscaled", image.shape[0], image.shape[1], image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = image + if active_node is not None: + active_node.image = image + Generator.shared().upscale(input_image_path, "brick wall", context.scene.dream_textures_upscale_full_precision).add_done_callback(image_done) + + return {"FINISHED"} \ No newline at end of file diff --git a/preferences.py b/preferences.py index 874827d8..c3dc2567 100644 --- a/preferences.py +++ b/preferences.py @@ -4,8 +4,6 @@ import os import webbrowser import shutil -from concurrent.futures import Future -import functools from .absolute_path import WEIGHTS_PATH, absolute_path from .operators.install_dependencies import InstallDependencies @@ -13,8 +11,6 @@ from .property_groups.dream_prompt import DreamPrompt from .ui.presets import RestoreDefaultPresets, default_presets_missing from .generator_process import Generator -from .generator_process.actions.huggingface_hub import Model -from typing import List class OpenHuggingFace(bpy.types.Operator): bl_idname = "dream_textures.open_hugging_face" @@ -158,9 +154,7 @@ class StableDiffusionPreferences(bpy.types.AddonPreferences): @staticmethod def register(): - print("register") set_model_list('installed_models', Generator.shared().hf_list_installed_models().result()) - print("done") def draw(self, context): layout = self.layout diff --git a/requirements.txt b/requirements.txt index ab3914df..5dae8140 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -git+https://github.com/huggingface/diffusers@fix-sd2-attention-slice#egg=diffusers # SD2 attention slicing fix +git+https://github.com/huggingface/diffusers@main#egg=diffusers transformers accelerate huggingface_hub diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index cc675d99..340cba4f 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -314,7 +314,7 @@ def optimization(prop): optimization("sequential_cpu_offload") optimization("channels_last_memory_format") optimization("cpu_only") - # optimization("xformers_attention") # FIXME: xFormers is not currently supported due to a lack of official Windows binaries. + # optimization("xformers_attention") # FIXME: xFormers is not yet available. yield MemoryOptimizationPanel def actions_panel(sub_panel, space_type, get_prompt): diff --git a/ui/panels/upscaling.py b/ui/panels/upscaling.py index 7786058c..f6157b34 100644 --- a/ui/panels/upscaling.py +++ b/ui/panels/upscaling.py @@ -8,39 +8,6 @@ from ...absolute_path import REAL_ESRGAN_WEIGHTS_PATH from ..space_types import SPACE_TYPES import os -import webbrowser -import shutil - -class OpenRealESRGANDownload(bpy.types.Operator): - bl_idname = "stable_diffusion.open_realesrgan_download" - bl_label = "Download Weights from GitHub" - bl_description = ("Opens to the latest release of Real-ESRGAN, where the weights can be downloaded.") - bl_options = {"REGISTER", "INTERNAL"} - - def execute(self, context): - webbrowser.open("https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.3.0") - return {"FINISHED"} - -class OpenRealESRGANWeightsDirectory(bpy.types.Operator, ImportHelper): - bl_idname = "stable_diffusion.open_realesrgan_weights_directory" - bl_label = "Import Model Weights" - bl_description = ("Opens the directory that should contain the 'realesr-general-x4v3.pth' file") - - filename_ext = ".pth" - filter_glob: bpy.props.StringProperty( - default="*.pth", - options={'HIDDEN'}, - maxlen=255, - ) - - def execute(self, context): - _, extension = os.path.splitext(self.filepath) - if extension != '.pth': - self.report({"ERROR"}, "Select a valid Real-ESRGAN '.pth' file.") - return {"FINISHED"} - shutil.copy(self.filepath, REAL_ESRGAN_WEIGHTS_PATH) - - return {"FINISHED"} def upscaling_panels(): for space_type in SPACE_TYPES: @@ -54,10 +21,10 @@ class UpscalingPanel(Panel): bl_options = {'DEFAULT_CLOSED'} @classmethod - def poll(self, context): + def poll(cls, context): if not BackendTarget[context.scene.dream_textures_prompt.backend].upscaling(): return False - if self.bl_space_type == 'NODE_EDITOR': + if cls.bl_space_type == 'NODE_EDITOR': return context.area.ui_type == "ShaderNodeTree" or context.area.ui_type == "CompositorNodeTree" else: return True @@ -65,21 +32,16 @@ def poll(self, context): def draw(self, context): layout = self.layout layout.use_property_split = True - if not os.path.exists(REAL_ESRGAN_WEIGHTS_PATH): - layout.label(text="Real-ESRGAN model weights not installed.") - layout.label(text="1. Download the file 'realesr-general-x4v3.pth' from GitHub") - layout.operator(OpenRealESRGANDownload.bl_idname, icon="URL") - layout.label(text="2. Select the downloaded weights to install.") - layout.operator(OpenRealESRGANWeightsDirectory.bl_idname, icon="IMPORT") - layout = layout.column() - layout.enabled = os.path.exists(REAL_ESRGAN_WEIGHTS_PATH) + layout.prop(context.scene, "dream_textures_upscale_outscale") layout.prop(context.scene, "dream_textures_upscale_full_precision") layout.prop(context.scene, "dream_textures_upscale_seamless") + if not context.scene.dream_textures_upscale_full_precision: box = layout.box() box.label(text="Note: Some GPUs do not support mixed precision math", icon="ERROR") box.label(text="If you encounter an error, enable full precision.") + if context.scene.dream_textures_info != "": layout.label(text=context.scene.dream_textures_info, icon="INFO") else: From 81830f510ea3ce75177cc3d354b8ae03ff4a5679 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 27 Nov 2022 17:47:19 -0500 Subject: [PATCH 08/36] Add tiling to stable diffusion upscaler --- .github/workflows/package-release.yml | 4 +- __init__.py | 5 +- generator_process/actions/prompt_to_image.py | 68 ++++++++++++------ generator_process/actions/upscale.py | 60 ++++++++++++---- generator_process/actor.py | 3 +- operators/upscale.py | 49 ++++++------- ui/panels/upscaling.py | 74 +++++++++++++++----- 7 files changed, 175 insertions(+), 88 deletions(-) diff --git a/.github/workflows/package-release.yml b/.github/workflows/package-release.yml index 4a9c119f..83bfaccd 100644 --- a/.github/workflows/package-release.yml +++ b/.github/workflows/package-release.yml @@ -33,7 +33,6 @@ jobs: - name: Checkout repository uses: actions/checkout@v3 with: - submodules: recursive path: dream_textures - name: Setup Python uses: actions/setup-python@v4 @@ -41,8 +40,7 @@ jobs: python-version: '3.10' - name: Install dependencies into target shell: bash - run: 'python -m pip install -r ../requirements-lin-win-colab-CUDA.txt --no-cache-dir --target ../.python_dependencies' - working-directory: dream_textures/stable_diffusion + run: 'python -m pip install -r requirements.txt --no-cache-dir --target .python_dependencies' - name: Archive Release uses: thedoctor0/zip-release@main with: diff --git a/__init__.py b/__init__.py index 682bade4..c841091e 100644 --- a/__init__.py +++ b/__init__.py @@ -42,6 +42,7 @@ def clear_modules(): from .classes import CLASSES, PREFERENCE_CLASSES from .tools import TOOLS from .operators.dream_texture import DreamTexture, kill_generator + from .operators.upscale import upscale_options from .property_groups.dream_prompt import DreamPrompt from .preferences import StableDiffusionPreferences from .ui.presets import register_default_presets @@ -85,8 +86,8 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_render_properties_enabled = BoolProperty(default=False) bpy.types.Scene.dream_textures_render_properties_prompt = PointerProperty(type=DreamPrompt) - bpy.types.Scene.dream_textures_upscale_full_precision = bpy.props.BoolProperty(name="Full Precision", default=True) - bpy.types.Scene.dream_textures_upscale_seamless = bpy.props.BoolProperty(name="Seamless", default=False) + bpy.types.Scene.dream_textures_upscale_prompt = PointerProperty(type=DreamPrompt) + bpy.types.Scene.dream_textures_upscale_tile_size = IntProperty(name="Tile Size", default=128, step=64, min=64, max=512) for cls in CLASSES: bpy.utils.register_class(cls) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 429320bb..f310898e 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -66,6 +66,41 @@ def can_use(self, property, device) -> bool: annotation: _AnnotatedAlias = self.__annotations__[property] return annotation.__metadata__ == device return True + + def apply(self, pipeline, device): + """ + Apply the optimizations to a diffusers pipeline. + + All exceptions are ignored to make this more general purpose across different pipelines. + """ + import torch + + torch.backends.cudnn.benchmark = self.can_use("cudnn_benchmark", device) + torch.backends.cuda.matmul.allow_tf32 = self.can_use("tf32", device) + + try: + if self.can_use("attention_slicing", device): + pipeline.enable_attention_slicing(self.attention_slice_size) + else: + pipeline.disable_attention_slicing() + except: pass + + try: + if self.can_use("sequential_cpu_offload", device): + pipeline.enable_sequential_cpu_offload() + except: pass + + try: + if self.can_use("channels_last_memory_format", device): + pipeline.unet.to(memory_format=torch.channels_last) + else: + pipeline.unet.to(memory_format=torch.contiguous_format) + except: pass + + # FIXME: xFormers wheels are not yet available (https://github.com/facebookresearch/xformers/issues/533) + # if self.can_use("xformers_attention", device): + # pipeline.enable_xformers_memory_efficient_attention() + return pipeline def choose_device(self) -> str: """ @@ -95,9 +130,15 @@ def prompt_to_image( height: int, seed: int, + cfg_scale: float, + use_negative_prompt: bool, + negative_prompt: str, + seamless: bool, seamless_axes: list[str], + iterations: int, + **kwargs ) -> Generator[NDArray, None, None]: match pipeline: @@ -260,26 +301,7 @@ def _approximate_decoded_latents(self, latents): } if is_stable_diffusion_2 else None) # Optimizations - - torch.backends.cudnn.benchmark = optimizations.can_use("cudnn_benchmark", device) - torch.backends.cuda.matmul.allow_tf32 = optimizations.can_use("tf32", device) - - if optimizations.can_use("attention_slicing", device): - pipe.enable_attention_slicing(optimizations.attention_slice_size) - else: - pipe.disable_attention_slicing() - - if use_cpu_offload: - pipe.enable_sequential_cpu_offload() - - if optimizations.can_use("channels_last_memory_format", device): - pipe.unet.to(memory_format=torch.channels_last) - else: - pipe.unet.to(memory_format=torch.contiguous_format) - - # FIXME: xFormers wheels are not yet available (https://github.com/facebookresearch/xformers/issues/533) - # if optimizations.can_use("xformers_attention", device): - # pipe.enable_xformers_memory_efficient_attention() + pipe = optimizations.apply(pipe, device) # RNG generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) @@ -296,9 +318,9 @@ def _approximate_decoded_latents(self, latents): height=height, width=width, num_inference_steps=steps, - guidance_scale=7.5, - negative_prompt=None, - num_images_per_prompt=1, + guidance_scale=cfg_scale, + negative_prompt=negative_prompt if use_negative_prompt else None, + num_images_per_prompt=iterations, eta=0.0, generator=generator, latents=None, diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index 4574c66f..f97e6d59 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -1,26 +1,60 @@ -from numpy.typing import NDArray import numpy as np +from .prompt_to_image import Optimizations, Scheduler def upscale( self, image: str, + prompt: str, + steps: int, + seed: int, + cfg_scale: float, + scheduler: Scheduler, + + tile_size: int, + + optimizations: Optimizations, - half_precision: bool -) -> NDArray: + **kwargs +): + from PIL import Image, ImageOps import torch import diffusers - from PIL import Image, ImageOps - model_id = "stabilityai/stable-diffusion-x4-upscaler" + if optimizations.cpu_only: + device = "cpu" + else: + device = self.choose_device() + pipe = diffusers.StableDiffusionUpscalePipeline.from_pretrained( - model_id, - revision="fp16" if half_precision else None, - torch_dtype=torch.float16 if half_precision else torch.float32 + "stabilityai/stable-diffusion-x4-upscaler", + revision="fp16" if optimizations.can_use("half_precision", device) else None, + torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32 ) - pipe = pipe.to(self.choose_device()) - - pipe.enable_attention_slicing() + pipe.scheduler = scheduler.create(pipe, None) + pipe = pipe.to(device) + pipe = optimizations.apply(pipe, device) + + low_res_image = Image.open(image) + + generator = torch.Generator() if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + initial_seed = generator.initial_seed() + + final = Image.new('RGB', (low_res_image.size[0] * 4, low_res_image.size[1] * 4)) + for x in range(low_res_image.size[0] // tile_size): + for y in range(low_res_image.size[1] // tile_size): + x_offset = x * tile_size + y_offset = y * tile_size + tile = low_res_image.crop((x_offset, y_offset, x_offset + tile_size, y_offset + tile_size)) + upscaled = pipe( + prompt=prompt, + image=tile, + num_inference_steps=steps, + generator=torch.manual_seed(initial_seed), + guidance_scale=cfg_scale, + + ).images[0] + final.paste(upscaled, (x_offset * 4, y_offset * 4)) + yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. - result = pipe(prompt=prompt, image=Image.open(image).convert('RGB').resize((128, 128))).images[0] - return np.asarray(ImageOps.flip(result).convert('RGBA'), dtype=np.float32) / 255. \ No newline at end of file + yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. \ No newline at end of file diff --git a/generator_process/actor.py b/generator_process/actor.py index 10e82c11..5de93950 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -68,7 +68,6 @@ def _done(_): return self._exception def cancel(self): - self.done = True self._cancelled = True self.set_done() @@ -209,7 +208,7 @@ def start(self: T) -> T: """ match self.context: case ActorContext.FRONTEND: - self.process = Process(target=_start_backend, args=(self.__class__, self._message_queue, self._response_queue), name="__actor__") + self.process = Process(target=_start_backend, args=(self.__class__, self._message_queue, self._response_queue), name="__actor__", daemon=True) self.process.start() case ActorContext.BACKEND: self._backend_loop() diff --git a/operators/upscale.py b/operators/upscale.py index 1448a3c6..61951650 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -1,10 +1,14 @@ import bpy import tempfile -from multiprocessing.shared_memory import SharedMemory -import numpy as np -import sys +from ..prompt_engineering import custom_structure from ..generator_process import Generator +upscale_options = [ + ("2", "2x", "", 2), + ("4", "4x", "", 4), + ("8", "8x", "", 8), +] + class Upscale(bpy.types.Operator): bl_idname = "shade.dream_textures_upscale" bl_label = "Upscale" @@ -16,7 +20,6 @@ def poll(cls, context): return Generator.shared().can_use() def execute(self, context): - scene = context.scene screen = context.screen 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 @@ -72,41 +75,29 @@ def bpy_image(name, width, height, pixels): image.pack() return image - def image_callback(shared_memory_name, seed, width, height): - scene.dream_textures_info = "" - shared_memory = SharedMemory(shared_memory_name) - image = bpy_image(seed + ' (Upscaled)', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) + def on_tile_complete(_, tile): + image = bpy_image("diffusers-upscaled", tile.shape[0], tile.shape[1], tile.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = image - if active_node is not None: - active_node.image = image - shared_memory.close() - - def info_callback(msg=""): - scene.dream_textures_info = msg - def exception_callback(fatal, msg, trace): - scene.dream_textures_info = "" - self.report({'ERROR'}, msg) - if trace: - print(trace, file=sys.stderr) - - # args = { - # 'input': input_image_path, - # 'name': input_image.name, - # 'outscale': int(context.scene.dream_textures_upscale_outscale), - # 'full_precision': context.scene.dream_textures_upscale_full_precision, - # 'seamless': context.scene.dream_textures_upscale_seamless - # } def image_done(future): - image = future.result() + image = future.result()[-1] image = bpy_image("diffusers-upscaled", image.shape[0], image.shape[1], image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = image if active_node is not None: active_node.image = image - Generator.shared().upscale(input_image_path, "brick wall", context.scene.dream_textures_upscale_full_precision).add_done_callback(image_done) + gen = Generator.shared() + context.scene.dream_textures_upscale_prompt.prompt_structure = custom_structure.id + f = gen.upscale( + image=input_image_path, + tile_size=context.scene.dream_textures_upscale_tile_size, + **context.scene.dream_textures_upscale_prompt.generate_args() + ) + f.add_response_callback(on_tile_complete) + f.add_done_callback(image_done) + gen._active_generation_future = f return {"FINISHED"} \ No newline at end of file diff --git a/ui/panels/upscaling.py b/ui/panels/upscaling.py index f6157b34..bc5b1c3c 100644 --- a/ui/panels/upscaling.py +++ b/ui/panels/upscaling.py @@ -1,13 +1,10 @@ -import bpy from bpy.types import Panel -from bpy_extras.io_utils import ImportHelper from ...generator_process.registrar import BackendTarget from ...pil_to_image import * from ...prompt_engineering import * from ...operators.upscale import Upscale -from ...absolute_path import REAL_ESRGAN_WEIGHTS_PATH +from .dream_texture import create_panel, advanced_panel from ..space_types import SPACE_TYPES -import os def upscaling_panels(): for space_type in SPACE_TYPES: @@ -32,20 +29,65 @@ def poll(cls, context): def draw(self, context): layout = self.layout layout.use_property_split = True + layout.use_property_decorate = False + + prompt = context.scene.dream_textures_upscale_prompt - layout.prop(context.scene, "dream_textures_upscale_outscale") - layout.prop(context.scene, "dream_textures_upscale_full_precision") - layout.prop(context.scene, "dream_textures_upscale_seamless") - - if not context.scene.dream_textures_upscale_full_precision: - box = layout.box() - box.label(text="Note: Some GPUs do not support mixed precision math", icon="ERROR") - box.label(text="If you encounter an error, enable full precision.") + layout.prop(prompt, "prompt_structure_token_subject") + layout.prop(context.scene, "dream_textures_upscale_tile_size") + + if context.scene.dream_textures_upscale_tile_size > 128: + warning_box = layout.box() + warning_box.label(text="Warning", icon="ERROR") + warning_box.label(text="Large tile sizes consume more VRAM.") + + UpscalingPanel.__name__ = UpscalingPanel.bl_idname + class ActionsPanel(Panel): + """Panel for AI Upscaling Actions""" + bl_category = "Dream" + bl_label = "Actions" + bl_idname = f"DREAM_PT_dream_upscaling_actions_panel_{space_type}" + bl_space_type = space_type + bl_region_type = 'UI' + bl_parent_id = UpscalingPanel.bl_idname + bl_options = {'HIDE_HEADER'} + + @classmethod + def poll(cls, context): + if not BackendTarget[context.scene.dream_textures_prompt.backend].upscaling(): + return False + if cls.bl_space_type == 'NODE_EDITOR': + return context.area.ui_type == "ShaderNodeTree" or context.area.ui_type == "CompositorNodeTree" + else: + return True + + def draw(self, context): + layout = self.layout + layout.use_property_split = True if context.scene.dream_textures_info != "": layout.label(text=context.scene.dream_textures_info, icon="INFO") else: - layout.operator(Upscale.bl_idname, icon="FULLSCREEN_ENTER") - - UpscalingPanel.__name__ = f"DREAM_PT_dream_troubleshooting_panel_{space_type}" - yield UpscalingPanel \ No newline at end of file + image = None + for area in context.screen.areas: + if area.type == 'IMAGE_EDITOR': + image = area.spaces.active.image + + row = layout.row() + row.scale_y = 1.5 + row.operator( + Upscale.bl_idname, + text=f"Upscale to {image.size[0] * 4}x{image.size[1] * 4}" if image is not None else "Upscale", + icon="FULLSCREEN_ENTER" + ) + yield UpscalingPanel + advanced_panels = [*create_panel(space_type, 'UI', UpscalingPanel.bl_idname, advanced_panel, lambda context: context.scene.dream_textures_upscale_prompt)] + outer_panel = advanced_panels[0] + outer_original_idname = outer_panel.bl_idname + outer_panel.bl_idname += "_upscaling" + for panel in advanced_panels: + panel.bl_idname += "_upscaling" + if panel.bl_parent_id == outer_original_idname: + panel.bl_parent_id = outer_panel.bl_idname + yield panel + yield ActionsPanel \ No newline at end of file From fb07290904f71c6aa0c1508354a929e08a95edd7 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 27 Nov 2022 19:15:49 -0500 Subject: [PATCH 09/36] Add image_to_image and inpaint --- generator_process/__init__.py | 2 + generator_process/actions/image_to_image.py | 212 +++++++++++++++ generator_process/actions/inpaint.py | 255 +++++++++++++++++++ generator_process/actions/prompt_to_image.py | 72 ++++-- operators/dream_texture.py | 147 ++++++----- property_groups/dream_prompt.py | 7 +- render_pass.py | 15 +- ui/panels/dream_texture.py | 2 +- 8 files changed, 605 insertions(+), 107 deletions(-) create mode 100644 generator_process/actions/image_to_image.py create mode 100644 generator_process/actions/inpaint.py diff --git a/generator_process/__init__.py b/generator_process/__init__.py index 835c7997..401544fb 100644 --- a/generator_process/__init__.py +++ b/generator_process/__init__.py @@ -6,6 +6,8 @@ class Generator(Actor): """ from .actions.prompt_to_image import prompt_to_image, choose_device + from .actions.image_to_image import image_to_image + from .actions.inpaint import inpaint from .actions.upscale import upscale from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models from .actions.ocio_transform import ocio_transform \ No newline at end of file diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py new file mode 100644 index 00000000..75e5496e --- /dev/null +++ b/generator_process/actions/image_to_image.py @@ -0,0 +1,212 @@ +from typing import Union, Generator, Callable, List, Optional +import os +from contextlib import nullcontext +from numpy.typing import NDArray +import numpy as np +from .prompt_to_image import Pipeline, Scheduler, Optimizations, StepPreviewMode, approximate_decoded_latents, _configure_model_padding + +def image_to_image( + self, + pipeline: Pipeline, + + model: str, + + scheduler: Scheduler, + + optimizations: Optimizations, + + image: NDArray | str, + fit: bool, + strength: float, + prompt: str, + steps: int, + width: int, + height: int, + seed: int, + + cfg_scale: float, + use_negative_prompt: bool, + negative_prompt: str, + + seamless: bool, + seamless_axes: list[str], + + iterations: int, + + step_preview_mode: StepPreviewMode, + + **kwargs +) -> Generator[NDArray, None, None]: + match pipeline: + case Pipeline.STABLE_DIFFUSION: + import diffusers + import torch + from PIL import Image, ImageOps + from ...absolute_path import WEIGHTS_PATH + + # Mostly copied from `diffusers.StableDiffusionImg2ImgPipeline`, with slight modifications to yield the latents at each step. + class GeneratorPipeline(diffusers.StableDiffusionImg2ImgPipeline): + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + init_image: Union[torch.FloatTensor, Image.Image], + strength: float = 0.8, + num_inference_steps: Optional[int] = 50, + guidance_scale: Optional[float] = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: Optional[float] = 0.0, + generator: Optional[torch.Generator] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + **kwargs, + ): + # 1. Check inputs + self.check_inputs(prompt, strength, callback_steps) + + # 2. Define call parameters + batch_size = 1 if isinstance(prompt, str) else len(prompt) + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_embeddings = self._encode_prompt( + prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt + ) + + # 4. Preprocess image + if isinstance(init_image, Image.Image): + init_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess(init_image) + + # 5. set timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.get_timesteps(num_inference_steps, strength, device) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + + # 6. Prepare latent variables + latents = self.prepare_latents( + init_image, latent_timestep, batch_size, num_images_per_prompt, text_embeddings.dtype, device, generator + ) + + # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 8. Denoising loop + for i, t in enumerate(self.progress_bar(timesteps)): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # NOTE: Modified to yield the latents instead of calling a callback. + match kwargs['step_preview_mode']: + case StepPreviewMode.NONE: + pass + case StepPreviewMode.FAST: + yield np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255. + case StepPreviewMode.ACCURATE: + yield from [ + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for image in self.numpy_to_pil(self.decode_latents(latents)) + ] + + # 9. Post-processing + image = self.decode_latents(latents) + + # TODO: Add UI to enable this + # 10. Run safety checker + # image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype) + + # NOTE: Modified to yield the decoded image as a numpy array. + yield from [ + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for image in self.numpy_to_pil(image) + ] + + if optimizations.cpu_only: + device = "cpu" + else: + device = self.choose_device() + + use_cpu_offload = optimizations.can_use("sequential_cpu_offload", device) + + # StableDiffusionPipeline w/ caching + if hasattr(self, "_cached_img2img_pipe") and self._cached_img2img_pipe[1] == model and use_cpu_offload == self._cached_img2img_pipe[2]: + pipe = self._cached_img2img_pipe[0] + else: + storage_folder = os.path.join(WEIGHTS_PATH, model) + revision = "main" + ref_path = os.path.join(storage_folder, "refs", revision) + with open(ref_path) as f: + commit_hash = f.read() + + snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) + pipe = GeneratorPipeline.from_pretrained( + snapshot_folder, + revision="fp16" if optimizations.can_use("half_precision", device) else None, + torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32, + ) + pipe = pipe.to(device) + setattr(self, "_cached_img2img_pipe", (pipe, model, use_cpu_offload, snapshot_folder)) + + # Scheduler + is_stable_diffusion_2 = 'stabilityai--stable-diffusion-2' in model + pipe.scheduler = scheduler.create(pipe, { + 'model_path': self._cached_img2img_pipe[3], + 'subfolder': 'scheduler', + } if is_stable_diffusion_2 else None) + + # Optimizations + pipe = optimizations.apply(pipe, device) + + # RNG + generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + + # Seamless + _configure_model_padding(pipe.unet, seamless, seamless_axes) + _configure_model_padding(pipe.vae, seamless, seamless_axes) + + # Inference + with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + init_image = (Image.open(image) if isinstance(image, str) else Image.fromarray(image)).convert('RGB') + yield from pipe( + prompt=prompt, + init_image=init_image, + strength=strength, + height=init_image.size[1] if fit else height, + width=init_image.size[0] if fit else width, + num_inference_steps=steps, + guidance_scale=cfg_scale, + negative_prompt=negative_prompt if use_negative_prompt else None, + num_images_per_prompt=iterations, + eta=0.0, + generator=generator, + latents=None, + output_type="pil", + return_dict=True, + callback=None, + callback_steps=1, + step_preview_mode=step_preview_mode + ) + case Pipeline.STABILITY_SDK: + import stability_sdk + raise NotImplementedError() + case _: + raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actions/inpaint.py b/generator_process/actions/inpaint.py new file mode 100644 index 00000000..a80b4209 --- /dev/null +++ b/generator_process/actions/inpaint.py @@ -0,0 +1,255 @@ +from typing import Union, Generator, Callable, List, Optional +import os +from contextlib import nullcontext +from numpy.typing import NDArray +import numpy as np +from .prompt_to_image import Pipeline, Scheduler, Optimizations, StepPreviewMode, approximate_decoded_latents, _configure_model_padding + +def inpaint( + self, + pipeline: Pipeline, + + model: str, + + scheduler: Scheduler, + + optimizations: Optimizations, + + image: NDArray | str, + fit: bool, + strength: float, + prompt: str, + steps: int, + width: int, + height: int, + seed: int, + + cfg_scale: float, + use_negative_prompt: bool, + negative_prompt: str, + + seamless: bool, + seamless_axes: list[str], + + iterations: int, + + step_preview_mode: StepPreviewMode, + + **kwargs +) -> Generator[NDArray, None, None]: + match pipeline: + case Pipeline.STABLE_DIFFUSION: + import diffusers + import torch + from PIL import Image, ImageOps + from ...absolute_path import WEIGHTS_PATH + + # Mostly copied from `diffusers.StableDiffusionInpaintPipeline`, with slight modifications to yield the latents at each step. + class GeneratorPipeline(diffusers.StableDiffusionInpaintPipeline): + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + image: Union[torch.FloatTensor, Image.Image], + mask_image: Union[torch.FloatTensor, Image.Image], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[torch.Generator] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + **kwargs, + ): + # 0. Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # 1. Check inputs + self.check_inputs(prompt, height, width, callback_steps) + + # 2. Define call parameters + batch_size = 1 if isinstance(prompt, str) else len(prompt) + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_embeddings = self._encode_prompt( + prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt + ) + + # 4. Preprocess mask and image + if isinstance(image, Image.Image) and isinstance(mask_image, Image.Image): + mask, masked_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.prepare_mask_and_masked_image(image, mask_image) + + # 5. set timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps_tensor = self.scheduler.timesteps + + # 6. Prepare latent variables + num_channels_latents = self.vae.config.latent_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + + # 7. Prepare mask latent variables + mask, masked_image_latents = self.prepare_mask_latents( + mask, + masked_image, + batch_size * num_images_per_prompt, + height, + width, + text_embeddings.dtype, + device, + generator, + do_classifier_free_guidance, + ) + + # 8. Check that sizes of mask, masked image and latents match + num_channels_mask = mask.shape[1] + num_channels_masked_image = masked_image_latents.shape[1] + if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: + raise ValueError( + f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" + f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" + f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" + f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of" + " `pipeline.unet` or your `mask_image` or `image` input." + ) + + # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 10. Denoising loop + for i, t in enumerate(self.progress_bar(timesteps_tensor)): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + # concat latents, mask, masked_image_latents in the channel dimension + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) + + # predict the noise residual + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # NOTE: Modified to yield the latents instead of calling a callback. + match kwargs['step_preview_mode']: + case StepPreviewMode.NONE: + pass + case StepPreviewMode.FAST: + yield np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255. + case StepPreviewMode.ACCURATE: + yield from [ + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for image in self.numpy_to_pil(self.decode_latents(latents)) + ] + + # 11. Post-processing + image = self.decode_latents(latents) + + # TODO: Add UI to enable this. + # 12. Run safety checker + # image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype) + + # NOTE: Modified to yield the decoded image as a numpy array. + yield from [ + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for image in self.numpy_to_pil(image) + ] + + if optimizations.cpu_only: + device = "cpu" + else: + device = self.choose_device() + + use_cpu_offload = optimizations.can_use("sequential_cpu_offload", device) + + # StableDiffusionPipeline w/ caching + if hasattr(self, "_cached_img2img_pipe") and self._cached_img2img_pipe[1] == model and use_cpu_offload == self._cached_img2img_pipe[2]: + pipe = self._cached_img2img_pipe[0] + else: + storage_folder = os.path.join(WEIGHTS_PATH, model) + revision = "main" + ref_path = os.path.join(storage_folder, "refs", revision) + with open(ref_path) as f: + commit_hash = f.read() + + snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) + pipe = GeneratorPipeline.from_pretrained( + snapshot_folder, + revision="fp16" if optimizations.can_use("half_precision", device) else None, + torch_dtype=torch.float16 if optimizations.can_use("half_precision", device) else torch.float32, + ) + pipe = pipe.to(device) + setattr(self, "_cached_img2img_pipe", (pipe, model, use_cpu_offload, snapshot_folder)) + + # Scheduler + is_stable_diffusion_2 = 'stabilityai--stable-diffusion-2' in model + pipe.scheduler = scheduler.create(pipe, { + 'model_path': self._cached_img2img_pipe[3], + 'subfolder': 'scheduler', + } if is_stable_diffusion_2 else None) + + # Optimizations + pipe = optimizations.apply(pipe, device) + + # RNG + generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + + # Seamless + _configure_model_padding(pipe.unet, seamless, seamless_axes) + _configure_model_padding(pipe.vae, seamless, seamless_axes) + + # Inference + with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + init_image = Image.open(image) if isinstance(image, str) else Image.fromarray(image) + yield from pipe( + prompt=prompt, + image=init_image.convert('RGB'), + mask_image=ImageOps.invert(init_image.getchannel('A')), + strength=strength, + height=init_image.size[1] if fit else height, + width=init_image.size[0] if fit else width, + num_inference_steps=steps, + guidance_scale=cfg_scale, + negative_prompt=negative_prompt if use_negative_prompt else None, + num_images_per_prompt=iterations, + eta=0.0, + generator=generator, + latents=None, + output_type="pil", + return_dict=True, + callback=None, + callback_steps=1, + step_preview_mode=step_preview_mode + ) + case Pipeline.STABILITY_SDK: + import stability_sdk + raise NotImplementedError() + case _: + raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index f310898e..2e96c6ac 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -42,7 +42,6 @@ def scheduler_class(): return scheduler_class().from_pretrained(pretrained['model_path'], subfolder=pretrained['subfolder']) else: return scheduler_class().from_config(pipeline.scheduler.config) - @dataclass(eq=True) class Optimizations: @@ -102,6 +101,11 @@ def apply(self, pipeline, device): # pipeline.enable_xformers_memory_efficient_attention() return pipeline +class StepPreviewMode(enum.Enum): + NONE = "None" + FAST = "Fast" + ACCURATE = "Accurate" + def choose_device(self) -> str: """ Automatically select which PyTorch device to use. @@ -114,6 +118,31 @@ def choose_device(self) -> str: else: return "cpu" +def approximate_decoded_latents(latents): + """ + Approximate the decoded latents without using the VAE. + """ + import torch + # origingally adapted from code by @erucipe and @keturn here: + # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7 + + # these updated numbers for v1.5 are from @torridgristle + v1_5_latent_rgb_factors = torch.tensor([ + # R G B + [ 0.3444, 0.1385, 0.0670], # L1 + [ 0.1247, 0.4027, 0.1494], # L2 + [-0.3192, 0.2513, 0.2103], # L3 + [-0.1307, -0.1874, -0.7445] # L4 + ], 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() + def prompt_to_image( self, pipeline: Pipeline, @@ -139,6 +168,8 @@ def prompt_to_image( iterations: int, + step_preview_mode: StepPreviewMode, + **kwargs ) -> Generator[NDArray, None, None]: match pipeline: @@ -227,7 +258,16 @@ def __call__( latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # NOTE: Modified to yield the latents instead of calling a callback. - yield np.asarray(ImageOps.flip(Image.fromarray(self._approximate_decoded_latents(latents))).convert('RGBA'), dtype=np.float32) / 255. + match kwargs['step_preview_mode']: + case StepPreviewMode.NONE: + pass + case StepPreviewMode.FAST: + yield np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255. + case StepPreviewMode.ACCURATE: + yield from [ + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for image in self.numpy_to_pil(self.decode_latents(latents)) + ] # 8. Post-processing image = self.decode_latents(latents) @@ -241,31 +281,6 @@ def __call__( np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. for image in self.numpy_to_pil(image) ] - - - def _approximate_decoded_latents(self, latents): - """ - Approximate the decoded latents without using the VAE. - """ - # origingally adapted from code by @erucipe and @keturn here: - # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7 - - # these updated numbers for v1.5 are from @torridgristle - v1_5_latent_rgb_factors = torch.tensor([ - # R G B - [ 0.3444, 0.1385, 0.0670], # L1 - [ 0.1247, 0.4027, 0.1494], # L2 - [-0.3192, 0.2513, 0.2103], # L3 - [-0.1307, -0.1874, -0.7445] # L4 - ], 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() if optimizations.cpu_only: device = "cpu" @@ -327,7 +342,8 @@ def _approximate_decoded_latents(self, latents): output_type="pil", return_dict=True, callback=None, - callback_steps=1 + callback_steps=1, + step_preview_mode=step_preview_mode ) case Pipeline.STABILITY_SDK: import stability_sdk diff --git a/operators/dream_texture.py b/operators/dream_texture.py index c7421f8a..0d2308a3 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -22,6 +22,32 @@ last_data_block = None timer = None +def save_temp_image(img, path=None): + path = path if path is not None else tempfile.NamedTemporaryFile().name + + settings = bpy.context.scene.render.image_settings + file_format = settings.file_format + mode = settings.color_mode + depth = settings.color_depth + + settings.file_format = 'PNG' + settings.color_mode = 'RGBA' + settings.color_depth = '8' + + img.save_render(path) + + settings.file_format = file_format + settings.color_mode = mode + settings.color_depth = depth + + return path + +def bpy_image(name, width, height, pixels): + image = bpy.data.images.new(name, width=width, height=height) + image.pixels[:] = pixels + image.pack() + return image + class DreamTexture(bpy.types.Operator): bl_idname = "shade.dream_texture" bl_label = "Dream Texture" @@ -53,79 +79,74 @@ def execute(self, context): history_entry.prompt_structure = custom_structure.id history_entry.prompt_structure_token_subject = file_batch_lines[i].body - def bpy_image(name, width, height, pixels): - image = bpy.data.images.new(name, width=width, height=height) - image.pixels[:] = pixels - image.pack() - return image - - node_tree = context.material.node_tree if hasattr(context, 'material') else None + node_tree = context.material.node_tree if hasattr(context, 'material') and hasattr(context.material, 'node_tree') else None screen = context.screen scene = context.scene - iteration = 0 - def image_writer(shared_memory_name, seed, width, height, upscaled=False): - nonlocal iteration - global last_data_block - # Only use the non-upscaled texture, as upscaling is currently unsupported by the addon. - if not upscaled: - if last_data_block is not None: - bpy.data.images.remove(last_data_block) - last_data_block = None - generator = GeneratorProcess.shared(create=False) - if generator is None or generator.process.poll() or width == 0 or height == 0: - return # process was closed - shared_memory = SharedMemory(shared_memory_name) - image = bpy_image(f"{seed}", width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) - shared_memory.close() - if node_tree is not None: - nodes = node_tree.nodes - texture_node = nodes.new("ShaderNodeTexImage") - texture_node.image = image - nodes.active = texture_node - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = image - scene.dream_textures_prompt.seed = str(seed) # update property in case seed was sourced randomly or from hash - history_entries[iteration].seed = str(seed) - history_entries[iteration].random_seed = False - iteration += 1 - - def view_step(step, width=None, height=None, shared_memory_name=None): - scene.dream_textures_progress = step + 1 - if shared_memory_name is None: - return # show steps disabled - global last_data_block - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - shared_memory = SharedMemory(shared_memory_name) - step_image = bpy_image(f'Step {step + 1}/{scene.dream_textures_prompt.steps}', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) - shared_memory.close() - area.spaces.active.image = step_image - if last_data_block is not None: - bpy.data.images.remove(last_data_block) - last_data_block = step_image - return # Only perform this on the first image editor found. - def step_callback(future, step_image): - image = bpy_image("diffusers-image", step_image.shape[0], step_image.shape[1], step_image.ravel()) + generated_args = scene.dream_textures_prompt.generate_args() + + init_image = None + if generated_args['use_init_img']: + match generated_args['init_img_src']: + case 'file': + init_image = save_temp_image(scene.init_img) + case 'open_editor': + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + if area.spaces.active.image is not None: + init_image = save_temp_image(area.spaces.active.image) + + last_data_block = None + def step_callback(_, step_image): + nonlocal last_data_block + if last_data_block is not None: + bpy.data.images.remove(last_data_block) + last_data_block = bpy_image("Step", step_image.shape[0], step_image.shape[1], step_image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = image - def image_done(future): + area.spaces.active.image = last_data_block + + def done_callback(future): + nonlocal last_data_block del gen._active_generation_future - image = future.result()[-1] - image = bpy_image("diffusers-image", image.shape[0], image.shape[1], image.ravel()) + image = future.result() + if isinstance(image, list): + image = image[-1] + if last_data_block is not None: + bpy.data.images.remove(last_data_block) + last_data_block = bpy_image("Final", image.shape[0], image.shape[1], image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = image + area.spaces.active.image = last_data_block + if node_tree is not None: + # TODO: Create Image Texture node. + pass + gen = Generator.shared() - f = gen.prompt_to_image( - Pipeline.STABLE_DIFFUSION, - **scene.dream_textures_prompt.generate_args(), - ) + if init_image is not None: + match generated_args['init_img_action']: + case 'modify': + f = gen.image_to_image( + Pipeline.STABLE_DIFFUSION, + image=init_image, + **generated_args + ) + case 'inpaint': + f = gen.inpaint( + Pipeline.STABLE_DIFFUSION, + image=init_image, + **generated_args + ) + case 'outpaint': + pass + else: + f = gen.prompt_to_image( + Pipeline.STABLE_DIFFUSION, + **generated_args, + ) gen._active_generation_future = f f.add_response_callback(step_callback) - f.add_done_callback(image_done) + f.add_done_callback(done_callback) return {"FINISHED"} headless_prompt = None @@ -325,7 +346,7 @@ class CancelGenerator(bpy.types.Operator): @classmethod def poll(self, context): gen = Generator.shared() - return hasattr(gen, "_active_generation_future") and gen._active_generation_future is not None + return hasattr(gen, "_active_generation_future") and gen._active_generation_future is not None and not gen._active_generation_future.cancelled and not gen._active_generation_future.done def execute(self, context): gen = Generator.shared() diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 264f26d0..6ddd1d37 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -5,12 +5,14 @@ from typing import _AnnotatedAlias from ..absolute_path import absolute_path from ..generator_process.registrar import BackendTarget -from ..generator_process.actions.prompt_to_image import Optimizations, Scheduler +from ..generator_process.actions.prompt_to_image import Optimizations, Scheduler, StepPreviewMode from ..generator_process import Generator from ..prompt_engineering import * scheduler_options = [(scheduler.value, scheduler.value, '') for scheduler in Scheduler] +step_preview_mode_options = [(mode.value, mode.value, '') for mode in StepPreviewMode] + precision_options = [ ('auto', 'Automatic', "", 1), ('float32', 'Full Precision (float32)', "", 2), @@ -98,7 +100,7 @@ def seed_clamp(self, ctx): "steps": IntProperty(name="Steps", default=25, min=1), "cfg_scale": FloatProperty(name="CFG Scale", default=7.5, min=1, soft_min=1.01, description="How strongly the prompt influences the image"), "scheduler": EnumProperty(name="Scheduler", items=scheduler_options, default=0), - "show_steps": BoolProperty(name="Show Steps", description="Displays intermediate steps in the Image Viewer. Disabling can speed up generation", default=False), + "step_preview_mode": EnumProperty(name="Step Preview", description="Displays intermediate steps in the Image Viewer. Disabling can speed up generation", items=step_preview_mode_options, default=1), # Init Image "use_init_img": BoolProperty(name="Use Init Image", default=False), @@ -212,6 +214,7 @@ def generate_args(self): args['seed'] = self.get_seed() args['optimizations'] = self.get_optimizations() args['scheduler'] = Scheduler(args['scheduler']) + args['step_preview_mode'] = StepPreviewMode(args['step_preview_mode']) return args DreamPrompt.generate_prompt = generate_prompt diff --git a/render_pass.py b/render_pass.py index 77d85eb5..2ce6173a 100644 --- a/render_pass.py +++ b/render_pass.py @@ -51,19 +51,8 @@ def render(self, depsgraph): render_pass = pass_i if render_pass.name == "Dream Textures": self.update_stats("Dream Textures", "Starting") - def image_callback(set_pixels, shared_memory_name, seed, width, height, upscaled=False): - # Only use the non-upscaled texture, as upscaling is currently unsupported by the addon. - if not upscaled: - shared_memory = SharedMemory(shared_memory_name) - set_pixels(np.frombuffer(shared_memory.buf, dtype=np.float32).copy().reshape((size_x * size_y, 4))) - - shared_memory.close() step_count = int(scene.dream_textures_render_properties_prompt.strength * scene.dream_textures_render_properties_prompt.steps) - def step_callback(step, width=None, height=None, shared_memory_name=None): - self.update_stats("Dream Textures", f"Step {step + 1}/{step_count}") - self.update_progress(step / step_count) - return self.update_stats("Dream Textures", "Creating temporary image") combined_pass_image = bpy.data.images.new("dream_textures_post_processing_temp", width=size_x, height=size_y) @@ -82,7 +71,7 @@ def step_callback(step, width=None, height=None, shared_memory_name=None): display_device=scene.display_settings.display_device, look=scene.view_settings.look, inverse=False - ).get() + ).result() combined_pass_image.pixels[:] = combined_pixels.ravel() @@ -90,7 +79,7 @@ def step_callback(step, width=None, height=None, shared_memory_name=None): pixels = Generator.shared().prompt_to_image( **scene.dream_textures_render_properties_prompt.generate_args() - ) + ).result() # Perform an inverse transform so when Blender applies its transform everything looks correct. event = threading.Event() diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 340cba4f..5e583989 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -263,7 +263,7 @@ def draw(self, context): layout.prop(get_prompt(context), "steps") layout.prop(get_prompt(context), "cfg_scale") layout.prop(get_prompt(context), "scheduler") - layout.prop(get_prompt(context), "show_steps") + layout.prop(get_prompt(context), "step_preview_mode") yield AdvancedPanel From aa0132b42dd14ddbf9491c13a7a46a01da2c880a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 27 Nov 2022 21:23:18 -0500 Subject: [PATCH 10/36] Update render_pass.py --- generator_process/actions/image_to_image.py | 4 ++ generator_process/actor.py | 9 ++-- render_pass.py | 51 ++++++++++++--------- requirements.txt | 2 +- 4 files changed, 40 insertions(+), 26 deletions(-) diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index 75e5496e..148b54c4 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -186,6 +186,10 @@ def __call__( with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): init_image = (Image.open(image) if isinstance(image, str) else Image.fromarray(image)).convert('RGB') + print(fit) + print(width, height) + print(init_image) + print(init_image.size) yield from pipe( prompt=prompt, init_image=init_image, diff --git a/generator_process/actor.py b/generator_process/actor.py index 5de93950..5664f4e6 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -275,7 +275,7 @@ def _receive(self, message: Message): self._response_queue.put(Message.END) def _send(self, name): - def _send(*args, **kwargs): + def _send(*args, _block=False, **kwargs): future = Future() def _send_thread(future: Future): self._lock.acquire() @@ -296,8 +296,11 @@ def _send_thread(future: Future): future.add_response(response) self._lock.release() - thread = threading.Thread(target=_send_thread, args=(future,), daemon=True) - thread.start() + if _block: + _send_thread(future) + else: + thread = threading.Thread(target=_send_thread, args=(future,), daemon=True) + thread.start() return future return _send diff --git a/render_pass.py b/render_pass.py index 2ce6173a..0866fa2c 100644 --- a/render_pass.py +++ b/render_pass.py @@ -1,16 +1,10 @@ import bpy import cycles -import threading -import threading -import functools import numpy as np import os -from multiprocessing.shared_memory import SharedMemory - +from .generator_process.actions.prompt_to_image import Pipeline, StepPreviewMode from .generator_process import Generator -from .operators.dream_texture import dream_texture - update_render_passes_original = cycles.CyclesRender.update_render_passes render_original = cycles.CyclesRender.render # del_original = cycles.CyclesRender.__del__ @@ -52,7 +46,7 @@ def render(self, depsgraph): if render_pass.name == "Dream Textures": self.update_stats("Dream Textures", "Starting") - step_count = int(scene.dream_textures_render_properties_prompt.strength * scene.dream_textures_render_properties_prompt.steps) + # step_count = int(scene.dream_textures_render_properties_prompt.strength * scene.dream_textures_render_properties_prompt.steps) self.update_stats("Dream Textures", "Creating temporary image") combined_pass_image = bpy.data.images.new("dream_textures_post_processing_temp", width=size_x, height=size_y) @@ -62,7 +56,9 @@ def render(self, depsgraph): combined_pixels = np.empty((size_x * size_y, 4), dtype=np.float32) rect.foreach_get(combined_pixels) - combined_pixels = Generator.shared().ocio_transform( + 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, @@ -70,23 +66,36 @@ def render(self, depsgraph): view_transform=scene.view_settings.view_transform, display_device=scene.display_settings.display_device, look=scene.view_settings.look, - inverse=False + inverse=False, + _block=True ).result() - - combined_pass_image.pixels[:] = combined_pixels.ravel() - self.update_stats("Dream Textures", "Starting...") + self.update_stats("Dream Textures", "Generating...") - pixels = Generator.shared().prompt_to_image( - **scene.dream_textures_render_properties_prompt.generate_args() + generated_args = scene.dream_textures_render_properties_prompt.generate_args() + generated_args['step_preview_mode'] = StepPreviewMode.NONE + generated_args['width'] = size_x + generated_args['height'] = size_y + pixels = gen.image_to_image( + pipeline=Pipeline.STABLE_DIFFUSION, + image=(combined_pixels.reshape((size_x, size_y, 4)) * 255).astype(np.uint8), + **generated_args, + _block=True ).result() # Perform an inverse transform so when Blender applies its transform everything looks correct. - event = threading.Event() - buf = pixels.tobytes() - combined_pixels_memory.buf[:] = buf - bpy.app.timers.register(functools.partial(do_ocio_transform, event, pixels, combined_pixels_memory, True)) - event.wait() + self.update_stats("Dream Textures", "Applying inverse color management transforms") + pixels = gen.ocio_transform( + pixels, + config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'), + exposure=scene.view_settings.exposure, + gamma=scene.view_settings.gamma, + view_transform=scene.view_settings.view_transform, + display_device=scene.display_settings.display_device, + look=scene.view_settings.look, + inverse=True, + _block=True + ).result() reshaped = pixels.reshape((size_x * size_y, 4)) render_pass.rect.foreach_set(reshaped) @@ -96,8 +105,6 @@ def render(self, depsgraph): del combined_pixels del reshaped - combined_pixels_memory.close() - def cleanup(): bpy.data.images.remove(combined_pass_image) bpy.app.timers.register(cleanup) diff --git a/requirements.txt b/requirements.txt index 5dae8140..1519363c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,4 +9,4 @@ scipy # LMSDiscreteScheduler stability-sdk==0.2.6 # DreamStudio -opencolorio # color management \ No newline at end of file +opencolorio==2.1.2 # color management \ No newline at end of file From 1eb8d1a8e41fbd0d339fa41aca073f5533120b67 Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Tue, 29 Nov 2022 22:55:25 -0500 Subject: [PATCH 11/36] fix activating when diffusers is not installed --- generator_process/registrar.py | 6 ++++++ preferences.py | 6 ++++-- property_groups/dream_prompt.py | 7 ++++--- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/generator_process/registrar.py b/generator_process/registrar.py index 9bcebe77..6eea82cd 100644 --- a/generator_process/registrar.py +++ b/generator_process/registrar.py @@ -1,3 +1,5 @@ +import os +from ..absolute_path import absolute_path from .intent import Intent from enum import IntEnum @@ -6,6 +8,10 @@ class BackendTarget(IntEnum): LOCAL = 0 STABILITY_SDK = 1 + @staticmethod + def local_available(): + return os.path.exists(absolute_path(".python_dependencies/diffusers")) + def __str__(self): return self.name diff --git a/preferences.py b/preferences.py index c3dc2567..fb51588f 100644 --- a/preferences.py +++ b/preferences.py @@ -11,6 +11,7 @@ from .property_groups.dream_prompt import DreamPrompt from .ui.presets import RestoreDefaultPresets, default_presets_missing from .generator_process import Generator +from .generator_process.registrar import BackendTarget class OpenHuggingFace(bpy.types.Operator): bl_idname = "dream_textures.open_hugging_face" @@ -154,7 +155,8 @@ class StableDiffusionPreferences(bpy.types.AddonPreferences): @staticmethod def register(): - set_model_list('installed_models', Generator.shared().hf_list_installed_models().result()) + if BackendTarget.local_available(): + set_model_list('installed_models', Generator.shared().hf_list_installed_models().result()) def draw(self, context): layout = self.layout @@ -166,7 +168,7 @@ def draw(self, context): has_dependencies = len(os.listdir(absolute_path(".python_dependencies"))) > 2 if has_dependencies: - has_local = os.path.exists(absolute_path(".python_dependencies/diffusers")) > 0 + has_local = BackendTarget.local_available() if has_local: search_box = layout.box() diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 6ddd1d37..918d575d 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -54,13 +54,14 @@ def inpaint_mask_sources_filtered(self, context): def _on_model_options(future): global _model_options _model_options = future.result() -Generator.shared().hf_list_installed_models().add_done_callback(_on_model_options) +if BackendTarget.local_available(): + Generator.shared().hf_list_installed_models().add_done_callback(_on_model_options) def model_options(self, context): return [(m.id, os.path.basename(m.id).replace('models--', '').replace('--', '/'), '', i) for i, m in enumerate(_model_options)] def backend_options(self, context): def options(): - if os.path.exists(absolute_path(".python_dependencies/diffusers")): + if BackendTarget.local_available(): yield (BackendTarget.LOCAL.name, 'Local', 'Run on your own hardware', 1) if len(context.preferences.addons[__package__.split('.')[0]].preferences.dream_studio_key) > 0: yield (BackendTarget.STABILITY_SDK.name, 'DreamStudio', 'Run in the cloud with DreamStudio', 2) @@ -76,7 +77,7 @@ def seed_clamp(self, ctx): pass # will get hashed once generated attributes = { - "backend": EnumProperty(name="Backend", items=backend_options, default=1 if os.path.exists(absolute_path(".python_dependencies/diffusers")) else 2, description="Fill in a few simple options to create interesting images quickly"), + "backend": EnumProperty(name="Backend", items=backend_options, default=1 if BackendTarget.local_available() else 2, description="Fill in a few simple options to create interesting images quickly"), "model": EnumProperty(name="Model", items=model_options, description="Specify which model to use for inference"), # Prompt From b5f273b52748814e176f0cd992c59093a3a32059 Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Wed, 30 Nov 2022 00:32:05 -0500 Subject: [PATCH 12/36] load dependencies sooner --- generator_process/actions/huggingface_hub.py | 2 ++ generator_process/actor.py | 16 ++++++++++------ 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/generator_process/actions/huggingface_hub.py b/generator_process/actions/huggingface_hub.py index e11ff1eb..fe8b3d93 100644 --- a/generator_process/actions/huggingface_hub.py +++ b/generator_process/actions/huggingface_hub.py @@ -31,6 +31,8 @@ def hf_list_models( def hf_list_installed_models(self) -> list[Model]: from diffusers.utils import DIFFUSERS_CACHE + if not os.path.exists(DIFFUSERS_CACHE): + return [] return list( filter( lambda x: os.path.isdir(x.id), diff --git a/generator_process/actor.py b/generator_process/actor.py index 5664f4e6..1212b624 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -1,4 +1,4 @@ -from multiprocessing import Queue, Process, Lock +from multiprocessing import Queue, Process, Lock, current_process import multiprocessing.synchronize import enum import traceback @@ -6,6 +6,15 @@ from typing import Type, TypeVar, Callable, Any, MutableSet, Generator # from concurrent.futures import Future import site +import sys + +def _load_dependencies(): + from ..absolute_path import absolute_path + site.addsitedir(absolute_path(".python_dependencies")) + deps = sys.path.pop(-1) + sys.path.insert(0, deps) +if current_process().name == "__actor__": + _load_dependencies() class Future: """ @@ -244,13 +253,8 @@ def can_use(self): if result := self._lock.acquire(block=False): self._lock.release() return result - - def _load_dependencies(self): - from ..absolute_path import absolute_path - site.addsitedir(absolute_path(".python_dependencies")) def _backend_loop(self): - self._load_dependencies() while True: self._receive(self._message_queue.get()) From d5b1efd86483c807b3d82511cf7845eb6c347261 Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Thu, 1 Dec 2022 13:26:59 -0500 Subject: [PATCH 13/36] re-add other platform requirements --- __init__.py | 10 +++--- operators/install_dependencies.py | 31 ++++++++----------- .../dreamstudio.txt | 0 requirements/linux-rocm.txt | 13 ++++++++ .../mac-mps-cpu.txt | 0 requirements/win-linux-cuda.txt | 13 ++++++++ 6 files changed, 44 insertions(+), 23 deletions(-) rename requirements-dreamstudio.txt => requirements/dreamstudio.txt (100%) create mode 100644 requirements/linux-rocm.txt rename requirements.txt => requirements/mac-mps-cpu.txt (100%) create mode 100644 requirements/win-linux-cuda.txt diff --git a/__init__.py b/__init__.py index c841091e..b3e69870 100644 --- a/__init__.py +++ b/__init__.py @@ -48,10 +48,10 @@ def clear_modules(): from .ui.presets import register_default_presets requirements_path_items = ( - # Use the old version of requirements-win.txt to fix installation issues with Blender + PyTorch 1.12.1 - ('requirements-lin-win-colab-CUDA.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'), - ('requirements-mac-MPS-CPU.txt', 'Apple Silicon', 'Apple M1/M2'), - ('requirements-lin-AMD.txt', 'Linux (AMD)', 'Linux with AMD GPU'), + ('requirements/win-linux-cuda.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'), + ('requirements/mac-mps-cpu.txt', 'Apple Silicon', 'Apple M1/M2'), + ('requirements/linux-rocm.txt', 'Linux (AMD)', 'Linux with AMD GPU'), + ('requirements/dreamstudio.txt', 'DreamStudio', 'Cloud Compute Service') ) def register(): @@ -61,7 +61,7 @@ def register(): if hasattr(bpy.types, dt_op.idname()): # objects under bpy.ops are created on the fly, have to check that it actually exists a little differently raise RuntimeError("Another instance of Dream Textures is already running.") - bpy.types.Scene.dream_textures_requirements_path = EnumProperty(name="Platform", items=requirements_path_items, description="Specifies which set of dependencies to install", default='requirements-mac-MPS-CPU.txt' if sys.platform == 'darwin' else 'requirements-lin-win-colab-CUDA.txt') + bpy.types.Scene.dream_textures_requirements_path = EnumProperty(name="Platform", items=requirements_path_items, description="Specifies which set of dependencies to install", default='requirements/mac-mps-cpu.txt' if sys.platform == 'darwin' else 'requirements/win-linux-cuda.txt') for cls in PREFERENCE_CLASSES: bpy.utils.register_class(cls) diff --git a/operators/install_dependencies.py b/operators/install_dependencies.py index 855e0c20..4fc544e9 100644 --- a/operators/install_dependencies.py +++ b/operators/install_dependencies.py @@ -112,27 +112,22 @@ def install_and_import_requirements(requirements_txt=None, pip_install=PipInstal print("downloading additional include files") python_devel_tgz_path = absolute_path('python-devel.tgz') response = requests.get(f"https://www.python.org/ftp/python/{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}/Python-{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}.tgz") - open(python_devel_tgz_path, 'wb').write(response.content) - python_devel_tgz = tarfile.open(python_devel_tgz_path) - def members(tf): - prefix = f"Python-{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}/Include/" - l = len(prefix) - for member in tf.getmembers(): - if member.path.startswith(prefix): - member.path = member.path[l:] - yield member - python_devel_tgz.extractall(path=python_include_dir, members=members(python_devel_tgz)) + with open(python_devel_tgz_path, 'wb') as f: + f.write(response.content) + with tarfile.open(python_devel_tgz_path) as python_devel_tgz: + def members(tf): + prefix = f"Python-{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}/Include/" + l = len(prefix) + for member in tf.getmembers(): + if member.path.startswith(prefix): + member.path = member.path[l:] + yield member + python_devel_tgz.extractall(path=python_include_dir, members=members(python_devel_tgz)) + os.remove(python_devel_tgz_path) else: print(f"skipping include files, can't write to {python_include_dir}",file=sys.stderr) - requirements_path = requirements_txt - if requirements_path is None: - if sys.platform == 'darwin': # Use MPS dependencies list on macOS - requirements_path = 'stable_diffusion/requirements-mac-MPS-CPU.txt' - else: # Use CUDA dependencies by default on Linux/Windows. - # These are not the submodule dependencies from the `development` branch, but use the `main` branch deps for PyTorch 1.11.0. - requirements_path = 'requirements-win-torch-1-11-0.txt' - subprocess.run([sys.executable, "-m", "pip", "install", "-r", absolute_path(requirements_path), "--upgrade", "--no-cache-dir", "--target", absolute_path('.python_dependencies')], check=True, env=environ_copy, cwd=absolute_path("stable_diffusion/")) + subprocess.run([sys.executable, "-m", "pip", "install", "-r", absolute_path(requirements_txt), "--upgrade", "--no-cache-dir", "--target", absolute_path('.python_dependencies')], check=True, env=environ_copy, cwd=absolute_path("")) class InstallDependencies(bpy.types.Operator): bl_idname = "stable_diffusion.install_dependencies" diff --git a/requirements-dreamstudio.txt b/requirements/dreamstudio.txt similarity index 100% rename from requirements-dreamstudio.txt rename to requirements/dreamstudio.txt diff --git a/requirements/linux-rocm.txt b/requirements/linux-rocm.txt new file mode 100644 index 00000000..3d627876 --- /dev/null +++ b/requirements/linux-rocm.txt @@ -0,0 +1,13 @@ +git+https://github.com/huggingface/diffusers@main#egg=diffusers +transformers +accelerate +huggingface_hub + +--extra-index-url https://download.pytorch.org/whl/rocm5.2/ +torch>=1.13 + +scipy # LMSDiscreteScheduler + +stability-sdk==0.2.6 # DreamStudio + +opencolorio==2.1.2 # color management \ No newline at end of file diff --git a/requirements.txt b/requirements/mac-mps-cpu.txt similarity index 100% rename from requirements.txt rename to requirements/mac-mps-cpu.txt diff --git a/requirements/win-linux-cuda.txt b/requirements/win-linux-cuda.txt new file mode 100644 index 00000000..0fb9b569 --- /dev/null +++ b/requirements/win-linux-cuda.txt @@ -0,0 +1,13 @@ +git+https://github.com/huggingface/diffusers@main#egg=diffusers +transformers +accelerate +huggingface_hub + +--extra-index-url https://download.pytorch.org/whl/cu117 +torch>=1.13 + +scipy # LMSDiscreteScheduler + +stability-sdk==0.2.6 # DreamStudio + +opencolorio==2.1.2 # color management \ No newline at end of file From c778402b940944382a9727c4f02688d117f4a37f Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Thu, 1 Dec 2022 23:14:10 -0500 Subject: [PATCH 14/36] fix can_use() and cancel/stop operators --- generator_process/actions/prompt_to_image.py | 2 +- generator_process/actor.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 2e96c6ac..b673a7dc 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -63,7 +63,7 @@ def can_use(self, property, device) -> bool: return False if isinstance(self.__annotations__.get(property, None), _AnnotatedAlias): annotation: _AnnotatedAlias = self.__annotations__[property] - return annotation.__metadata__ == device + return annotation.__metadata__[0] == device return True def apply(self, pipeline, device): diff --git a/generator_process/actor.py b/generator_process/actor.py index 1212b624..bbf84249 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -38,6 +38,7 @@ def __init__(self): self._exception = None self.done = False self.cancelled = False + self.actor = actor def result(self): """ @@ -77,8 +78,7 @@ def _done(_): return self._exception def cancel(self): - self._cancelled = True - self.set_done() + self.cancelled = True def add_response(self, response): """ @@ -188,10 +188,10 @@ class Actor: "shared" } - def __init__(self, context: ActorContext, message_queue: Queue = Queue(maxsize=1), response_queue: Queue = Queue(maxsize=1)): + def __init__(self, context: ActorContext, message_queue: Queue = None, response_queue: Queue = None): self.context = context - self._message_queue = message_queue - self._response_queue = response_queue + self._message_queue = message_queue if message_queue is not None else Queue(maxsize=1) + self._response_queue = response_queue if response_queue is not None else Queue(maxsize=1) self._setup() self.__class__._shared_instance = self @@ -265,7 +265,7 @@ def _receive(self, message: Message): for res in iter(response): extra_message = None try: - self._message_queue.get(block=False) + extra_message = self._message_queue.get(block=False) except: pass if extra_message == Message.CANCEL: From c44b8ac33b75c3b60039121ecdf2aae9ccb46d7a Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Fri, 2 Dec 2022 18:41:02 -0500 Subject: [PATCH 15/36] fix non-square images --- generator_process/actor.py | 1 - operators/dream_texture.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/generator_process/actor.py b/generator_process/actor.py index bbf84249..59e451e7 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -38,7 +38,6 @@ def __init__(self): self._exception = None self.done = False self.cancelled = False - self.actor = actor def result(self): """ diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 0d2308a3..bbd735dc 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -101,7 +101,7 @@ def step_callback(_, step_image): nonlocal last_data_block if last_data_block is not None: bpy.data.images.remove(last_data_block) - last_data_block = bpy_image("Step", step_image.shape[0], step_image.shape[1], step_image.ravel()) + last_data_block = bpy_image("Step", step_image.shape[1], step_image.shape[0], step_image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = last_data_block @@ -114,7 +114,7 @@ def done_callback(future): image = image[-1] if last_data_block is not None: bpy.data.images.remove(last_data_block) - last_data_block = bpy_image("Final", image.shape[0], image.shape[1], image.ravel()) + last_data_block = bpy_image("Final", image.shape[1], image.shape[0], image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = last_data_block From 2af731febe825c30bedab6a8432096fb2c178a54 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 13:51:39 -0500 Subject: [PATCH 16/36] Update optimization options --- generator_process/actions/image_to_image.py | 2 +- generator_process/actions/inpaint.py | 2 +- generator_process/actions/prompt_to_image.py | 16 +++++++++++----- ui/panels/dream_texture.py | 3 +-- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index 148b54c4..7505465f 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -183,7 +183,7 @@ def __call__( _configure_model_padding(pipe.vae, seamless, seamless_axes) # Inference - with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + with (torch.inference_mode() if device != 'mps' else nullcontext()), \ (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): init_image = (Image.open(image) if isinstance(image, str) else Image.fromarray(image)).convert('RGB') print(fit) diff --git a/generator_process/actions/inpaint.py b/generator_process/actions/inpaint.py index a80b4209..033d7a4e 100644 --- a/generator_process/actions/inpaint.py +++ b/generator_process/actions/inpaint.py @@ -225,7 +225,7 @@ def __call__( _configure_model_padding(pipe.vae, seamless, seamless_axes) # Inference - with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ + with (torch.inference_mode() if device != 'mps' else nullcontext()), \ (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): init_image = Image.open(image) if isinstance(image, str) else Image.fromarray(image) yield from pipe( diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index b673a7dc..c8f706ee 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -47,7 +47,6 @@ def scheduler_class(): class Optimizations: attention_slicing: bool = True attention_slice_size: Union[str, int] = "auto" - inference_mode: Annotated[bool, "cuda"] = True cudnn_benchmark: Annotated[bool, "cuda"] = False tf32: Annotated[bool, "cuda"] = False amp: Annotated[bool, "cuda"] = False @@ -106,6 +105,11 @@ class StepPreviewMode(enum.Enum): FAST = "Fast" ACCURATE = "Accurate" +@dataclass +class ImageGenerationOutput: + image: NDArray + seed: int + def choose_device(self) -> str: """ Automatically select which PyTorch device to use. @@ -278,7 +282,7 @@ def __call__( # NOTE: Modified to yield the decoded image as a numpy array. yield from [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + ImageGenerationOutput(np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., generator.initial_seed()) for image in self.numpy_to_pil(image) ] @@ -319,15 +323,17 @@ def __call__( pipe = optimizations.apply(pipe, device) # RNG - generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + generator = torch.Generator(device=device) + if seed is not None: + generator = generator.manual_seed(seed) # Seamless _configure_model_padding(pipe.unet, seamless, seamless_axes) _configure_model_padding(pipe.vae, seamless, seamless_axes) # Inference - with (torch.inference_mode() if optimizations.can_use("inference_mode", device) else nullcontext()), \ - (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + with (torch.inference_mode() if device != 'mps' else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): yield from pipe( prompt=prompt, height=height, diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 5e583989..842a63fb 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -283,11 +283,11 @@ def optimization(prop): if hasattr(prompt, f"optimizations_{prop}"): layout.prop(prompt, f"optimizations_{prop}") - optimization("inference_mode") optimization("cudnn_benchmark") optimization("tf32") optimization("amp") optimization("half_precision") + optimization("channels_last_memory_format") yield SpeedOptimizationPanel class MemoryOptimizationPanel(sub_panel): @@ -312,7 +312,6 @@ def optimization(prop): if prompt.optimizations_attention_slice_size_src == 'manual': slice_size_row.prop(prompt, "optimizations_attention_slice_size", text="Size") optimization("sequential_cpu_offload") - optimization("channels_last_memory_format") optimization("cpu_only") # optimization("xformers_attention") # FIXME: xFormers is not yet available. yield MemoryOptimizationPanel From 0c61680e980541c4bac4f51559c059f52e4cb120 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 14:14:08 -0500 Subject: [PATCH 17/36] Add proper file naming --- generator_process/actions/prompt_to_image.py | 35 +++++++++++++++----- operators/dream_texture.py | 10 +++--- 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index c8f706ee..1966a809 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -5,6 +5,7 @@ from contextlib import nullcontext from numpy.typing import NDArray import numpy as np +import random class Pipeline(enum.IntEnum): STABLE_DIFFUSION = 0 @@ -106,9 +107,11 @@ class StepPreviewMode(enum.Enum): ACCURATE = "Accurate" @dataclass -class ImageGenerationOutput: +class ImageGenerationResult: image: NDArray seed: int + step: int + final: bool def choose_device(self) -> str: """ @@ -175,7 +178,7 @@ def prompt_to_image( step_preview_mode: StepPreviewMode, **kwargs -) -> Generator[NDArray, None, None]: +) -> Generator[ImageGenerationResult, None, None]: match pipeline: case Pipeline.STABLE_DIFFUSION: import diffusers @@ -266,10 +269,20 @@ def __call__( case StepPreviewMode.NONE: pass case StepPreviewMode.FAST: - yield np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255. + yield ImageGenerationResult( + np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + i, + False + ) case StepPreviewMode.ACCURATE: yield from [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + i, + False + ) for image in self.numpy_to_pil(self.decode_latents(latents)) ] @@ -282,7 +295,12 @@ def __call__( # NOTE: Modified to yield the decoded image as a numpy array. yield from [ - ImageGenerationOutput(np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., generator.initial_seed()) + ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + num_inference_steps, + True + ) for image in self.numpy_to_pil(image) ] @@ -323,9 +341,10 @@ def __call__( pipe = optimizations.apply(pipe, device) # RNG - generator = torch.Generator(device=device) - if seed is not None: - generator = generator.manual_seed(seed) + generator = torch.Generator(device="cpu" if device == "mps" else device) # MPS does not support the `Generator` API + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + generator = generator.manual_seed(seed) # Seamless _configure_model_padding(pipe.unet, seamless, seamless_axes) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index bbd735dc..120c5eb6 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -14,7 +14,7 @@ from ..prompt_engineering import * from ..absolute_path import WEIGHTS_PATH, CLIPSEG_WEIGHTS_PATH from ..generator_process import Generator -from ..generator_process.actions.prompt_to_image import Pipeline, Optimizations +from ..generator_process.actions.prompt_to_image import Pipeline, Optimizations, ImageGenerationResult import tempfile @@ -97,11 +97,11 @@ def execute(self, context): init_image = save_temp_image(area.spaces.active.image) last_data_block = None - def step_callback(_, step_image): + def step_callback(_, step_image: ImageGenerationResult): nonlocal last_data_block if last_data_block is not None: bpy.data.images.remove(last_data_block) - last_data_block = bpy_image("Step", step_image.shape[1], step_image.shape[0], step_image.ravel()) + last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = last_data_block @@ -109,12 +109,12 @@ def step_callback(_, step_image): def done_callback(future): nonlocal last_data_block del gen._active_generation_future - image = future.result() + image: ImageGenerationResult | list = future.result() if isinstance(image, list): image = image[-1] if last_data_block is not None: bpy.data.images.remove(last_data_block) - last_data_block = bpy_image("Final", image.shape[1], image.shape[0], image.ravel()) + last_data_block = bpy_image(str(image.seed), image.image.shape[1], image.image.shape[0], image.image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = last_data_block From e08eaad6155c24b55b75a8d62dfa94fb3cc71307 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 14:23:17 -0500 Subject: [PATCH 18/36] Fix data block removal error --- operators/dream_texture.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 104e36d3..5f6e04e1 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -102,6 +102,7 @@ def step_callback(_, step_image: ImageGenerationResult): nonlocal last_data_block if last_data_block is not None: bpy.data.images.remove(last_data_block) + last_data_block = None last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': @@ -117,6 +118,7 @@ def done_callback(future): image_result = image_result[-1] if last_data_block is not None: bpy.data.images.remove(last_data_block) + last_data_block = None image = bpy_image(str(image_result.seed), image_result.image.shape[1], image_result.image.shape[0], image_result.image.ravel()) if node_tree is not None: nodes = node_tree.nodes From eb41b2269a552e9f6feb4bff1cb2e2a576e90a8c Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 15:53:55 -0500 Subject: [PATCH 19/36] Add upscale tile blending with the 'tiler' --- __init__.py | 1 + generator_process/actions/upscale.py | 53 ++++++++++++++++++---------- operators/upscale.py | 3 +- requirements/mac-mps-cpu.txt | 6 ++-- ui/panels/upscaling.py | 1 + 5 files changed, 42 insertions(+), 22 deletions(-) diff --git a/__init__.py b/__init__.py index b3e69870..41969c34 100644 --- a/__init__.py +++ b/__init__.py @@ -88,6 +88,7 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_upscale_prompt = PointerProperty(type=DreamPrompt) bpy.types.Scene.dream_textures_upscale_tile_size = IntProperty(name="Tile Size", default=128, step=64, min=64, max=512) + bpy.types.Scene.dream_textures_upscale_blend = IntProperty(name="Blend", default=32, step=8, min=0, max=512) for cls in CLASSES: bpy.utils.register_class(cls) diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index f97e6d59..619d0073 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -1,5 +1,6 @@ import numpy as np from .prompt_to_image import Optimizations, Scheduler +import random def upscale( self, @@ -12,6 +13,7 @@ def upscale( scheduler: Scheduler, tile_size: int, + blend: int, optimizations: Optimizations, @@ -20,6 +22,7 @@ def upscale( from PIL import Image, ImageOps import torch import diffusers + from tiler import Tiler, Merger if optimizations.cpu_only: device = "cpu" @@ -37,24 +40,36 @@ def upscale( low_res_image = Image.open(image) - generator = torch.Generator() if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + generator = torch.Generator(device="cpu" if device == "mps" else device) # MPS does not support the `Generator` API + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + generator = generator.manual_seed(seed) initial_seed = generator.initial_seed() - final = Image.new('RGB', (low_res_image.size[0] * 4, low_res_image.size[1] * 4)) - for x in range(low_res_image.size[0] // tile_size): - for y in range(low_res_image.size[1] // tile_size): - x_offset = x * tile_size - y_offset = y * tile_size - tile = low_res_image.crop((x_offset, y_offset, x_offset + tile_size, y_offset + tile_size)) - upscaled = pipe( - prompt=prompt, - image=tile, - num_inference_steps=steps, - generator=torch.manual_seed(initial_seed), - guidance_scale=cfg_scale, - - ).images[0] - final.paste(upscaled, (x_offset * 4, y_offset * 4)) - yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. - - yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. \ No newline at end of file + # final = Image.new('RGB', (low_res_image.size[0] * 4, low_res_image.size[1] * 4)) + tiler = Tiler( + data_shape=(low_res_image.size[0], low_res_image.size[1], len(low_res_image.getbands())), + tile_shape=(tile_size, tile_size, len(low_res_image.getbands())), + overlap=(blend, blend, 0), + channel_dimension=2 + ) + merger = Merger(Tiler( + data_shape=(low_res_image.size[0] * 4, low_res_image.size[1] * 4, 3), + tile_shape=(tile_size * 4, tile_size * 4, 3), + overlap=(blend * 4, blend * 4, 0), + channel_dimension=2 + )) + input_array = np.array(low_res_image) + for id, tile in tiler(input_array, progress_bar=True): + merger.add(id, np.array(pipe( + prompt=prompt, + image=Image.fromarray(tile), + num_inference_steps=steps, + generator=torch.manual_seed(initial_seed), + guidance_scale=cfg_scale, + ).images[0])) + final = Image.fromarray(merger.merge().astype(np.uint8)) + yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. + + # final = Image.fromarray(merger.merge().astype(np.uint8)) + # yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. \ No newline at end of file diff --git a/operators/upscale.py b/operators/upscale.py index 61951650..68cba213 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -12,7 +12,7 @@ class Upscale(bpy.types.Operator): bl_idname = "shade.dream_textures_upscale" bl_label = "Upscale" - bl_description = ("Upscale with Real-ESRGAN") + bl_description = ("Upscale with Stable Diffusion x4 Upscaler") bl_options = {"REGISTER"} @classmethod @@ -94,6 +94,7 @@ def image_done(future): f = gen.upscale( image=input_image_path, tile_size=context.scene.dream_textures_upscale_tile_size, + blend=context.scene.dream_textures_upscale_blend, **context.scene.dream_textures_upscale_prompt.generate_args() ) f.add_response_callback(on_tile_complete) diff --git a/requirements/mac-mps-cpu.txt b/requirements/mac-mps-cpu.txt index 1519363c..9ef0f72a 100644 --- a/requirements/mac-mps-cpu.txt +++ b/requirements/mac-mps-cpu.txt @@ -1,6 +1,6 @@ git+https://github.com/huggingface/diffusers@main#egg=diffusers transformers -accelerate +accelerate==0.14.0 huggingface_hub torch>=1.13 @@ -9,4 +9,6 @@ scipy # LMSDiscreteScheduler stability-sdk==0.2.6 # DreamStudio -opencolorio==2.1.2 # color management \ No newline at end of file +opencolorio==2.1.2 # color management + +tiler # Upscaler tiling \ No newline at end of file diff --git a/ui/panels/upscaling.py b/ui/panels/upscaling.py index bc5b1c3c..c83a5262 100644 --- a/ui/panels/upscaling.py +++ b/ui/panels/upscaling.py @@ -35,6 +35,7 @@ def draw(self, context): layout.prop(prompt, "prompt_structure_token_subject") layout.prop(context.scene, "dream_textures_upscale_tile_size") + layout.prop(context.scene, "dream_textures_upscale_blend") if context.scene.dream_textures_upscale_tile_size > 128: warning_box = layout.box() From 333baf4a727f70ecba59a496484714e8325d54c7 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 15:54:23 -0500 Subject: [PATCH 20/36] Add 'tiler' to other requirements files --- requirements/linux-rocm.txt | 4 +++- requirements/win-linux-cuda.txt | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/requirements/linux-rocm.txt b/requirements/linux-rocm.txt index 3d627876..65a7f068 100644 --- a/requirements/linux-rocm.txt +++ b/requirements/linux-rocm.txt @@ -10,4 +10,6 @@ scipy # LMSDiscreteScheduler stability-sdk==0.2.6 # DreamStudio -opencolorio==2.1.2 # color management \ No newline at end of file +opencolorio==2.1.2 # color management + +tiler # Upscaler tiling \ No newline at end of file diff --git a/requirements/win-linux-cuda.txt b/requirements/win-linux-cuda.txt index 0fb9b569..540458a2 100644 --- a/requirements/win-linux-cuda.txt +++ b/requirements/win-linux-cuda.txt @@ -10,4 +10,6 @@ scipy # LMSDiscreteScheduler stability-sdk==0.2.6 # DreamStudio -opencolorio==2.1.2 # color management \ No newline at end of file +opencolorio==2.1.2 # color management + +tiler # Upscaler tiling \ No newline at end of file From 191e00997c9c195b082ce2591b0378d4d5e166a6 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 20:38:18 -0500 Subject: [PATCH 21/36] Add upscale naming --- generator_process/actions/upscale.py | 32 +++++++++++++++++++++++----- operators/upscale.py | 24 ++++++++++++++++----- 2 files changed, 46 insertions(+), 10 deletions(-) diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index 619d0073..e2033609 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -1,6 +1,15 @@ import numpy as np -from .prompt_to_image import Optimizations, Scheduler +from .prompt_to_image import Optimizations, Scheduler, StepPreviewMode import random +from dataclasses import dataclass +from numpy.typing import NDArray + +@dataclass +class ImageUpscaleResult: + image: NDArray + tile: int + total: int + final: bool def upscale( self, @@ -17,6 +26,8 @@ def upscale( optimizations: Optimizations, + step_preview_mode: StepPreviewMode, + **kwargs ): from PIL import Image, ImageOps @@ -68,8 +79,19 @@ def upscale( generator=torch.manual_seed(initial_seed), guidance_scale=cfg_scale, ).images[0])) + if step_preview_mode != StepPreviewMode.NONE: + step = Image.fromarray(merger.merge().astype(np.uint8)) + yield ImageUpscaleResult( + np.asarray(ImageOps.flip(step).convert('RGBA'), dtype=np.float32) / 255., + id + 1, + tiler.n_tiles, + (id + 1) == tiler.n_tiles + ) + if step_preview_mode == StepPreviewMode.NONE: final = Image.fromarray(merger.merge().astype(np.uint8)) - yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. - - # final = Image.fromarray(merger.merge().astype(np.uint8)) - # yield np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255. \ No newline at end of file + yield ImageUpscaleResult( + np.asarray(ImageOps.flip(final).convert('RGBA'), dtype=np.float32) / 255., + tiler.n_tiles, + tiler.n_tiles, + True + ) \ No newline at end of file diff --git a/operators/upscale.py b/operators/upscale.py index 68cba213..7c1d597c 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -2,6 +2,7 @@ import tempfile from ..prompt_engineering import custom_structure from ..generator_process import Generator +from ..generator_process.actions.upscale import ImageUpscaleResult upscale_options = [ ("2", "2x", "", 2), @@ -75,15 +76,28 @@ def bpy_image(name, width, height, pixels): image.pack() return image - def on_tile_complete(_, tile): - image = bpy_image("diffusers-upscaled", tile.shape[0], tile.shape[1], tile.ravel()) + last_data_block = None + def on_tile_complete(_, tile: ImageUpscaleResult): + nonlocal last_data_block + if last_data_block is not None: + bpy.data.images.remove(last_data_block) + last_data_block = None + if tile.final: + return + last_data_block = bpy_image(f"Tile {tile.tile}/{tile.total}", tile.image.shape[0], tile.image.shape[1], tile.image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = image + area.spaces.active.image = last_data_block def image_done(future): - image = future.result()[-1] - image = bpy_image("diffusers-upscaled", image.shape[0], image.shape[1], image.ravel()) + nonlocal last_data_block + if last_data_block is not None: + bpy.data.images.remove(last_data_block) + last_data_block = None + tile: ImageUpscaleResult = future.result() + if isinstance(tile, list): + tile = tile[-1] + image = bpy_image(f"{input_image.name} (Upscaled)", tile.image.shape[0], tile.image.shape[1], tile.image.ravel()) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = image From 61cd13aa3b4a8e51fe59d2f73dc0e06790ff4dcb Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 21:00:00 -0500 Subject: [PATCH 22/36] Add progress bar --- classes.py | 3 +- generator_process/actions/prompt_to_image.py | 9 +- generator_process/actions/upscale.py | 15 +- operators/dream_texture.py | 194 ++----------------- 4 files changed, 29 insertions(+), 192 deletions(-) diff --git a/classes.py b/classes.py index 30c7327c..c7e37958 100644 --- a/classes.py +++ b/classes.py @@ -1,6 +1,6 @@ from .operators.install_dependencies import InstallDependencies from .operators.open_latest_version import OpenLatestVersion -from .operators.dream_texture import DreamTexture, ReleaseGenerator, HeadlessDreamTexture, CancelGenerator +from .operators.dream_texture import DreamTexture, ReleaseGenerator, CancelGenerator from .operators.view_history import SCENE_UL_HistoryList, RecallHistoryEntry, ClearHistory, RemoveHistorySelection, ExportHistorySelection, ImportPromptFile from .operators.inpaint_area_brush import InpaintAreaBrushActivated from .operators.upscale import Upscale @@ -11,7 +11,6 @@ from .ui.presets import DREAM_PT_AdvancedPresets, DREAM_MT_AdvancedPresets, AddAdvancedPreset, RestoreDefaultPresets CLASSES = ( - HeadlessDreamTexture, *render_properties.render_properties_panels(), DreamTexture, diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index b3d45d61..8739f752 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -108,7 +108,7 @@ class StepPreviewMode(enum.Enum): @dataclass class ImageGenerationResult: - image: NDArray + image: NDArray | None seed: int step: int final: bool @@ -267,7 +267,12 @@ def __call__( # NOTE: Modified to yield the latents instead of calling a callback. match kwargs['step_preview_mode']: case StepPreviewMode.NONE: - pass + yield ImageGenerationResult( + None, + generator.initial_seed(), + i, + False + ) case StepPreviewMode.FAST: yield ImageGenerationResult( np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255., diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index e2033609..5d759ab2 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -6,7 +6,7 @@ @dataclass class ImageUpscaleResult: - image: NDArray + image: NDArray | None tile: int total: int final: bool @@ -57,7 +57,6 @@ def upscale( generator = generator.manual_seed(seed) initial_seed = generator.initial_seed() - # final = Image.new('RGB', (low_res_image.size[0] * 4, low_res_image.size[1] * 4)) tiler = Tiler( data_shape=(low_res_image.size[0], low_res_image.size[1], len(low_res_image.getbands())), tile_shape=(tile_size, tile_size, len(low_res_image.getbands())), @@ -81,12 +80,12 @@ def upscale( ).images[0])) if step_preview_mode != StepPreviewMode.NONE: step = Image.fromarray(merger.merge().astype(np.uint8)) - yield ImageUpscaleResult( - np.asarray(ImageOps.flip(step).convert('RGBA'), dtype=np.float32) / 255., - id + 1, - tiler.n_tiles, - (id + 1) == tiler.n_tiles - ) + yield ImageUpscaleResult( + (np.asarray(ImageOps.flip(step).convert('RGBA'), dtype=np.float32) / 255.) if step is not None else None, + id + 1, + tiler.n_tiles, + (id + 1) == tiler.n_tiles + ) if step_preview_mode == StepPreviewMode.NONE: final = Image.fromarray(merger.merge().astype(np.uint8)) yield ImageUpscaleResult( diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 5f6e04e1..0f6a14bb 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -97,16 +97,22 @@ def execute(self, context): if area.spaces.active.image is not None: init_image = save_temp_image(area.spaces.active.image) + # Setup the progress indicator + bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=0, min=0, max=generated_args['steps']) + scene.dream_textures_info = "Starting..." + last_data_block = None def step_callback(_, step_image: ImageGenerationResult): nonlocal last_data_block if last_data_block is not None: bpy.data.images.remove(last_data_block) last_data_block = None - last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = last_data_block + scene.dream_textures_progress = step_image.step + if step_image.image is not None: + last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = last_data_block iteration = 0 def done_callback(future): @@ -140,6 +146,9 @@ def done_callback(future): iteration += 1 if iteration < generated_args['iterations']: generate_next() + else: + scene.dream_textures_info = "" + scene.dream_textures_progress = 0 gen = Generator.shared() def generate_next(): @@ -170,183 +179,8 @@ def generate_next(): generate_next() return {"FINISHED"} -headless_prompt = None -headless_step_callback = None -headless_image_callback = None -headless_init_img = None -headless_args = None -def dream_texture(prompt, step_callback, image_callback, init_img=None, **kwargs): - global headless_prompt - headless_prompt = prompt - global headless_step_callback - headless_step_callback = step_callback - global headless_image_callback - headless_image_callback = image_callback - global headless_init_img - headless_init_img = init_img - global headless_args - headless_args = kwargs - bpy.ops.shade.dream_texture_headless() - -class HeadlessDreamTexture(bpy.types.Operator): - bl_idname = "shade.dream_texture_headless" - bl_label = "Headless Dream Texture" - bl_description = "Generate a texture with AI" - bl_options = {'REGISTER'} - - @classmethod - def poll(cls, context): - return GeneratorProcess.can_use() - - def modal(self, context, event): - if event.type != 'TIMER': - return {'PASS_THROUGH'} - try: - next(generator_advance) - except StopIteration: - modal_stopped(context) - return {'FINISHED'} - except Exception as e: - modal_stopped(context) - raise e - return {'RUNNING_MODAL'} - - def execute(self, context): - global headless_prompt - screen = context.screen - scene = context.scene - - global headless_init_img - init_img = headless_init_img or (scene.init_img if headless_prompt.use_init_img and headless_prompt.init_img_src == 'file' else None) - - def info(msg=""): - scene.dream_textures_info = msg - - def handle_exception(fatal, msg, trace): - info() # clear variable - if fatal: - kill_generator() - self.report({'ERROR'},msg) - if trace: - print(trace, file=sys.stderr) - if msg == MISSING_DEPENDENCIES_ERROR: - from .open_latest_version import do_force_show_download - do_force_show_download() - - def step_progress_update(self, context): - if hasattr(context.area, "regions"): - for region in context.area.regions: - if region.type == "UI": - region.tag_redraw() - return None - - bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty( - name="", - default=0, - min=0, - max=(int(headless_prompt.strength * headless_prompt.steps) if init_img is not None else headless_prompt.steps) + 1, - update=step_progress_update - ) - bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info", update=step_progress_update) - - info("Waiting For Process") - if len(backend_options(self, context)) <= 1: - headless_prompt.backend = backend_options(self, context)[0][0] - generator = GeneratorProcess.shared(backend=BackendTarget[headless_prompt.backend]) - - if not generator.backend.color_correction(): - headless_prompt.use_init_img_color = False - - def save_temp_image(img, path=None): - path = path if path is not None else tempfile.NamedTemporaryFile().name - - settings = scene.render.image_settings - file_format = settings.file_format - mode = settings.color_mode - depth = settings.color_depth - - settings.file_format = 'PNG' - settings.color_mode = 'RGBA' - settings.color_depth = '8' - - img.save_render(path) - - settings.file_format = file_format - settings.color_mode = mode - settings.color_depth = depth - - return path - - if headless_prompt.use_init_img and headless_prompt.init_img_src == 'open_editor': - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - if area.spaces.active.image is not None: - init_img = area.spaces.active.image - init_img_path = None - if init_img is not None: - init_img_path = save_temp_image(init_img) - - args = headless_prompt.generate_args() - args.update(headless_args) - if headless_init_img is not None: - args['use_init_img'] = True - if args['prompt_structure'] == file_batch_structure.id: - args['prompt'] = [line.body for line in scene.dream_textures_prompt_file.lines if len(line.body.strip()) > 0] - args['init_img'] = init_img_path - if args['use_init_img_color']: - args['init_color'] = init_img_path - if args['backend'] == BackendTarget.STABILITY_SDK.name: - args['dream_studio_key'] = context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.dream_studio_key - - def step_callback(step, width=None, height=None, shared_memory_name=None): - global headless_step_callback - info() # clear variable - scene.dream_textures_progress = step + 1 - headless_step_callback(step, width, height, shared_memory_name) - - received_noncolorized = False - def image_callback(shared_memory_name, seed, width, height, upscaled=False): - global headless_image_callback - info() # clear variable - nonlocal received_noncolorized - if args['use_init_img'] and args['use_init_img_color'] and not received_noncolorized: - received_noncolorized = True - return - received_noncolorized = False - headless_image_callback(shared_memory_name, seed, width, height, upscaled) - - global generator_advance - generator_advance = generator.prompt2image(args, - # a function or method that will be called each step - step_callback=step_callback, - # a function or method that will be called each time an image is generated - image_callback=image_callback, - # a function or method that will recieve messages - info_callback=info, - exception_callback=handle_exception - ) - context.window_manager.modal_handler_add(self) - global timer - timer = context.window_manager.event_timer_add(1 / 15, window=context.window) - return {'RUNNING_MODAL'} - -def modal_stopped(context): - global timer - if timer: - context.window_manager.event_timer_remove(timer) - timer = None - if not hasattr(context,'scene'): - context = bpy.context # modal context is sometimes missing scene? - context.scene.dream_textures_progress = 0 - context.scene.dream_textures_info = "" - global last_data_block - if last_data_block is not None: - bpy.data.images.remove(last_data_block) - last_data_block = None - def kill_generator(context=bpy.context): Generator.shared_close() - modal_stopped(context) class ReleaseGenerator(bpy.types.Operator): bl_idname = "shade.dream_textures_release_generator" @@ -365,7 +199,7 @@ class CancelGenerator(bpy.types.Operator): bl_options = {'REGISTER'} @classmethod - def poll(self, context): + def poll(cls, context): gen = Generator.shared() return hasattr(gen, "_active_generation_future") and gen._active_generation_future is not None and not gen._active_generation_future.cancelled and not gen._active_generation_future.done From 69623eeeea2f6c1a57d0abe86a0c9c17701ccec5 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 4 Dec 2022 21:15:03 -0500 Subject: [PATCH 23/36] Add progress bar to upscaling --- operators/dream_texture.py | 29 ++++++++++++++++++++++------- operators/upscale.py | 30 +++++++++++++++++++++++++----- ui/panels/upscaling.py | 38 ++++++++++++++++++++++++-------------- 3 files changed, 71 insertions(+), 26 deletions(-) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 0f6a14bb..cfc53e51 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -98,7 +98,13 @@ def execute(self, context): init_image = save_temp_image(area.spaces.active.image) # Setup the progress indicator - bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=0, min=0, max=generated_args['steps']) + def step_progress_update(self, context): + if hasattr(context.area, "regions"): + for region in context.area.regions: + if region.type == "UI": + region.tag_redraw() + return None + bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=0, min=0, max=generated_args['steps'], update=step_progress_update) scene.dream_textures_info = "Starting..." last_data_block = None @@ -107,12 +113,14 @@ def step_callback(_, step_image: ImageGenerationResult): if last_data_block is not None: bpy.data.images.remove(last_data_block) last_data_block = None - scene.dream_textures_progress = step_image.step - if step_image.image is not None: - last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = last_data_block + def update_progress(): + scene.dream_textures_progress = step_image.step + if step_image.image is not None: + last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = last_data_block + bpy.app.timers.register(update_progress) iteration = 0 def done_callback(future): @@ -181,6 +189,11 @@ def generate_next(): def kill_generator(context=bpy.context): Generator.shared_close() + try: + context.scene.dream_textures_info = "" + context.scene.dream_textures_progress = 0 + except: + pass class ReleaseGenerator(bpy.types.Operator): bl_idname = "shade.dream_textures_release_generator" @@ -206,4 +219,6 @@ def poll(cls, context): def execute(self, context): gen = Generator.shared() gen._active_generation_future.cancel() + context.scene.dream_textures_info = "" + context.scene.dream_textures_progress = 0 return {'FINISHED'} diff --git a/operators/upscale.py b/operators/upscale.py index 7c1d597c..e52c4586 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -22,6 +22,7 @@ def poll(cls, context): def execute(self, context): screen = context.screen + scene = context.scene 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 @@ -76,18 +77,35 @@ def bpy_image(name, width, height, pixels): image.pack() return image + generated_args = context.scene.dream_textures_upscale_prompt.generate_args() + + # Setup the progress indicator + def step_progress_update(self, context): + if hasattr(context.area, "regions"): + for region in context.area.regions: + if region.type == "UI": + region.tag_redraw() + return None + bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=0, min=0, max=generated_args['steps'], update=step_progress_update) + scene.dream_textures_info = "Starting..." + last_data_block = None def on_tile_complete(_, tile: ImageUpscaleResult): nonlocal last_data_block if last_data_block is not None: bpy.data.images.remove(last_data_block) last_data_block = None + else: + bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=tile.tile, min=0, max=tile.total, update=step_progress_update) if tile.final: return - last_data_block = bpy_image(f"Tile {tile.tile}/{tile.total}", tile.image.shape[0], tile.image.shape[1], tile.image.ravel()) - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = last_data_block + def update_progress(): + scene.dream_textures_progress = tile.tile + last_data_block = bpy_image(f"Tile {tile.tile}/{tile.total}", tile.image.shape[0], tile.image.shape[1], tile.image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = last_data_block + bpy.app.timers.register(update_progress) def image_done(future): nonlocal last_data_block @@ -103,13 +121,15 @@ def image_done(future): area.spaces.active.image = image if active_node is not None: active_node.image = image + scene.dream_textures_info = "" + scene.dream_textures_progress = 0 gen = Generator.shared() context.scene.dream_textures_upscale_prompt.prompt_structure = custom_structure.id f = gen.upscale( image=input_image_path, tile_size=context.scene.dream_textures_upscale_tile_size, blend=context.scene.dream_textures_upscale_blend, - **context.scene.dream_textures_upscale_prompt.generate_args() + **generated_args ) f.add_response_callback(on_tile_complete) f.add_done_callback(image_done) diff --git a/ui/panels/upscaling.py b/ui/panels/upscaling.py index c83a5262..4300c301 100644 --- a/ui/panels/upscaling.py +++ b/ui/panels/upscaling.py @@ -3,6 +3,7 @@ from ...pil_to_image import * from ...prompt_engineering import * from ...operators.upscale import Upscale +from ...operators.dream_texture import CancelGenerator, ReleaseGenerator from .dream_texture import create_panel, advanced_panel from ..space_types import SPACE_TYPES @@ -65,22 +66,31 @@ def poll(cls, context): def draw(self, context): layout = self.layout layout.use_property_split = True + layout.use_property_decorate = False - if context.scene.dream_textures_info != "": - layout.label(text=context.scene.dream_textures_info, icon="INFO") + image = None + for area in context.screen.areas: + if area.type == 'IMAGE_EDITOR': + image = area.spaces.active.image + row = layout.row() + row.scale_y = 1.5 + if context.scene.dream_textures_progress <= 0: + if context.scene.dream_textures_info != "": + row.label(text=context.scene.dream_textures_info, icon="INFO") + else: + row.operator( + Upscale.bl_idname, + text=f"Upscale to {image.size[0] * 4}x{image.size[1] * 4}" if image is not None else "Upscale", + icon="FULLSCREEN_ENTER" + ) else: - image = None - for area in context.screen.areas: - if area.type == 'IMAGE_EDITOR': - image = area.spaces.active.image - - row = layout.row() - row.scale_y = 1.5 - row.operator( - Upscale.bl_idname, - text=f"Upscale to {image.size[0] * 4}x{image.size[1] * 4}" if image is not None else "Upscale", - icon="FULLSCREEN_ENTER" - ) + disabled_row = row.row() + disabled_row.use_property_split = True + disabled_row.prop(context.scene, 'dream_textures_progress', slider=True) + disabled_row.enabled = False + if CancelGenerator.poll(context): + row.operator(CancelGenerator.bl_idname, icon="CANCEL", text="") + row.operator(ReleaseGenerator.bl_idname, icon="X", text="") yield UpscalingPanel advanced_panels = [*create_panel(space_type, 'UI', UpscalingPanel.bl_idname, advanced_panel, lambda context: context.scene.dream_textures_upscale_prompt)] outer_panel = advanced_panels[0] From c01ba6b5d0e18c703a97c2119cda31891607bbdf Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 17:35:41 -0500 Subject: [PATCH 24/36] Fix image2image --- generator_process/actions/image_to_image.py | 108 ++++++++++++-------- operators/dream_texture.py | 16 +-- operators/upscale.py | 12 +-- 3 files changed, 79 insertions(+), 57 deletions(-) diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index 7505465f..379f6d1e 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -3,7 +3,8 @@ from contextlib import nullcontext from numpy.typing import NDArray import numpy as np -from .prompt_to_image import Pipeline, Scheduler, Optimizations, StepPreviewMode, approximate_decoded_latents, _configure_model_padding +from .prompt_to_image import Pipeline, Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, approximate_decoded_latents, _configure_model_padding +import random def image_to_image( self, @@ -42,6 +43,7 @@ def image_to_image( import diffusers import torch from PIL import Image, ImageOps + import PIL.Image from ...absolute_path import WEIGHTS_PATH # Mostly copied from `diffusers.StableDiffusionImg2ImgPipeline`, with slight modifications to yield the latents at each step. @@ -50,7 +52,7 @@ class GeneratorPipeline(diffusers.StableDiffusionImg2ImgPipeline): def __call__( self, prompt: Union[str, List[str]], - init_image: Union[torch.FloatTensor, Image.Image], + image: Union[torch.FloatTensor, PIL.Image.Image], strength: float = 0.8, num_inference_steps: Optional[int] = 50, guidance_scale: Optional[float] = 7.5, @@ -64,6 +66,8 @@ def __call__( callback_steps: Optional[int] = 1, **kwargs, ): + image = init_image or image + # 1. Check inputs self.check_inputs(prompt, strength, callback_steps) @@ -81,50 +85,67 @@ def __call__( ) # 4. Preprocess image - if isinstance(init_image, Image.Image): - init_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess(init_image) + if isinstance(image, PIL.Image.Image): + image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess(image) # 5. set timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.get_timesteps(num_inference_steps, strength, device) + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) # 6. Prepare latent variables latents = self.prepare_latents( - init_image, latent_timestep, batch_size, num_images_per_prompt, text_embeddings.dtype, device, generator + image, latent_timestep, batch_size, num_images_per_prompt, text_embeddings.dtype, device, generator ) # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 8. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # NOTE: Modified to yield the latents instead of calling a callback. - match kwargs['step_preview_mode']: - case StepPreviewMode.NONE: - pass - case StepPreviewMode.FAST: - yield np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255. - case StepPreviewMode.ACCURATE: - yield from [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in self.numpy_to_pil(self.decode_latents(latents)) - ] + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # NOTE: Modified to yield the latents instead of calling a callback. + match kwargs['step_preview_mode']: + case StepPreviewMode.NONE: + yield ImageGenerationResult( + None, + generator.initial_seed(), + i, + False + ) + case StepPreviewMode.FAST: + yield ImageGenerationResult( + np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + i, + False + ) + case StepPreviewMode.ACCURATE: + yield from [ + ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + i, + False + ) + for image in self.numpy_to_pil(self.decode_latents(latents)) + ] # 9. Post-processing image = self.decode_latents(latents) @@ -135,7 +156,12 @@ def __call__( # NOTE: Modified to yield the decoded image as a numpy array. yield from [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + num_inference_steps, + True + ) for image in self.numpy_to_pil(image) ] @@ -176,7 +202,10 @@ def __call__( pipe = optimizations.apply(pipe, device) # RNG - generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + generator = torch.Generator(device="cpu" if device == "mps" else device) # MPS does not support the `Generator` API + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + generator = generator.manual_seed(seed) # Seamless _configure_model_padding(pipe.unet, seamless, seamless_axes) @@ -186,23 +215,16 @@ def __call__( with (torch.inference_mode() if device != 'mps' else nullcontext()), \ (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): init_image = (Image.open(image) if isinstance(image, str) else Image.fromarray(image)).convert('RGB') - print(fit) - print(width, height) - print(init_image) - print(init_image.size) yield from pipe( prompt=prompt, - init_image=init_image, + image=init_image if fit else init_image.resize((width, height)), strength=strength, - height=init_image.size[1] if fit else height, - width=init_image.size[0] if fit else width, num_inference_steps=steps, guidance_scale=cfg_scale, negative_prompt=negative_prompt if use_negative_prompt else None, num_images_per_prompt=iterations, eta=0.0, generator=generator, - latents=None, output_type="pil", return_dict=True, callback=None, diff --git a/operators/dream_texture.py b/operators/dream_texture.py index cfc53e51..78632d42 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -113,14 +113,14 @@ def step_callback(_, step_image: ImageGenerationResult): if last_data_block is not None: bpy.data.images.remove(last_data_block) last_data_block = None - def update_progress(): - scene.dream_textures_progress = step_image.step - if step_image.image is not None: - last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = last_data_block - bpy.app.timers.register(update_progress) + if step_image.final: + return + scene.dream_textures_progress = step_image.step + if step_image.image is not None: + last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = last_data_block iteration = 0 def done_callback(future): diff --git a/operators/upscale.py b/operators/upscale.py index e52c4586..ea418e03 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -99,12 +99,12 @@ def on_tile_complete(_, tile: ImageUpscaleResult): bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="", default=tile.tile, min=0, max=tile.total, update=step_progress_update) if tile.final: return - def update_progress(): - scene.dream_textures_progress = tile.tile - last_data_block = bpy_image(f"Tile {tile.tile}/{tile.total}", tile.image.shape[0], tile.image.shape[1], tile.image.ravel()) - for area in screen.areas: - if area.type == 'IMAGE_EDITOR': - area.spaces.active.image = last_data_block + + scene.dream_textures_progress = tile.tile + last_data_block = bpy_image(f"Tile {tile.tile}/{tile.total}", tile.image.shape[0], tile.image.shape[1], tile.image.ravel()) + for area in screen.areas: + if area.type == 'IMAGE_EDITOR': + area.spaces.active.image = last_data_block bpy.app.timers.register(update_progress) def image_done(future): From 5b5f8d3aa0470c8a66f1e843fec9195c4ee48c11 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 18:50:47 -0500 Subject: [PATCH 25/36] Fix inpaint --- generator_process/actions/image_to_image.py | 2 +- generator_process/actions/inpaint.py | 113 ++++++++++++-------- generator_process/actor.py | 9 +- operators/dream_texture.py | 12 ++- 4 files changed, 87 insertions(+), 49 deletions(-) diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index 379f6d1e..ceb25d9f 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -222,7 +222,7 @@ def __call__( num_inference_steps=steps, guidance_scale=cfg_scale, negative_prompt=negative_prompt if use_negative_prompt else None, - num_images_per_prompt=iterations, + num_images_per_prompt=1, eta=0.0, generator=generator, output_type="pil", diff --git a/generator_process/actions/inpaint.py b/generator_process/actions/inpaint.py index 033d7a4e..fd23c941 100644 --- a/generator_process/actions/inpaint.py +++ b/generator_process/actions/inpaint.py @@ -3,7 +3,8 @@ from contextlib import nullcontext from numpy.typing import NDArray import numpy as np -from .prompt_to_image import Pipeline, Scheduler, Optimizations, StepPreviewMode, approximate_decoded_latents, _configure_model_padding +import random +from .prompt_to_image import Pipeline, Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, approximate_decoded_latents, _configure_model_padding def inpaint( self, @@ -42,6 +43,7 @@ def inpaint( import diffusers import torch from PIL import Image, ImageOps + import PIL.Image from ...absolute_path import WEIGHTS_PATH # Mostly copied from `diffusers.StableDiffusionInpaintPipeline`, with slight modifications to yield the latents at each step. @@ -50,8 +52,8 @@ class GeneratorPipeline(diffusers.StableDiffusionInpaintPipeline): def __call__( self, prompt: Union[str, List[str]], - image: Union[torch.FloatTensor, Image.Image], - mask_image: Union[torch.FloatTensor, Image.Image], + image: Union[torch.FloatTensor, PIL.Image.Image], + mask_image: Union[torch.FloatTensor, PIL.Image.Image], height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, @@ -88,12 +90,12 @@ def __call__( ) # 4. Preprocess mask and image - if isinstance(image, Image.Image) and isinstance(mask_image, Image.Image): + if isinstance(image, PIL.Image.Image) and isinstance(mask_image, PIL.Image.Image): mask, masked_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.prepare_mask_and_masked_image(image, mask_image) # 5. set timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps_tensor = self.scheduler.timesteps + timesteps = self.scheduler.timesteps # 6. Prepare latent variables num_channels_latents = self.vae.config.latent_channels @@ -126,58 +128,76 @@ def __call__( num_channels_masked_image = masked_image_latents.shape[1] if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: raise ValueError( - f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" - f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" - f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of" - " `pipeline.unet` or your `mask_image` or `image` input." + f"Select an inpainting model, such as 'stabilityai/stable-diffusion-2-inpainting'" ) # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 10. Denoising loop - for i, t in enumerate(self.progress_bar(timesteps_tensor)): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) - - # predict the noise residual - noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # NOTE: Modified to yield the latents instead of calling a callback. - match kwargs['step_preview_mode']: - case StepPreviewMode.NONE: - pass - case StepPreviewMode.FAST: - yield np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255. - case StepPreviewMode.ACCURATE: - yield from [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. - for image in self.numpy_to_pil(self.decode_latents(latents)) - ] + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + # concat latents, mask, masked_image_latents in the channel dimension + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) + + # predict the noise residual + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # NOTE: Modified to yield the latents instead of calling a callback. + match kwargs['step_preview_mode']: + case StepPreviewMode.NONE: + yield ImageGenerationResult( + None, + generator.initial_seed(), + i, + False + ) + case StepPreviewMode.FAST: + yield ImageGenerationResult( + np.asarray(ImageOps.flip(Image.fromarray(approximate_decoded_latents(latents))).resize((width, height), Image.Resampling.NEAREST).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + i, + False + ) + case StepPreviewMode.ACCURATE: + yield from [ + ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + i, + False + ) + for image in self.numpy_to_pil(self.decode_latents(latents)) + ] # 11. Post-processing image = self.decode_latents(latents) - # TODO: Add UI to enable this. - # 12. Run safety checker + # TODO: Add UI to enable this + # 10. Run safety checker # image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype) # NOTE: Modified to yield the decoded image as a numpy array. yield from [ - np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + generator.initial_seed(), + num_inference_steps, + True + ) for image in self.numpy_to_pil(image) ] @@ -218,7 +238,10 @@ def __call__( pipe = optimizations.apply(pipe, device) # RNG - generator = None if seed is None else (torch.manual_seed(seed) if device == "mps" else torch.Generator(device=device).manual_seed(seed)) + generator = torch.Generator(device="cpu" if device == "mps" else device) # MPS does not support the `Generator` API + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + generator = generator.manual_seed(seed) # Seamless _configure_model_padding(pipe.unet, seamless, seamless_axes) @@ -238,7 +261,7 @@ def __call__( num_inference_steps=steps, guidance_scale=cfg_scale, negative_prompt=negative_prompt if use_negative_prompt else None, - num_images_per_prompt=iterations, + num_images_per_prompt=1, eta=0.0, generator=generator, latents=None, diff --git a/generator_process/actor.py b/generator_process/actor.py index 59e451e7..403faf69 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -29,6 +29,7 @@ class Future: _exception: BaseException | None = None done: bool = False cancelled: bool = False + call_done_on_exception: bool = True def __init__(self): self._response_callbacks = set() @@ -38,6 +39,7 @@ def __init__(self): self._exception = None self.done = False self.cancelled = False + self.call_done_on_exception = True def result(self): """ @@ -92,6 +94,8 @@ def set_exception(self, exception: BaseException): Set the exception. """ self._exception = exception + for exception_callback in self._exception_callbacks: + exception_callback(self, exception) def set_done(self): """ @@ -99,8 +103,9 @@ def set_done(self): """ assert not self.done self.done = True - for done_callback in self._done_callbacks: - done_callback(self) + if self._exception is None or self.call_done_on_exception: + for done_callback in self._done_callbacks: + done_callback(self) def add_response_callback(self, callback: Callable[['Future', Any], None]): """ diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 78632d42..21bb14c0 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -126,7 +126,8 @@ def step_callback(_, step_image: ImageGenerationResult): def done_callback(future): nonlocal last_data_block nonlocal iteration - del gen._active_generation_future + if hasattr(gen, '_active_generation_future'): + del gen._active_generation_future image_result: ImageGenerationResult | list = future.result() if isinstance(image_result, list): image_result = image_result[-1] @@ -158,6 +159,13 @@ def done_callback(future): scene.dream_textures_info = "" scene.dream_textures_progress = 0 + def exception_callback(_, exception): + scene.dream_textures_info = "" + scene.dream_textures_progress = 0 + if hasattr(gen, '_active_generation_future'): + del gen._active_generation_future + self.report({'ERROR'}, str(exception)) + gen = Generator.shared() def generate_next(): if init_image is not None: @@ -182,7 +190,9 @@ def generate_next(): **generated_args, ) gen._active_generation_future = f + f.call_done_on_exception = False f.add_response_callback(step_callback) + f.add_exception_callback(exception_callback) f.add_done_callback(done_callback) generate_next() return {"FINISHED"} From 5b001519b71812493547e230273707bd8bec6636 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 19:00:00 -0500 Subject: [PATCH 26/36] Raise the error in the exception_callback --- operators/dream_texture.py | 1 + 1 file changed, 1 insertion(+) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 21bb14c0..bee615eb 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -165,6 +165,7 @@ def exception_callback(_, exception): if hasattr(gen, '_active_generation_future'): del gen._active_generation_future self.report({'ERROR'}, str(exception)) + raise exception gen = Generator.shared() def generate_next(): From ca934eeb7015e6e6df5e21d9a3478c0127a8b06a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 20:22:46 -0500 Subject: [PATCH 27/36] Support Stability SDK --- generator_process/actions/prompt_to_image.py | 52 +++++++++++++++++++- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 8739f752..09742394 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -43,6 +43,22 @@ def scheduler_class(): return scheduler_class().from_pretrained(pretrained['model_path'], subfolder=pretrained['subfolder']) else: return scheduler_class().from_config(pipeline.scheduler.config) + + def stability_sdk(self): + import stability_sdk.interfaces.gooseai.generation.generation_pb2 + match self: + case Scheduler.LMS_DISCRETE: + return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_K_LMS + case Scheduler.DDIM: + return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_DDIM + case Scheduler.DDPM: + return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_DDPM + case Scheduler.EULER_DISCRETE: + return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_K_EULER + case Scheduler.EULER_ANCESTRAL_DISCRETE: + return stability_sdk.interfaces.gooseai.generation.generation_pb2.SAMPLER_K_EULER_ANCESTRAL + case _: + raise ValueError(f"{self} cannot be used with DreamStudio.") @dataclass(eq=True) class Optimizations: @@ -177,6 +193,9 @@ def prompt_to_image( step_preview_mode: StepPreviewMode, + # Stability SDK + key: str | None = None, + **kwargs ) -> Generator[ImageGenerationResult, None, None]: match pipeline: @@ -376,8 +395,37 @@ def __call__( step_preview_mode=step_preview_mode ) case Pipeline.STABILITY_SDK: - import stability_sdk - raise NotImplementedError() + import stability_sdk.client + import stability_sdk.interfaces.gooseai.generation.generation_pb2 + from PIL import Image, ImageOps + import io + + if key is None: + raise ValueError("DreamStudio key not provided. Enter your key in the add-on preferences.") + client = stability_sdk.client.StabilityInference(key=key, engine=model) + + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + + answers = client.generate( + prompt=prompt, + width=width, + height=height, + cfg_scale=cfg_scale, + sampler=scheduler.stability_sdk(), + steps=steps, + seed=seed + ) + for answer in answers: + for artifact in answer.artifacts: + if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE: + image = Image.open(io.BytesIO(artifact.binary)) + yield ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + seed, + steps, + True + ) case _: raise Exception(f"Unsupported pipeline {pipeline}.") From fac7d1308d5862d7a2ab524af9d113cc1886fdd7 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 20:23:02 -0500 Subject: [PATCH 28/36] Remove pinned version from requirements --- requirements/dreamstudio.txt | 2 +- requirements/linux-rocm.txt | 2 +- requirements/mac-mps-cpu.txt | 2 +- requirements/win-linux-cuda.txt | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/requirements/dreamstudio.txt b/requirements/dreamstudio.txt index 6f8f619d..cf7d8a33 100644 --- a/requirements/dreamstudio.txt +++ b/requirements/dreamstudio.txt @@ -1,2 +1,2 @@ -stability-sdk==0.2.6 +stability-sdk opencolorio \ No newline at end of file diff --git a/requirements/linux-rocm.txt b/requirements/linux-rocm.txt index 65a7f068..6efffdb1 100644 --- a/requirements/linux-rocm.txt +++ b/requirements/linux-rocm.txt @@ -8,7 +8,7 @@ torch>=1.13 scipy # LMSDiscreteScheduler -stability-sdk==0.2.6 # DreamStudio +stability-sdk # DreamStudio opencolorio==2.1.2 # color management diff --git a/requirements/mac-mps-cpu.txt b/requirements/mac-mps-cpu.txt index 9ef0f72a..3e4615d2 100644 --- a/requirements/mac-mps-cpu.txt +++ b/requirements/mac-mps-cpu.txt @@ -7,7 +7,7 @@ torch>=1.13 scipy # LMSDiscreteScheduler -stability-sdk==0.2.6 # DreamStudio +stability-sdk # DreamStudio opencolorio==2.1.2 # color management diff --git a/requirements/win-linux-cuda.txt b/requirements/win-linux-cuda.txt index 540458a2..8e146e20 100644 --- a/requirements/win-linux-cuda.txt +++ b/requirements/win-linux-cuda.txt @@ -8,7 +8,7 @@ torch>=1.13 scipy # LMSDiscreteScheduler -stability-sdk==0.2.6 # DreamStudio +stability-sdk # DreamStudio opencolorio==2.1.2 # color management From b5bb9470332e881d50a79a1e9ad9a6d0d4eb3254 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 20:38:01 -0500 Subject: [PATCH 29/36] Pipeline switching --- generator_process/actions/prompt_to_image.py | 53 ++++++++++++ generator_process/registrar.py | 91 -------------------- operators/dream_texture.py | 7 +- preferences.py | 6 +- property_groups/dream_prompt.py | 36 +++++--- render_pass.py | 1 - ui/panels/dream_texture.py | 19 ++-- ui/panels/render_properties.py | 10 +-- ui/panels/upscaling.py | 6 +- 9 files changed, 97 insertions(+), 132 deletions(-) delete mode 100644 generator_process/registrar.py diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 8739f752..8f356e84 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -12,6 +12,59 @@ class Pipeline(enum.IntEnum): STABILITY_SDK = 1 + @staticmethod + def local_available(): + from ...absolute_path import absolute_path + return os.path.exists(absolute_path(".python_dependencies/diffusers")) + + def __str__(self): + return self.name + + def model(self): + return True + + def init_img_actions(self): + match self: + case Pipeline.STABLE_DIFFUSION: + return ['modify', 'inpaint', 'outpaint'] + case Pipeline.STABILITY_SDK: + return ['modify', 'inpaint'] + + def inpaint_mask_sources(self): + match self: + case Pipeline.STABLE_DIFFUSION: + return ['alpha', 'prompt'] + case Pipeline.STABILITY_SDK: + return ['alpha'] + + def color_correction(self): + match self: + case Pipeline.STABLE_DIFFUSION: + return True + case Pipeline.STABILITY_SDK: + return False + + def negative_prompts(self): + match self: + case Pipeline.STABLE_DIFFUSION: + return True + case Pipeline.STABILITY_SDK: + return False + + def seamless(self): + match self: + case Pipeline.STABLE_DIFFUSION: + return True + case Pipeline.STABILITY_SDK: + return False + + def upscaling(self): + match self: + case Pipeline.STABLE_DIFFUSION: + return True + case Pipeline.STABILITY_SDK: + return False + class Scheduler(enum.Enum): LMS_DISCRETE = "LMS Discrete" DDIM = "DDIM" diff --git a/generator_process/registrar.py b/generator_process/registrar.py deleted file mode 100644 index 6eea82cd..00000000 --- a/generator_process/registrar.py +++ /dev/null @@ -1,91 +0,0 @@ -import os -from ..absolute_path import absolute_path -from .intent import Intent -from enum import IntEnum - -class BackendTarget(IntEnum): - """Which generator backend to use""" - LOCAL = 0 - STABILITY_SDK = 1 - - @staticmethod - def local_available(): - return os.path.exists(absolute_path(".python_dependencies/diffusers")) - - def __str__(self): - return self.name - - def init_img_actions(self): - match self: - case BackendTarget.LOCAL: - return ['modify', 'inpaint', 'outpaint'] - case BackendTarget.STABILITY_SDK: - return ['modify', 'inpaint'] - - def inpaint_mask_sources(self): - match self: - case BackendTarget.LOCAL: - return ['alpha', 'prompt'] - case BackendTarget.STABILITY_SDK: - return ['alpha'] - - def color_correction(self): - match self: - case BackendTarget.LOCAL: - return True - case BackendTarget.STABILITY_SDK: - return False - - def negative_prompts(self): - match self: - case BackendTarget.LOCAL: - return True - case BackendTarget.STABILITY_SDK: - return False - - def seamless(self): - match self: - case BackendTarget.LOCAL: - return True - case BackendTarget.STABILITY_SDK: - return False - - def upscaling(self): - match self: - case BackendTarget.LOCAL: - return True - case BackendTarget.STABILITY_SDK: - return False - -class _GeneratorIntent: - def __init__(self, func): - self.name = func.__name__ - self.func = func - -class _IntentBackend: - def __init__(self, intent: Intent, backend: BackendTarget | None, func): - self.intent = intent - self.backend = backend - self.func = func - -class _IntentRegistrar: - def __init__(self): - self._generator_intents: list[_GeneratorIntent] = [] - self._intent_backends: list[_IntentBackend] = [] - - def generator_intent(self, func): - ''' - Registers an intent as a function on the `GeneratorProcess` class. - ''' - intent = _GeneratorIntent(func) - self._generator_intents.append(intent) - return intent - - def intent_backend(self, intent: Intent, backend_target: BackendTarget | None = None): - def decorator(func): - backend = _IntentBackend(intent, backend_target, func) - self._intent_backends.append(backend) - return backend - return decorator - -registrar = _IntentRegistrar() \ No newline at end of file diff --git a/operators/dream_texture.py b/operators/dream_texture.py index bee615eb..0cf79471 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -6,9 +6,7 @@ from numpy.typing import NDArray from multiprocessing.shared_memory import SharedMemory -from ..property_groups.dream_prompt import backend_options - -from ..generator_process.registrar import BackendTarget +from ..property_groups.dream_prompt import pipeline_options from ..preferences import StableDiffusionPreferences from ..pil_to_image import * @@ -173,13 +171,11 @@ def generate_next(): match generated_args['init_img_action']: case 'modify': f = gen.image_to_image( - Pipeline.STABLE_DIFFUSION, image=init_image, **generated_args ) case 'inpaint': f = gen.inpaint( - Pipeline.STABLE_DIFFUSION, image=init_image, **generated_args ) @@ -187,7 +183,6 @@ def generate_next(): raise NotImplementedError() else: f = gen.prompt_to_image( - Pipeline.STABLE_DIFFUSION, **generated_args, ) gen._active_generation_future = f diff --git a/preferences.py b/preferences.py index fb51588f..ea99fbbb 100644 --- a/preferences.py +++ b/preferences.py @@ -11,7 +11,7 @@ from .property_groups.dream_prompt import DreamPrompt from .ui.presets import RestoreDefaultPresets, default_presets_missing from .generator_process import Generator -from .generator_process.registrar import BackendTarget +from .generator_process.actions.prompt_to_image import Pipeline class OpenHuggingFace(bpy.types.Operator): bl_idname = "dream_textures.open_hugging_face" @@ -155,7 +155,7 @@ class StableDiffusionPreferences(bpy.types.AddonPreferences): @staticmethod def register(): - if BackendTarget.local_available(): + if Pipeline.local_available(): set_model_list('installed_models', Generator.shared().hf_list_installed_models().result()) def draw(self, context): @@ -168,7 +168,7 @@ def draw(self, context): has_dependencies = len(os.listdir(absolute_path(".python_dependencies"))) > 2 if has_dependencies: - has_local = BackendTarget.local_available() + has_local = Pipeline.local_available() if has_local: search_box = layout.box() diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 2cac98b0..ebf12ab7 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -4,8 +4,7 @@ import sys from typing import _AnnotatedAlias from ..absolute_path import absolute_path -from ..generator_process.registrar import BackendTarget -from ..generator_process.actions.prompt_to_image import Optimizations, Scheduler, StepPreviewMode +from ..generator_process.actions.prompt_to_image import Optimizations, Scheduler, StepPreviewMode, Pipeline from ..generator_process import Generator from ..prompt_engineering import * @@ -32,7 +31,7 @@ ] def init_image_actions_filtered(self, context): - available = BackendTarget[self.backend].init_img_actions() + available = Pipeline[self.pipeline].init_img_actions() return list(filter(lambda x: x[0] in available, init_image_actions)) inpaint_mask_sources = [ @@ -41,7 +40,7 @@ def init_image_actions_filtered(self, context): ] def inpaint_mask_sources_filtered(self, context): - available = BackendTarget[self.backend].inpaint_mask_sources() + available = Pipeline[self.pipeline].inpaint_mask_sources() return list(filter(lambda x: x[0] in available, inpaint_mask_sources)) seamless_axes = [ @@ -54,17 +53,28 @@ def inpaint_mask_sources_filtered(self, context): def _on_model_options(future): global _model_options _model_options = future.result() -if BackendTarget.local_available(): +if Pipeline.local_available(): Generator.shared().hf_list_installed_models().add_done_callback(_on_model_options) def model_options(self, context): - return [(m.id, os.path.basename(m.id).replace('models--', '').replace('--', '/'), '', i) for i, m in enumerate(_model_options)] - -def backend_options(self, context): + match Pipeline[self.pipeline]: + case Pipeline.STABLE_DIFFUSION: + return [(m.id, os.path.basename(m.id).replace('models--', '').replace('--', '/'), '', i) for i, m in enumerate(_model_options)] + case Pipeline.STABILITY_SDK: + return [(x, x, '') for x in [ + "stable-diffusion-v1", + "stable-diffusion-v1-5", + "stable-diffusion-512-v2-0", + "stable-diffusion-768-v2-0", + "stable-inpainting-v1-0", + "stable-inpainting-512-v2-0" + ]] + +def pipeline_options(self, context): def options(): - if BackendTarget.local_available(): - yield (BackendTarget.LOCAL.name, 'Local', 'Run on your own hardware', 1) + if Pipeline.local_available(): + yield (Pipeline.STABLE_DIFFUSION.name, 'Stable Diffusion', 'Stable Diffusion on your own hardware', 1) if len(context.preferences.addons[__package__.split('.')[0]].preferences.dream_studio_key) > 0: - yield (BackendTarget.STABILITY_SDK.name, 'DreamStudio', 'Run in the cloud with DreamStudio', 2) + yield (Pipeline.STABILITY_SDK.name, 'DreamStudio', 'Cloud compute via DreamStudio', 2) return [*options()] def seed_clamp(self, ctx): @@ -77,7 +87,7 @@ def seed_clamp(self, ctx): pass # will get hashed once generated attributes = { - "backend": EnumProperty(name="Backend", items=backend_options, default=1 if BackendTarget.local_available() else 2, description="Fill in a few simple options to create interesting images quickly"), + "pipeline": EnumProperty(name="Pipeline", items=pipeline_options, default=1 if Pipeline.local_available() else 2, description="Specify which model and target should be used."), "model": EnumProperty(name="Model", items=model_options, description="Specify which model to use for inference"), # Prompt @@ -219,6 +229,8 @@ def generate_args(self): args['optimizations'] = self.get_optimizations() args['scheduler'] = Scheduler(args['scheduler']) args['step_preview_mode'] = StepPreviewMode(args['step_preview_mode']) + args['pipeline'] = Pipeline[args['pipeline']] + args['key'] = bpy.context.preferences.addons[__package__.split('.')[0]].preferences.dream_studio_key return args DreamPrompt.generate_prompt = generate_prompt diff --git a/render_pass.py b/render_pass.py index 0866fa2c..5f14556f 100644 --- a/render_pass.py +++ b/render_pass.py @@ -77,7 +77,6 @@ def render(self, depsgraph): generated_args['width'] = size_x generated_args['height'] = size_y pixels = gen.image_to_image( - pipeline=Pipeline.STABLE_DIFFUSION, image=(combined_pixels.reshape((size_x, size_y, 4)) * 255).astype(np.uint8), **generated_args, _block=True diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 842a63fb..ff68e1e9 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -5,8 +5,6 @@ import os import shutil -from ...generator_process.registrar import BackendTarget - from ...absolute_path import CLIPSEG_WEIGHTS_PATH from ..presets import DREAM_PT_AdvancedPresets from ...pil_to_image import * @@ -15,9 +13,8 @@ from ...operators.open_latest_version import OpenLatestVersion, is_force_show_download, new_version_available from ...operators.view_history import ImportPromptFile from ..space_types import SPACE_TYPES -from ...property_groups.dream_prompt import DreamPrompt, backend_options -from ...generator_process.registrar import BackendTarget -from ...generator_process.actions.prompt_to_image import Optimizations +from ...property_groups.dream_prompt import DreamPrompt, pipeline_options +from ...generator_process.actions.prompt_to_image import Optimizations, Pipeline def dream_texture_panels(): for space_type in SPACE_TYPES: @@ -46,9 +43,9 @@ def draw(self, context): layout.use_property_split = True layout.use_property_decorate = False - if len(backend_options(self, context)) > 1: - layout.prop(context.scene.dream_textures_prompt, "backend") - if context.scene.dream_textures_prompt.backend == BackendTarget.LOCAL.name: + if len(pipeline_options(self, context)) > 1: + layout.prop(context.scene.dream_textures_prompt, "pipeline") + if Pipeline[context.scene.dream_textures_prompt.pipeline].model(): layout.prop(context.scene.dream_textures_prompt, 'model') if is_force_show_download(): @@ -111,7 +108,7 @@ def draw(self, context): segment_row.prop(get_prompt(context), enum_prop, icon_only=is_custom) if get_prompt(context).prompt_structure == file_batch_structure.id: layout.template_ID(context.scene, "dream_textures_prompt_file", open="text.open") - if BackendTarget[get_prompt(context).backend].seamless(): + if Pipeline[get_prompt(context).pipeline].seamless(): layout.prop(get_prompt(context), "seamless") if get_prompt(context).seamless: layout.prop(get_prompt(context), "seamless_axes") @@ -125,7 +122,7 @@ class NegativePromptPanel(sub_panel): @classmethod def poll(self, context): - return get_prompt(context).prompt_structure != file_batch_structure.id and BackendTarget[get_prompt(context).backend].negative_prompts() + return get_prompt(context).prompt_structure != file_batch_structure.id and Pipeline[get_prompt(context).pipeline].negative_prompts() def draw_header(self, context): layout = self.layout @@ -237,7 +234,7 @@ def draw(self, context): elif prompt.init_img_action == 'modify': layout.prop(prompt, "fit") layout.prop(prompt, "strength") - if BackendTarget[prompt.backend].color_correction(): + if Pipeline[prompt.pipeline].color_correction(): layout.prop(prompt, "use_init_img_color") yield InitImagePanel diff --git a/ui/panels/render_properties.py b/ui/panels/render_properties.py index b18201c8..4127e6de 100644 --- a/ui/panels/render_properties.py +++ b/ui/panels/render_properties.py @@ -1,7 +1,7 @@ import bpy from .dream_texture import create_panel, prompt_panel, advanced_panel -from ...property_groups.dream_prompt import backend_options -from ...generator_process.registrar import BackendTarget +from ...property_groups.dream_prompt import pipeline_options +from ...generator_process.actions.prompt_to_image import Pipeline class RenderPropertiesPanel(bpy.types.Panel): """Panel for Dream Textures render properties""" @@ -25,9 +25,9 @@ def draw(self, context): layout.use_property_decorate = False layout.active = context.scene.dream_textures_render_properties_enabled - if len(backend_options(self, context)) > 1: - layout.prop(context.scene.dream_textures_render_properties_prompt, "backend") - if context.scene.dream_textures_render_properties_prompt.backend == BackendTarget.LOCAL.name: + if len(pipeline_options(self, context)) > 1: + layout.prop(context.scene.dream_textures_render_properties_prompt, "pipeline") + if Pipeline[context.scene.dream_textures_render_properties_prompt.pipeline].model(): layout.prop(context.scene.dream_textures_render_properties_prompt, 'model') layout.prop(context.scene.dream_textures_render_properties_prompt, "strength") diff --git a/ui/panels/upscaling.py b/ui/panels/upscaling.py index 4300c301..667260cd 100644 --- a/ui/panels/upscaling.py +++ b/ui/panels/upscaling.py @@ -1,9 +1,9 @@ from bpy.types import Panel -from ...generator_process.registrar import BackendTarget from ...pil_to_image import * from ...prompt_engineering import * from ...operators.upscale import Upscale from ...operators.dream_texture import CancelGenerator, ReleaseGenerator +from ...generator_process.actions.prompt_to_image import Pipeline from .dream_texture import create_panel, advanced_panel from ..space_types import SPACE_TYPES @@ -20,7 +20,7 @@ class UpscalingPanel(Panel): @classmethod def poll(cls, context): - if not BackendTarget[context.scene.dream_textures_prompt.backend].upscaling(): + if not Pipeline[context.scene.dream_textures_prompt.pipeline].upscaling(): return False if cls.bl_space_type == 'NODE_EDITOR': return context.area.ui_type == "ShaderNodeTree" or context.area.ui_type == "CompositorNodeTree" @@ -56,7 +56,7 @@ class ActionsPanel(Panel): @classmethod def poll(cls, context): - if not BackendTarget[context.scene.dream_textures_prompt.backend].upscaling(): + if not Pipeline[context.scene.dream_textures_prompt.pipeline].upscaling(): return False if cls.bl_space_type == 'NODE_EDITOR': return context.area.ui_type == "ShaderNodeTree" or context.area.ui_type == "CompositorNodeTree" From 4b6a84dfd07995fb17fdf17809ab175c9f6a2c0f Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Mon, 5 Dec 2022 20:40:34 -0500 Subject: [PATCH 30/36] fix upscale --- generator_process/actions/upscale.py | 17 ++++++++--------- operators/upscale.py | 1 - 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index 5d759ab2..e484b3e1 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -54,8 +54,6 @@ def upscale( generator = torch.Generator(device="cpu" if device == "mps" else device) # MPS does not support the `Generator` API if seed is None: seed = random.randrange(0, np.iinfo(np.uint32).max) - generator = generator.manual_seed(seed) - initial_seed = generator.initial_seed() tiler = Tiler( data_shape=(low_res_image.size[0], low_res_image.size[1], len(low_res_image.getbands())), @@ -71,13 +69,14 @@ def upscale( )) input_array = np.array(low_res_image) for id, tile in tiler(input_array, progress_bar=True): - merger.add(id, np.array(pipe( - prompt=prompt, - image=Image.fromarray(tile), - num_inference_steps=steps, - generator=torch.manual_seed(initial_seed), - guidance_scale=cfg_scale, - ).images[0])) + with torch.no_grad(): + merger.add(id, np.array(pipe( + prompt=prompt, + image=Image.fromarray(tile), + num_inference_steps=steps, + generator=generator.manual_seed(seed), + guidance_scale=cfg_scale, + ).images[0])) if step_preview_mode != StepPreviewMode.NONE: step = Image.fromarray(merger.merge().astype(np.uint8)) yield ImageUpscaleResult( diff --git a/operators/upscale.py b/operators/upscale.py index ea418e03..5216cf82 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -105,7 +105,6 @@ def on_tile_complete(_, tile: ImageUpscaleResult): for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = last_data_block - bpy.app.timers.register(update_progress) def image_done(future): nonlocal last_data_block From dd196359a6d7275e48875756ac56fcc58334a345 Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Mon, 5 Dec 2022 20:50:52 -0500 Subject: [PATCH 31/36] no_grad already used in pipe --- generator_process/actions/upscale.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/generator_process/actions/upscale.py b/generator_process/actions/upscale.py index e484b3e1..de1201b4 100644 --- a/generator_process/actions/upscale.py +++ b/generator_process/actions/upscale.py @@ -69,14 +69,13 @@ def upscale( )) input_array = np.array(low_res_image) for id, tile in tiler(input_array, progress_bar=True): - with torch.no_grad(): - merger.add(id, np.array(pipe( - prompt=prompt, - image=Image.fromarray(tile), - num_inference_steps=steps, - generator=generator.manual_seed(seed), - guidance_scale=cfg_scale, - ).images[0])) + merger.add(id, np.array(pipe( + prompt=prompt, + image=Image.fromarray(tile), + num_inference_steps=steps, + generator=generator.manual_seed(seed), + guidance_scale=cfg_scale, + ).images[0])) if step_preview_mode != StepPreviewMode.NONE: step = Image.fromarray(merger.merge().astype(np.uint8)) yield ImageUpscaleResult( From 7f93e4e4882040e451460930116d8a49ae34a1b7 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 5 Dec 2022 21:49:22 -0500 Subject: [PATCH 32/36] Re-use the same datablock --- operators/dream_texture.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 0cf79471..b4409606 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -41,10 +41,15 @@ def save_temp_image(img, path=None): return path -def bpy_image(name, width, height, pixels): - image = bpy.data.images.new(name, width=width, height=height) +def bpy_image(name, width, height, pixels, existing_image): + if existing_image is None: + image = bpy.data.images.new(name, width=width, height=height) + else: + image = existing_image + image.name = name image.pixels[:] = pixels image.pack() + image.update() return image class DreamTexture(bpy.types.Operator): @@ -108,14 +113,11 @@ def step_progress_update(self, context): last_data_block = None def step_callback(_, step_image: ImageGenerationResult): nonlocal last_data_block - if last_data_block is not None: - bpy.data.images.remove(last_data_block) - last_data_block = None if step_image.final: return scene.dream_textures_progress = step_image.step if step_image.image is not None: - last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel()) + last_data_block = bpy_image(f"Step {step_image.step}/{generated_args['steps']}", step_image.image.shape[1], step_image.image.shape[0], step_image.image.ravel(), last_data_block) for area in screen.areas: if area.type == 'IMAGE_EDITOR': area.spaces.active.image = last_data_block @@ -129,10 +131,7 @@ def done_callback(future): image_result: ImageGenerationResult | list = future.result() if isinstance(image_result, list): image_result = image_result[-1] - if last_data_block is not None: - bpy.data.images.remove(last_data_block) - last_data_block = None - image = bpy_image(str(image_result.seed), image_result.image.shape[1], image_result.image.shape[0], image_result.image.ravel()) + image = bpy_image(str(image_result.seed), image_result.image.shape[1], image_result.image.shape[0], image_result.image.ravel(), last_data_block) if node_tree is not None: nodes = node_tree.nodes texture_node = nodes.new("ShaderNodeTexImage") From a8fca6c788d725723e2652fc507072cbc31f0075 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 12 Dec 2022 22:09:44 -0500 Subject: [PATCH 33/36] Add image_to_image and inpaint --- generator_process/actions/image_to_image.py | 40 +++++++++++++++++- generator_process/actions/inpaint.py | 44 +++++++++++++++++++- generator_process/actions/prompt_to_image.py | 2 + 3 files changed, 82 insertions(+), 4 deletions(-) diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index ceb25d9f..e5672929 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -36,6 +36,9 @@ def image_to_image( step_preview_mode: StepPreviewMode, + # Stability SDK + key: str | None = None, + **kwargs ) -> Generator[NDArray, None, None]: match pipeline: @@ -232,7 +235,40 @@ def __call__( step_preview_mode=step_preview_mode ) case Pipeline.STABILITY_SDK: - import stability_sdk - raise NotImplementedError() + import stability_sdk.client + import stability_sdk.interfaces.gooseai.generation.generation_pb2 + from PIL import Image, ImageOps + import io + + if key is None: + raise ValueError("DreamStudio key not provided. Enter your key in the add-on preferences.") + client = stability_sdk.client.StabilityInference(key=key, engine=model) + + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + + answers = client.generate( + prompt=prompt, + width=width, + height=height, + cfg_scale=cfg_scale, + sampler=scheduler.stability_sdk(), + steps=steps, + seed=seed, + init_image=(Image.open(image) if isinstance(image, str) else Image.fromarray(image)).convert('RGB'), + start_schedule=strength, + ) + for answer in answers: + for artifact in answer.artifacts: + if artifact.finish_reason == stability_sdk.interfaces.gooseai.generation.generation_pb2.FILTER: + raise ValueError("Your request activated DreamStudio's safety filter. Please modify your prompt and try again.") + if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE: + image = Image.open(io.BytesIO(artifact.binary)) + yield ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + seed, + steps, + True + ) case _: raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actions/inpaint.py b/generator_process/actions/inpaint.py index fd23c941..0981b7fc 100644 --- a/generator_process/actions/inpaint.py +++ b/generator_process/actions/inpaint.py @@ -36,6 +36,9 @@ def inpaint( step_preview_mode: StepPreviewMode, + # Stability SDK + key: str | None = None, + **kwargs ) -> Generator[NDArray, None, None]: match pipeline: @@ -272,7 +275,44 @@ def __call__( step_preview_mode=step_preview_mode ) case Pipeline.STABILITY_SDK: - import stability_sdk - raise NotImplementedError() + import stability_sdk.client + import stability_sdk.interfaces.gooseai.generation.generation_pb2 + from PIL import Image, ImageOps + import io + + if key is None: + raise ValueError("DreamStudio key not provided. Enter your key in the add-on preferences.") + client = stability_sdk.client.StabilityInference(key=key, engine=model) + + if seed is None: + seed = random.randrange(0, np.iinfo(np.uint32).max) + + init_image = Image.open(image) if isinstance(image, str) else Image.fromarray(image) + + answers = client.generate( + prompt=prompt, + width=width, + height=height, + cfg_scale=cfg_scale, + sampler=scheduler.stability_sdk(), + steps=steps, + seed=seed, + init_image=init_image.convert('RGB'), + mask_image=init_image.getchannel('A'), + start_schedule=strength, + ) + for answer in answers: + for artifact in answer.artifacts: + if artifact.finish_reason == stability_sdk.interfaces.gooseai.generation.generation_pb2.FILTER: + raise ValueError("Your request activated DreamStudio's safety filter. Please modify your prompt and try again.") + if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE: + image = Image.open(io.BytesIO(artifact.binary)) + yield ImageGenerationResult( + np.asarray(ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255., + seed, + steps, + True + ) + case _: raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index c1b94691..ea14dbeb 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -471,6 +471,8 @@ def __call__( ) for answer in answers: for artifact in answer.artifacts: + if artifact.finish_reason == stability_sdk.interfaces.gooseai.generation.generation_pb2.FILTER: + raise ValueError("Your request activated DreamStudio's safety filter. Please modify your prompt and try again.") if artifact.type == stability_sdk.interfaces.gooseai.generation.generation_pb2.ARTIFACT_IMAGE: image = Image.open(io.BytesIO(artifact.binary)) yield ImageGenerationResult( From 493b4844a64a5ac66d6306aef433734b550f96f9 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 14 Dec 2022 11:40:27 -0500 Subject: [PATCH 34/36] Update render pass --- generator_process/actions/image_to_image.py | 4 ++- render_pass.py | 27 ++++++--------------- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index e5672929..012506ba 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -34,7 +34,7 @@ def image_to_image( iterations: int, - step_preview_mode: StepPreviewMode, + step_preview_mode: StepPreviewMode | None, # Stability SDK key: str | None = None, @@ -125,6 +125,8 @@ def __call__( # NOTE: Modified to yield the latents instead of calling a callback. match kwargs['step_preview_mode']: + case None: + break case StepPreviewMode.NONE: yield ImageGenerationResult( None, diff --git a/render_pass.py b/render_pass.py index 5f14556f..f07e99b6 100644 --- a/render_pass.py +++ b/render_pass.py @@ -46,10 +46,7 @@ def render(self, depsgraph): if render_pass.name == "Dream Textures": self.update_stats("Dream Textures", "Starting") - # step_count = int(scene.dream_textures_render_properties_prompt.strength * scene.dream_textures_render_properties_prompt.steps) - self.update_stats("Dream Textures", "Creating temporary image") - combined_pass_image = bpy.data.images.new("dream_textures_post_processing_temp", width=size_x, height=size_y) rect = layer.passes["Combined"].rect @@ -73,19 +70,19 @@ def render(self, depsgraph): self.update_stats("Dream Textures", "Generating...") generated_args = scene.dream_textures_render_properties_prompt.generate_args() - generated_args['step_preview_mode'] = StepPreviewMode.NONE + generated_args['step_preview_mode'] = None generated_args['width'] = size_x generated_args['height'] = size_y - pixels = gen.image_to_image( - image=(combined_pixels.reshape((size_x, size_y, 4)) * 255).astype(np.uint8), + combined_pixels = gen.image_to_image( + image=np.flipud(combined_pixels.reshape((size_y, size_x, 4)) * 255).astype(np.uint8), **generated_args, _block=True - ).result() + ).result().image # Perform an inverse transform so when Blender applies its transform everything looks correct. self.update_stats("Dream Textures", "Applying inverse color management transforms") - pixels = gen.ocio_transform( - pixels, + combined_pixels = gen.ocio_transform( + combined_pixels.reshape((size_x * size_y, 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, @@ -96,17 +93,9 @@ def render(self, depsgraph): _block=True ).result() - reshaped = pixels.reshape((size_x * size_y, 4)) - render_pass.rect.foreach_set(reshaped) - - # delete pointers before closing shared memory - del pixels - del combined_pixels - del reshaped + combined_pixels = combined_pixels.reshape((size_x * size_y, 4)) + render_pass.rect.foreach_set(combined_pixels) - def cleanup(): - bpy.data.images.remove(combined_pass_image) - bpy.app.timers.register(cleanup) self.update_stats("Dream Textures", "Finished") else: pixels = np.empty((len(original_render_pass.rect), len(original_render_pass.rect[0])), dtype=np.float32) From 672489aaa34a2f7f141124061d8bd09b443d23ab Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 14 Dec 2022 12:18:35 -0500 Subject: [PATCH 35/36] Render pass improvements --- __init__.py | 2 +- generator_process/actions/image_to_image.py | 6 +- render_pass.py | 150 +++++++++++--------- 3 files changed, 87 insertions(+), 71 deletions(-) diff --git a/__init__.py b/__init__.py index 41969c34..7ae1b915 100644 --- a/__init__.py +++ b/__init__.py @@ -110,6 +110,6 @@ def unregister(): for tool in TOOLS: bpy.utils.unregister_tool(tool) - # unregister_render_pass() + unregister_render_pass() kill_generator() \ No newline at end of file diff --git a/generator_process/actions/image_to_image.py b/generator_process/actions/image_to_image.py index 012506ba..04224f18 100644 --- a/generator_process/actions/image_to_image.py +++ b/generator_process/actions/image_to_image.py @@ -34,13 +34,13 @@ def image_to_image( iterations: int, - step_preview_mode: StepPreviewMode | None, + step_preview_mode: StepPreviewMode, # Stability SDK key: str | None = None, **kwargs -) -> Generator[NDArray, None, None]: +) -> Generator[ImageGenerationResult, None, None]: match pipeline: case Pipeline.STABLE_DIFFUSION: import diffusers @@ -125,8 +125,6 @@ def __call__( # NOTE: Modified to yield the latents instead of calling a callback. match kwargs['step_preview_mode']: - case None: - break case StepPreviewMode.NONE: yield ImageGenerationResult( None, diff --git a/render_pass.py b/render_pass.py index f07e99b6..a48b2aab 100644 --- a/render_pass.py +++ b/render_pass.py @@ -2,8 +2,9 @@ import cycles import numpy as np import os -from .generator_process.actions.prompt_to_image import Pipeline, StepPreviewMode +from .generator_process.actions.prompt_to_image import Pipeline, StepPreviewMode, ImageGenerationResult from .generator_process import Generator +import threading update_render_passes_original = cycles.CyclesRender.update_render_passes render_original = cycles.CyclesRender.render @@ -33,73 +34,19 @@ def render(self, depsgraph): self.report({"ERROR"}, f"Image dimensions must be multiples of 64 (e.x. 512x512, 512x768, ...) closest is {round(size_x/64)*64}x{round(size_y/64)*64}") return result render_result = self.begin_result(0, 0, size_x, size_y) - for original_layer in original_result.layers: - layer = None - for layer_i in render_result.layers: - if layer_i.name == original_layer.name: - layer = layer_i - for original_render_pass in original_layer.passes: - render_pass = None - for pass_i in layer.passes: - if pass_i.name == original_render_pass.name: - render_pass = pass_i + for layer in render_result.layers: + for render_pass in layer.passes: if render_pass.name == "Dream Textures": - self.update_stats("Dream Textures", "Starting") - - self.update_stats("Dream Textures", "Creating temporary image") - - rect = layer.passes["Combined"].rect - - combined_pixels = np.empty((size_x * size_y, 4), dtype=np.float32) - rect.foreach_get(combined_pixels) - - gen = Generator.shared() - self.update_stats("Dream Textures", "Applying color management transforms") - combined_pixels = gen.ocio_transform( - combined_pixels, - config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'), - exposure=scene.view_settings.exposure, - gamma=scene.view_settings.gamma, - view_transform=scene.view_settings.view_transform, - display_device=scene.display_settings.display_device, - look=scene.view_settings.look, - inverse=False, - _block=True - ).result() - - self.update_stats("Dream Textures", "Generating...") - - generated_args = scene.dream_textures_render_properties_prompt.generate_args() - generated_args['step_preview_mode'] = None - generated_args['width'] = size_x - generated_args['height'] = size_y - combined_pixels = gen.image_to_image( - image=np.flipud(combined_pixels.reshape((size_y, size_x, 4)) * 255).astype(np.uint8), - **generated_args, - _block=True - ).result().image - - # 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_x * size_y, 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, - _block=True - ).result() - - combined_pixels = combined_pixels.reshape((size_x * size_y, 4)) - render_pass.rect.foreach_set(combined_pixels) - - self.update_stats("Dream Textures", "Finished") + self._render_dream_textures_pass(layer, (size_x, size_y), scene, render_pass, render_result) else: - pixels = np.empty((len(original_render_pass.rect), len(original_render_pass.rect[0])), dtype=np.float32) - original_render_pass.rect.foreach_get(pixels) + source_pass = None + for original_layer in original_result.layers: + if layer.name == original_layer.name: + for original_pass in original_layer.passes: + if original_pass.name == render_pass.name: + source_pass = original_pass + pixels = np.empty((len(source_pass.rect), len(source_pass.rect[0])), dtype=np.float32) + source_pass.rect.foreach_get(pixels) render_pass.rect[:] = pixels self.end_result(render_result) except Exception as e: @@ -107,6 +54,7 @@ def render(self, depsgraph): return result return render cycles.CyclesRender.render = render_decorator(cycles.CyclesRender.render) + cycles.CyclesRender._render_dream_textures_pass = _render_dream_textures_pass # def del_decorator(original): # def del_patch(self): @@ -121,5 +69,75 @@ def unregister_render_pass(): cycles.CyclesRender.update_render_passes = update_render_passes_original global render_original cycles.CyclesRender.render = render_original + del cycles.CyclesRender._render_dream_textures_pass # global del_original # cycles.CyclesRender.__del__ = del_original + +def _render_dream_textures_pass(self, layer, size, scene, render_pass, render_result): + self.update_stats("Dream Textures", "Starting") + self.update_stats("Dream Textures", "Creating temporary image") + + rect = layer.passes["Combined"].rect + + combined_pixels = np.empty((size[0] * size[1], 4), dtype=np.float32) + rect.foreach_get(combined_pixels) + + gen = Generator.shared() + self.update_stats("Dream Textures", "Applying color management transforms") + combined_pixels = gen.ocio_transform( + combined_pixels, + config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'), + exposure=scene.view_settings.exposure, + gamma=scene.view_settings.gamma, + view_transform=scene.view_settings.view_transform, + display_device=scene.display_settings.display_device, + look=scene.view_settings.look, + inverse=False + ).result() + + self.update_stats("Dream Textures", "Generating...") + + generated_args = scene.dream_textures_render_properties_prompt.generate_args() + generated_args['width'] = size[0] + generated_args['height'] = size[1] + f = gen.image_to_image( + image=np.flipud(combined_pixels.reshape((size[1], size[0], 4)) * 255).astype(np.uint8), + **generated_args + ) + event = threading.Event() + def on_step(_, step: ImageGenerationResult): + if step.final: + return + self.update_progress(step.step / generated_args['steps']) + if step.image is not None: + combined_pixels = step.image + render_pass.rect.foreach_set(combined_pixels.reshape((size[0] * size[1], 4))) + self.update_result(render_result) # This does not seem to have an effect. + def on_done(future): + nonlocal combined_pixels + result = future.result() + if isinstance(result, list): + result = result[-1] + combined_pixels = result.image + event.set() + f.add_response_callback(on_step) + f.add_done_callback(on_done) + event.wait() + + # Perform an inverse transform so when Blender applies its transform everything looks correct. + self.update_stats("Dream Textures", "Applying inverse color management transforms") + combined_pixels = gen.ocio_transform( + combined_pixels.reshape((size[0] * size[1], 4)), + config_path=os.path.join(bpy.utils.resource_path('LOCAL'), 'datafiles/colormanagement/config.ocio'), + exposure=scene.view_settings.exposure, + gamma=scene.view_settings.gamma, + view_transform=scene.view_settings.view_transform, + display_device=scene.display_settings.display_device, + look=scene.view_settings.look, + inverse=True + ).result() + + combined_pixels = combined_pixels.reshape((size[0] * size[1], 4)) + render_pass.rect.foreach_set(combined_pixels) + + self.update_stats("Dream Textures", "Finished") \ No newline at end of file From aefbe4fcf0082ebf73cc4dec30ceef72e67a90bb Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 14 Dec 2022 12:19:32 -0500 Subject: [PATCH 36/36] Remove incorrect stat --- render_pass.py | 1 - 1 file changed, 1 deletion(-) diff --git a/render_pass.py b/render_pass.py index a48b2aab..46b1e6c6 100644 --- a/render_pass.py +++ b/render_pass.py @@ -75,7 +75,6 @@ def unregister_render_pass(): def _render_dream_textures_pass(self, layer, size, scene, render_pass, render_result): self.update_stats("Dream Textures", "Starting") - self.update_stats("Dream Textures", "Creating temporary image") rect = layer.passes["Combined"].rect