From b33fc4fa0698a8cbfbedb202c2ac61ad0729b921 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 2 Oct 2022 21:45:55 -0400 Subject: [PATCH 1/6] Add Real-ESRGAN integration --- README.md | 2 + __init__.py | 4 + absolute_path.py | 3 +- classes.py | 7 +- generator_process.py | 209 ++++++++++++++++++++++-------- operators/dream_texture.py | 17 +-- operators/upscale.py | 100 ++++++++++++++ requirements-mac-realesrgan.txt | 10 ++ requirements-win-torch-1-11-0.txt | 1 + stable_diffusion | 2 +- ui/panel.py | 64 +++++++++ weights/realesrgan/.gitignore | 2 + 12 files changed, 353 insertions(+), 68 deletions(-) create mode 100644 operators/upscale.py create mode 100644 requirements-mac-realesrgan.txt create mode 100644 weights/realesrgan/.gitignore diff --git a/README.md b/README.md index d1af908b..8aa27931 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,8 @@ ## Installation Download the [latest release](https://github.com/carson-katri/dream-textures/releases/tag/0.0.6) and follow the instructions there to get up and running. +> On macOS, it is possible you will run into a quarantine issue with the dependencies. To work around this, run the following command in the app `Terminal`: `xattr -r -d com.apple.quarantine ~/Library/Application\ Support/Blender/3.3/scripts/addons/dream_textures/.python_dependencies`. This will allow the PyTorch `.dylib`s and `.so`s to load without having to manually allow each one in System Preferences. + ## Usage | Enter a prompt | Generate a unique texture in a few seconds | diff --git a/__init__.py b/__init__.py index 5c5e69ba..960667e2 100644 --- a/__init__.py +++ b/__init__.py @@ -35,6 +35,7 @@ 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 requirements_path_items = ( # Use the old version of requirements-win.txt to fix installation issues with Blender + PyTorch 1.12.1 @@ -67,6 +68,9 @@ def register(): bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="Progress", default=0, min=0, max=0) bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info") + bpy.types.Scene.dream_textures_upscale_target_size = bpy.props.EnumProperty(name="Target Size", items=upscale_options) + bpy.types.Scene.dream_textures_upscale_strength = bpy.props.FloatProperty(name="Strength", default=1, min=0, max=1) + for cls in CLASSES: bpy.utils.register_class(cls) diff --git a/absolute_path.py b/absolute_path.py index a6bdeff1..7ab08b58 100644 --- a/absolute_path.py +++ b/absolute_path.py @@ -8,4 +8,5 @@ def absolute_path(component): """ return os.path.join(os.path.dirname(os.path.realpath(__file__)), component) -WEIGHTS_PATH = absolute_path("stable_diffusion/models/ldm/stable-diffusion-v1/model.ckpt") \ No newline at end of file +WEIGHTS_PATH = absolute_path("stable_diffusion/models/ldm/stable-diffusion-v1/model.ckpt") +REAL_ESRGAN_WEIGHTS_PATH = absolute_path("weights/realesrgan/realesr-general-x4v3.pth") \ No newline at end of file diff --git a/classes.py b/classes.py index 7acd8541..fa731a42 100644 --- a/classes.py +++ b/classes.py @@ -3,8 +3,9 @@ from .operators.dream_texture import DreamTexture, ReleaseGenerator from .operators.view_history import RecallHistoryEntry, SCENE_UL_HistoryList from .operators.inpaint_area_brush import InpaintAreaStroke +from .operators.upscale import Upscale from .property_groups.dream_prompt import DreamPrompt -from .ui.panel import panels, history_panels, troubleshooting_panels +from .ui.panel import panels, history_panels, troubleshooting_panels, upscaling_panels, OpenRealESRGANDownload, OpenRealESRGANWeightsDirectory from .preferences import OpenGitDownloads, OpenHuggingFace, OpenWeightsDirectory, OpenRustInstaller, ValidateInstallation, StableDiffusionPreferences CLASSES = ( @@ -13,10 +14,14 @@ OpenLatestVersion, RecallHistoryEntry, InpaintAreaStroke, + Upscale, SCENE_UL_HistoryList, *panels(), + *upscaling_panels(), *history_panels(), *troubleshooting_panels(), + OpenRealESRGANDownload, + OpenRealESRGANWeightsDirectory ) PREFERENCE_CLASSES = ( diff --git a/generator_process.py b/generator_process.py index 35b865a7..f1912f90 100644 --- a/generator_process.py +++ b/generator_process.py @@ -7,11 +7,12 @@ import site import traceback import numpy as np -from enum import IntEnum as Lawsuit, auto +from enum import IntEnum +import tempfile MISSING_DEPENDENCIES_ERROR = "Python dependencies are missing. Click Download Latest Release to fix." -class Action(Lawsuit): # can't help myself +class Action(IntEnum): """IPC message types sent from backend to frontend""" UNKNOWN = -1 # placeholder so you can do Action(int).name or Action(int) == Action.UNKNOWN when int is invalid # don't add anymore negative actions @@ -21,6 +22,7 @@ class Action(Lawsuit): # can't help myself STEP_IMAGE = 3 STEP_NO_SHOW = 4 EXCEPTION = 5 + IMAGE_PATH = 6 # An image result that returns a path to a saved image instead of bytes @classmethod def _missing_(cls, value): @@ -28,6 +30,18 @@ def _missing_(cls, value): ACTION_BYTE_LENGTH = ceil(log(max(Action)+1,256)) # doubt there will ever be more than 255 actions, but just in case +class Intent(IntEnum): + """IPC messages types sent from frontend to backend""" + UNKNOWN = -1 + + PROMPT_TO_IMAGE = 0 + UPSCALE = 1 + + @classmethod + def _missing_(cls, value): + return cls.UNKNOWN + +_shared_instance = None class GeneratorProcess(): def __init__(self): self.process = subprocess.Popen([sys.executable,'generator_process.py'],cwd=os.path.dirname(os.path.realpath(__file__)),stdin=subprocess.PIPE,stdout=subprocess.PIPE) @@ -38,6 +52,21 @@ def __init__(self): self.thread = threading.Thread(target=self._run,daemon=True,name="BackgroundReader") self.thread.start() + @classmethod + def shared(self): + global _shared_instance + if _shared_instance is None: + _shared_instance = GeneratorProcess() + return _shared_instance + + @classmethod + def kill_shared(self): + global _shared_instance + if _shared_instance is None: + return + _shared_instance.kill() + _shared_instance = None + def kill(self): self.killed = True self.process.kill() @@ -45,6 +74,8 @@ def kill(self): def prompt2image(self, args, step_callback, image_callback, info_callback, exception_callback): self.args = args stdin = self.process.stdin + stdin.write(Intent.PROMPT_TO_IMAGE.to_bytes(ACTION_BYTE_LENGTH, sys.byteorder, signed=False)) + stdin.flush() b = bytes(json.dumps(args), encoding='utf-8') stdin.write(len(b).to_bytes(8,sys.byteorder,signed=False)) stdin.write(b) @@ -71,6 +102,33 @@ def prompt2image(self, args, step_callback, image_callback, info_callback, excep elif action == Action.EXCEPTION: return + def upscale(self, args, image_callback, info_callback, exception_callback): + stdin = self.process.stdin + stdin.write(Intent.UPSCALE.to_bytes(ACTION_BYTE_LENGTH, sys.byteorder, signed=False)) + # stdin.flush() + b = bytes(json.dumps(args), encoding='utf-8') + stdin.write(len(b).to_bytes(8,sys.byteorder,signed=False)) + stdin.write(b) + stdin.flush() + + queue = self.queue + callbacks = { + Action.INFO: info_callback, + Action.IMAGE_PATH: image_callback, + Action.EXCEPTION: exception_callback + } + + while True: + while len(queue) == 0: + yield + tup = queue.pop() + action = tup[0] + callbacks[action](**tup[1]) + if action == Action.IMAGE_PATH: + break + elif action == Action.EXCEPTION: + return + def _run(self): reader = self.reader def readUInt(length): @@ -91,7 +149,7 @@ def queue_exception_msg(msg): kwargs = {} if kwargs_len == 0 else json.loads(reader.read(kwargs_len)) payload_len = readUInt(8) - if action in [Action.INFO, Action.STEP_NO_SHOW]: + if action in [Action.INFO, Action.STEP_NO_SHOW, Action.IMAGE_PATH]: queue.append((action, kwargs)) elif action in [Action.IMAGE, Action.STEP_IMAGE]: expected_len = kwargs['width']*kwargs['height']*16 @@ -187,7 +245,7 @@ def send_exception(fatal = True, msg: str = None, trace: str = None): from stable_diffusion.ldm.generate import Generate from omegaconf import OmegaConf - from PIL import ImageOps + from PIL import Image, ImageOps from io import StringIO except ModuleNotFoundError as e: min_files = 10 # bump this up if more files get added to .python_dependencies in source @@ -212,17 +270,20 @@ def send_exception(fatal = True, msg: str = None, trace: str = None): def image_to_bytes(image): return (np.asarray(ImageOps.flip(image).convert('RGBA'),dtype=np.float32) * byte_to_normalized).tobytes() - def image_writer(image, seed, upscaled=False): + def image_writer(image, seed, upscaled=False, first_seed=None): # Only use the non-upscaled texture, as upscaling is currently unsupported by the addon. if not upscaled: send_action(Action.IMAGE, payload=image_to_bytes(image), seed=seed, width=image.width, height=image.height) - def view_step(samples, step): + step = 0 + def view_step(samples): + nonlocal step if args['show_steps']: image = generator._sample_to_image(samples) send_action(Action.STEP_IMAGE, payload=image_to_bytes(image), step=step, width=image.width, height=image.height) else: send_action(Action.STEP_NO_SHOW, step=step) + step += 1 def preload_models(): tqdm = None @@ -281,57 +342,103 @@ def update(self, n=1): generator = None while True: - json_len = int.from_bytes(stdin.read(8),sys.byteorder,signed=False) - if json_len == 0: - return # stdin closed - args = json.loads(stdin.read(json_len)) - - if generator is None or (generator.full_precision != args['full_precision'] and sys.platform != 'darwin'): - send_info("Loading Model") + intent = Intent.from_bytes(stdin.read(ACTION_BYTE_LENGTH), sys.byteorder, signed=False) + if intent == Intent.PROMPT_TO_IMAGE: + json_len = int.from_bytes(stdin.read(8),sys.byteorder,signed=False) + if json_len == 0: + return # stdin closed + args = json.loads(stdin.read(json_len)) + + if generator is None or (generator.full_precision != args['full_precision'] and sys.platform != 'darwin'): + send_info("Loading Model") + try: + generator = Generate( + conf=models_config, + model=model, + # These args are deprecated, but we need them to specify an absolute path to the weights. + weights=weights, + config=config, + full_precision=args['full_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() + except: + send_exception() + return + send_info("Starting") + try: - generator = Generate( - conf=models_config, + tmp_stderr = sys.stderr = StringIO() # prompt2image writes exceptions straight to stderr, intercepting + 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, + **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: + send_exception(False, f"{low_ram[1]}{' or disable full precision' if args['full_precision'] else ''}{low_ram[2]}", s) + elif s.find("CUDA out of memory. Tried to allocate") != -1: + send_exception(False, f"Not enough memory, use lower resolution{' or disable full precision' if args['full_precision'] else ''}", s) + else: + send_exception(True, msg=None, trace=s) # consider all unknown exceptions to be fatal so the generator process is fully restarted next time + return + except: + send_exception() + return + finally: + sys.stderr = stderr + elif intent == Intent.UPSCALE: + tmp_stderr = sys.stderr = StringIO() + json_len = int.from_bytes(stdin.read(8),sys.byteorder,signed=False) + if json_len == 0: + return # stdin closed + args = json.loads(stdin.read(json_len)) + send_info("Starting") + try: + from absolute_path import REAL_ESRGAN_WEIGHTS_PATH + import cv2 + from realesrgan import RealESRGANer + from realesrgan.archs.srvgg_arch import SRVGGNetCompact + # image = Image.open(args['input']) + image = cv2.imread(args['input'], cv2.IMREAD_UNCHANGED) + model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu') + netscale = 4 + send_info("Loading Upsampler") + upsampler = RealESRGANer( + scale=netscale, + model_path=REAL_ESRGAN_WEIGHTS_PATH, + dni_weight=1, model=model, - # These args are deprecated, but we need them to specify an absolute path to the weights. - weights=weights, - config=config, - full_precision=args['full_precision'] + tile=0, + tile_pad=10, + pre_pad=0, + half=False, + gpu_id=None ) - generator.load_model() + send_info("Enhancing") + output, _ = upsampler.enhance(image, outscale=4) + send_info("Upscaled") + output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) + output = Image.fromarray(output) + output_path = tempfile.NamedTemporaryFile().name + output.save(output_path, format="png") + send_action(Action.IMAGE_PATH, output_path=output_path) except: send_exception() return - send_info("Starting") - - try: - tmp_stderr = sys.stderr = StringIO() # prompt2image writes exceptions straight to stderr, intercepting - 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, - **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: - send_exception(False, f"{low_ram[1]}{' or disable full precision' if args['full_precision'] else ''}{low_ram[2]}", s) - elif s.find("CUDA out of memory. Tried to allocate") != -1: - send_exception(False, f"Not enough memory, use lower resolution{' or disable full precision' if args['full_precision'] else ''}", s) - else: - send_exception(True, msg=None, trace=s) # consider all unknown exceptions to be fatal so the generator process is fully restarted next time - return - except: - send_exception() - return - finally: - sys.stderr = stderr + finally: + sys.stderr = stderr + else: + send_exception(True, f"Unknown intent {intent} sent to process. Expected one of {Intent._member_names_}.", "") if __name__ == "__main__": main() \ No newline at end of file diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 94268dea..647901ec 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -13,10 +13,6 @@ import tempfile -# A shared `Generate` instance. -# This allows the slow model loading process to happen once, -# and re-use the model on subsequent calls. -generator = None generator_advance = None last_data_block = None timer = None @@ -87,12 +83,8 @@ def step_progress_update(self, context): bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="Progress", default=0, min=0, max=context.scene.dream_textures_prompt.steps + 1, update=step_progress_update) bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info", update=step_progress_update) - global generator - if generator is None: - info("Initializing Process") - generator = GeneratorProcess() - else: - info("Waiting For Process") + info("Waiting For Process") + generator = GeneratorProcess.shared() def bpy_image(name, width, height, pixels): image = bpy.data.images.new(name, width=width, height=height) @@ -195,10 +187,7 @@ def remove_timer(context): timer = None def kill_generator(context=bpy.context): - global generator - if generator: - generator.kill() - generator = None + GeneratorProcess.kill_shared() remove_timer(context) bpy.context.scene.dream_textures_progress = 0 bpy.context.scene.dream_textures_info = "" diff --git a/operators/upscale.py b/operators/upscale.py new file mode 100644 index 00000000..2f05248f --- /dev/null +++ b/operators/upscale.py @@ -0,0 +1,100 @@ +import bpy +import tempfile +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 + +class Upscale(bpy.types.Operator): + bl_idname = "shade.dream_textures_upscale" + bl_label = f"Upscale" + bl_description = ("Upscale with Real-ESRGAN") + bl_options = {"REGISTER"} + + @classmethod + def poll(cls, context): + global timer + return timer is None + + 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'} + + def execute(self, context): + def save_temp_image(img, path=None): + path = path if path is not None else tempfile.NamedTemporaryFile().name + + settings = 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 + + input_image = None + for area in context.screen.areas: + if area.type == 'IMAGE_EDITOR': + if area.spaces.active.image is not None: + input_image = save_temp_image(area.spaces.active.image) + + if input_image is None: + self.report({"ERROR"}, "No open image in the Image Editor space") + return {"FINISHED"} + + def image_callback(output_path): + print("Received image callback") + print(output_path) + image = bpy.data.images.load(output_path) + print("Saved") + image.pack() + print("Packed") + + def info_callback(msg=""): + print("Info", msg) + def exception_callback(fatal, msg, trace): + print("Exception Callback", fatal, msg, trace) + + generator = GeneratorProcess.shared() + + args = { + 'input': input_image, + 'level': int(context.scene.dream_textures_upscale_target_size), + 'strength': float(context.scene.dream_textures_upscale_strength), + } + print("Running", args) + 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 diff --git a/requirements-mac-realesrgan.txt b/requirements-mac-realesrgan.txt new file mode 100644 index 00000000..520c4967 --- /dev/null +++ b/requirements-mac-realesrgan.txt @@ -0,0 +1,10 @@ +-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 +-e . diff --git a/requirements-win-torch-1-11-0.txt b/requirements-win-torch-1-11-0.txt index 3f22122e..0a9e3b2f 100644 --- a/requirements-win-torch-1-11-0.txt +++ b/requirements-win-torch-1-11-0.txt @@ -30,4 +30,5 @@ git+https://github.com/lstein/GFPGAN@fix-dark-cast-images#egg=gfpgan torch==1.11.0 # Same as numpy - let pip do its thing torchvision +realesrgan==0.2.5.0 -e . diff --git a/stable_diffusion b/stable_diffusion index 61f46cac..8ba5e385 160000 --- a/stable_diffusion +++ b/stable_diffusion @@ -1 +1 @@ -Subproject commit 61f46cac31b53f3f043819358382102f1d1d0e78 +Subproject commit 8ba5e385ec979f4728c0995660e9762a478a11cd diff --git a/ui/panel.py b/ui/panel.py index 65aa80a8..1c7e92f5 100644 --- a/ui/panel.py +++ b/ui/panel.py @@ -5,9 +5,12 @@ from ..operators.dream_texture import DreamTexture, ReleaseGenerator from ..operators.view_history import SCENE_UL_HistoryList, RecallHistoryEntry from ..operators.open_latest_version import OpenLatestVersion, is_force_show_download, new_version_available +from ..operators.upscale import upscale_options, Upscale from ..help_section import help_section from ..preferences import StableDiffusionPreferences +from ..absolute_path import REAL_ESRGAN_WEIGHTS_PATH import sys +import os SPACE_TYPES = {'IMAGE_EDITOR', 'NODE_EDITOR'} @@ -164,3 +167,64 @@ def draw(self, context): help_section(self.layout, context) TroubleshootingPanel.__name__ = f"DREAM_PT_dream_troubleshooting_panel_{space_type}" yield TroubleshootingPanel + +import webbrowser + +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): + bl_idname = "stable_diffusion.open_realesrgan_weights_directory" + bl_label = "Open Target Directory" + bl_description = ("Opens the directory that should contain the 'realesr-general-x4v3.pth' file") + bl_options = {"REGISTER", "INTERNAL"} + + def execute(self, context): + path = os.path.dirname(REAL_ESRGAN_WEIGHTS_PATH) + if not os.path.exists(path): + os.mkdir(path) + webbrowser.open(f"file:///{os.path.realpath(path)}") + + return {"FINISHED"} + +def upscaling_panels(): + for space_type in SPACE_TYPES: + class UpscalingPanel(Panel): + """Panel for AI Upscaling""" + bl_label = "AI Upscaling" + bl_category = "Dream" + bl_idname = f"DREAM_PT_dream_upscaling_panel_{space_type}" + bl_space_type = space_type + bl_region_type = 'UI' + bl_options = {'DEFAULT_CLOSED'} + + @classmethod + def poll(self, context): + if self.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 + 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. Place it in the weights folder with the name 'realesr-general-x4v3.pth'") + layout.operator(OpenRealESRGANWeightsDirectory.bl_idname, icon="FOLDER_REDIRECT") + layout = layout.column() + layout.enabled = os.path.exists(REAL_ESRGAN_WEIGHTS_PATH) + layout.prop(context.scene, "dream_textures_upscale_target_size") + layout.prop(context.scene, "dream_textures_upscale_strength") + layout.operator(Upscale.bl_idname) + + UpscalingPanel.__name__ = f"DREAM_PT_dream_troubleshooting_panel_{space_type}" + yield UpscalingPanel \ No newline at end of file diff --git a/weights/realesrgan/.gitignore b/weights/realesrgan/.gitignore new file mode 100644 index 00000000..c96a04f0 --- /dev/null +++ b/weights/realesrgan/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore \ No newline at end of file From 404e69098818e3967b11729a5019e245640670c2 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 3 Oct 2022 19:49:38 -0400 Subject: [PATCH 2/6] Use SharedMemory to pass image data around and support outscale argument --- __init__.py | 3 +-- generator_process.py | 43 ++++++++++++++++++++----------------------- operators/upscale.py | 41 +++++++++++++++++++++++++++-------------- ui/panel.py | 8 +++++--- 4 files changed, 53 insertions(+), 42 deletions(-) diff --git a/__init__.py b/__init__.py index 960667e2..5edd0076 100644 --- a/__init__.py +++ b/__init__.py @@ -68,8 +68,7 @@ def register(): bpy.types.Scene.dream_textures_progress = bpy.props.IntProperty(name="Progress", default=0, min=0, max=0) bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info") - bpy.types.Scene.dream_textures_upscale_target_size = bpy.props.EnumProperty(name="Target Size", items=upscale_options) - bpy.types.Scene.dream_textures_upscale_strength = bpy.props.FloatProperty(name="Strength", default=1, min=0, max=1) + bpy.types.Scene.dream_textures_upscale_outscale = bpy.props.EnumProperty(name="Target Size", items=upscale_options) for cls in CLASSES: bpy.utils.register_class(cls) diff --git a/generator_process.py b/generator_process.py index f1912f90..d118473f 100644 --- a/generator_process.py +++ b/generator_process.py @@ -9,6 +9,7 @@ import numpy as np from enum import IntEnum import tempfile +from multiprocessing.shared_memory import SharedMemory MISSING_DEPENDENCIES_ERROR = "Python dependencies are missing. Click Download Latest Release to fix." @@ -22,7 +23,6 @@ class Action(IntEnum): STEP_IMAGE = 3 STEP_NO_SHOW = 4 EXCEPTION = 5 - IMAGE_PATH = 6 # An image result that returns a path to a saved image instead of bytes @classmethod def _missing_(cls, value): @@ -114,7 +114,7 @@ def upscale(self, args, image_callback, info_callback, exception_callback): queue = self.queue callbacks = { Action.INFO: info_callback, - Action.IMAGE_PATH: image_callback, + Action.IMAGE: image_callback, Action.EXCEPTION: exception_callback } @@ -124,7 +124,7 @@ def upscale(self, args, image_callback, info_callback, exception_callback): tup = queue.pop() action = tup[0] callbacks[action](**tup[1]) - if action == Action.IMAGE_PATH: + if action == Action.IMAGE: break elif action == Action.EXCEPTION: return @@ -149,18 +149,7 @@ def queue_exception_msg(msg): kwargs = {} if kwargs_len == 0 else json.loads(reader.read(kwargs_len)) payload_len = readUInt(8) - if action in [Action.INFO, Action.STEP_NO_SHOW, Action.IMAGE_PATH]: - queue.append((action, kwargs)) - elif action in [Action.IMAGE, Action.STEP_IMAGE]: - expected_len = kwargs['width']*kwargs['height']*16 - if payload_len != expected_len: - queue_exception_msg(f"Internal error, received image payload of {payload_len} bytes, expected {expected_len} bytes for {kwargs['width']}x{kwargs['height']} image") - return - if expected_len > len(image_buffer): - image_buffer = bytearray(expected_len) - m = memoryview(image_buffer)[:expected_len] - reader.readinto(m) - kwargs['pixels'] = np.frombuffer(m,dtype=np.float32) + if action in [Action.INFO, Action.STEP_NO_SHOW, Action.IMAGE, Action.STEP_IMAGE]: queue.append((action, kwargs)) elif action == Action.EXCEPTION: queue.append((action, kwargs)) @@ -171,6 +160,8 @@ def queue_exception_msg(msg): return def main(): + shared_memory: SharedMemory | None = None + stdin = sys.stdin.buffer stdout = sys.stdout.buffer sys.stdout = open(os.devnull, 'w') # prevent stable diffusion logs from breaking ipc @@ -271,9 +262,17 @@ def image_to_bytes(image): return (np.asarray(ImageOps.flip(image).convert('RGBA'),dtype=np.float32) * byte_to_normalized).tobytes() def image_writer(image, seed, upscaled=False, first_seed=None): - # Only use the non-upscaled texture, as upscaling is currently unsupported by the addon. + nonlocal shared_memory + # Only use the non-upscaled texture, as upscaling is a separate step in this addon. if not upscaled: - send_action(Action.IMAGE, payload=image_to_bytes(image), seed=seed, width=image.width, height=image.height) + image_bytes = image_to_bytes(image) + image_bytes_len = len(image_bytes) + if shared_memory is None or shared_memory.size != image_bytes_len: + if shared_memory is not None: + shared_memory.close() + shared_memory = SharedMemory(create=True, size=image_bytes_len) + shared_memory.buf[:] = image_bytes + send_action(Action.IMAGE, shared_memory_name=shared_memory.name, seed=seed, width=image.width, height=image.height) step = 0 def view_step(samples): @@ -424,14 +423,12 @@ def update(self, n=1): half=False, gpu_id=None ) - send_info("Enhancing") - output, _ = upsampler.enhance(image, outscale=4) - send_info("Upscaled") + send_info("Enhancing Input") + output, _ = upsampler.enhance(image, outscale=args['outscale']) + send_info("Converting Result") output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) output = Image.fromarray(output) - output_path = tempfile.NamedTemporaryFile().name - output.save(output_path, format="png") - send_action(Action.IMAGE_PATH, output_path=output_path) + image_writer(output, args['name']) except: send_exception() return diff --git a/operators/upscale.py b/operators/upscale.py index 2f05248f..b5759de1 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -1,5 +1,8 @@ import bpy import tempfile +from multiprocessing.shared_memory import SharedMemory +import numpy as np +import sys from ..generator_process import GeneratorProcess upscale_options = [ @@ -18,7 +21,7 @@ def remove_timer(context): class Upscale(bpy.types.Operator): bl_idname = "shade.dream_textures_upscale" - bl_label = f"Upscale" + bl_label = "Upscale" bl_description = ("Upscale with Real-ESRGAN") bl_options = {"REGISTER"} @@ -41,6 +44,8 @@ def modal(self, context, event): return {'RUNNING_MODAL'} def execute(self, context): + scene = context.scene + def save_temp_image(img, path=None): path = path if path is not None else tempfile.NamedTemporaryFile().name @@ -62,36 +67,44 @@ def save_temp_image(img, path=None): return path input_image = None + input_image_path = None for area in context.screen.areas: if area.type == 'IMAGE_EDITOR': if area.spaces.active.image is not None: - input_image = save_temp_image(area.spaces.active.image) + input_image = area.spaces.active.image + input_image_path = save_temp_image(input_image) if input_image is None: self.report({"ERROR"}, "No open image in the Image Editor space") return {"FINISHED"} - def image_callback(output_path): - print("Received image callback") - print(output_path) - image = bpy.data.images.load(output_path) - print("Saved") + def bpy_image(name, width, height, pixels): + image = bpy.data.images.new(name, width=width, height=height) + image.pixels[:] = pixels image.pack() - print("Packed") + return image + + def image_callback(shared_memory_name, seed, width, height): + scene.dream_textures_info = "" + shared_memory = SharedMemory(shared_memory_name) + bpy_image(seed + ' (Upscaled)', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) + shared_memory.close() def info_callback(msg=""): - print("Info", msg) + scene.dream_textures_info = msg def exception_callback(fatal, msg, trace): - print("Exception Callback", fatal, msg, trace) + scene.dream_textures_info = "" + self.report({'ERROR'}, msg) + if trace: + print(trace, file=sys.stderr) generator = GeneratorProcess.shared() args = { - 'input': input_image, - 'level': int(context.scene.dream_textures_upscale_target_size), - 'strength': float(context.scene.dream_textures_upscale_strength), + 'input': input_image_path, + 'name': input_image.name, + 'outscale': int(context.scene.dream_textures_upscale_outscale), } - print("Running", args) global generator_advance generator_advance = generator.upscale(args, image_callback, info_callback, exception_callback) context.window_manager.modal_handler_add(self) diff --git a/ui/panel.py b/ui/panel.py index 1c7e92f5..e8e6de30 100644 --- a/ui/panel.py +++ b/ui/panel.py @@ -222,9 +222,11 @@ def draw(self, context): layout.operator(OpenRealESRGANWeightsDirectory.bl_idname, icon="FOLDER_REDIRECT") layout = layout.column() layout.enabled = os.path.exists(REAL_ESRGAN_WEIGHTS_PATH) - layout.prop(context.scene, "dream_textures_upscale_target_size") - layout.prop(context.scene, "dream_textures_upscale_strength") - layout.operator(Upscale.bl_idname) + layout.prop(context.scene, "dream_textures_upscale_outscale") + 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 From 1d4f949afd1db18736ad40d8ae9ec714c97eff12 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 3 Oct 2022 20:03:58 -0400 Subject: [PATCH 3/6] Updates for latest submodule changes --- generator_process.py | 34 +++++++++++++++++---------------- operators/dream_texture.py | 18 ++++++++++------- property_groups/dream_prompt.py | 9 ++++++++- ui/panel.py | 3 +-- 4 files changed, 38 insertions(+), 26 deletions(-) diff --git a/generator_process.py b/generator_process.py index d118473f..7cc3fb78 100644 --- a/generator_process.py +++ b/generator_process.py @@ -138,7 +138,6 @@ def readUInt(length): def queue_exception_msg(msg): queue.append((Action.EXCEPTION, {'fatal': True, 'msg': msg, 'trace': None})) - image_buffer = bytearray(512*512*16) while not self.killed: action = readUInt(ACTION_BYTE_LENGTH) if action == Action.CLOSED: @@ -261,25 +260,28 @@ def send_exception(fatal = True, msg: str = None, trace: str = None): def image_to_bytes(image): return (np.asarray(ImageOps.flip(image).convert('RGBA'),dtype=np.float32) * byte_to_normalized).tobytes() - def image_writer(image, seed, upscaled=False, first_seed=None): + def share_image_memory(image): nonlocal shared_memory + image_bytes = image_to_bytes(image) + image_bytes_len = len(image_bytes) + if shared_memory is None or shared_memory.size != image_bytes_len: + if shared_memory is not None: + shared_memory.close() + shared_memory = SharedMemory(create=True, size=image_bytes_len) + shared_memory.buf[:] = image_bytes + return shared_memory.name + + 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: - image_bytes = image_to_bytes(image) - image_bytes_len = len(image_bytes) - if shared_memory is None or shared_memory.size != image_bytes_len: - if shared_memory is not None: - shared_memory.close() - shared_memory = SharedMemory(create=True, size=image_bytes_len) - shared_memory.buf[:] = image_bytes - send_action(Action.IMAGE, shared_memory_name=shared_memory.name, seed=seed, width=image.width, height=image.height) + send_action(Action.IMAGE, shared_memory_name=share_image_memory(image), seed=seed, width=image.width, height=image.height) step = 0 def view_step(samples): nonlocal step if args['show_steps']: - image = generator._sample_to_image(samples) - send_action(Action.STEP_IMAGE, payload=image_to_bytes(image), step=step, width=image.width, height=image.height) + image = generator.sample_to_image(samples) + send_action(Action.STEP_IMAGE, shared_memory_name=share_image_memory(image), step=step, width=image.width, height=image.height) else: send_action(Action.STEP_NO_SHOW, step=step) step += 1 @@ -348,7 +350,7 @@ def update(self, n=1): return # stdin closed args = json.loads(stdin.read(json_len)) - if generator is None or (generator.full_precision != args['full_precision'] and sys.platform != 'darwin'): + if generator is None or generator.precision != args['precision']: send_info("Loading Model") try: generator = Generate( @@ -357,7 +359,7 @@ def update(self, n=1): # These args are deprecated, but we need them to specify an absolute path to the weights. weights=weights, config=config, - full_precision=args['full_precision'] + 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() @@ -384,9 +386,9 @@ def update(self, n=1): import re low_ram = re.search(r"(Not enough memory, use lower resolution)( \(max approx. \d+x\d+\))",s,re.IGNORECASE) if low_ram: - send_exception(False, f"{low_ram[1]}{' or disable full precision' if args['full_precision'] else ''}{low_ram[2]}", s) + 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: - send_exception(False, f"Not enough memory, use lower resolution{' or disable full precision' if args['full_precision'] else ''}", s) + send_exception(False, f"Not enough memory, use lower resolution{' or disable full precision' if args['precision'] == 'float32' else ''}", s) else: send_exception(True, msg=None, trace=s) # consider all unknown exceptions to be fatal so the generator process is fully restarted next time return diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 647901ec..11d0d189 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -1,8 +1,8 @@ -from importlib.resources import path import sys import bpy import os -import math +import numpy as np +from multiprocessing.shared_memory import SharedMemory from ..preferences import StableDiffusionPreferences from ..pil_to_image import * @@ -92,7 +92,7 @@ def bpy_image(name, width, height, pixels): image.pack() return image - def image_writer(seed, width, height, pixels, upscaled=False): + def image_writer(shared_memory_name, seed, width, height, upscaled=False): info() # clear variable global last_data_block # Only use the non-upscaled texture, as upscaling is currently unsupported by the addon. @@ -102,7 +102,9 @@ def image_writer(seed, width, height, pixels, upscaled=False): last_data_block = None if generator is None or generator.process.poll() or width == 0 or height == 0: return # process was closed - image = bpy_image(f"{seed}", width, height, pixels) + 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") @@ -116,15 +118,17 @@ def image_writer(seed, width, height, pixels, upscaled=False): history_entry.seed = str(seed) history_entry.random_seed = False - def view_step(step, width=None, height=None, pixels=None): + def view_step(step, width=None, height=None, shared_memory_name=None): info() # clear variable scene.dream_textures_progress = step + 1 - if pixels is None: + if shared_memory_name is None: return # show steps disabled global last_data_block for area in screen.areas: if area.type == 'IMAGE_EDITOR': - step_image = bpy_image(f'Step {step + 1}/{scene.dream_textures_prompt.steps}', width, height, pixels) + 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) diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index a68148e3..59dd0cde 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -14,6 +14,13 @@ ("k_heun", "KHEUN", "", 8), ] +precision_options = [ + ('auto', 'Automatic', "", 1), + ('float32', 'Full Precision (float32)', "", 2), + ('autocast', 'Autocast', "", 3), + ('float16', 'Half Precision (float16)', "", 4), +] + def seed_clamp(self, ctx): # clamp seed right after input to make it clear what the limits are try: @@ -38,7 +45,7 @@ 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), - "full_precision": BoolProperty(name="Full Precision", default=True if sys.platform == 'darwin' else False, description="Whether to use full precision or half precision floats. Full precision is slower, but required by some GPUs"), + "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, description="How strongly the prompt influences the image"), diff --git a/ui/panel.py b/ui/panel.py index e8e6de30..2691989b 100644 --- a/ui/panel.py +++ b/ui/panel.py @@ -68,8 +68,7 @@ def draw_panel(self, context): advanced_box_heading.prop(scene.dream_textures_prompt, "show_advanced", icon="DOWNARROW_HLT" if scene.dream_textures_prompt.show_advanced else "RIGHTARROW_THIN", emboss=False, icon_only=True) advanced_box_heading.label(text="Advanced Configuration") if scene.dream_textures_prompt.show_advanced: - if sys.platform not in {'darwin'}: - advanced_box.prop(scene.dream_textures_prompt, "full_precision") + advanced_box.prop(scene.dream_textures_prompt, "precision") advanced_box.prop(scene.dream_textures_prompt, "random_seed") seed_row = advanced_box.row() seed_row.prop(scene.dream_textures_prompt, "seed") From 0666a23577c9f6469abc6d6f1eb406e8323e61ff Mon Sep 17 00:00:00 2001 From: NullSenseStudio <47096043+NullSenseStudio@users.noreply.github.com> Date: Mon, 3 Oct 2022 20:04:48 -0400 Subject: [PATCH 4/6] fix cv2 ImportError (#237) --- generator_process.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/generator_process.py b/generator_process.py index 7cc3fb78..f373c512 100644 --- a/generator_process.py +++ b/generator_process.py @@ -44,7 +44,8 @@ def _missing_(cls, value): _shared_instance = None class GeneratorProcess(): def __init__(self): - self.process = subprocess.Popen([sys.executable,'generator_process.py'],cwd=os.path.dirname(os.path.realpath(__file__)),stdin=subprocess.PIPE,stdout=subprocess.PIPE) + import bpy + self.process = subprocess.Popen([sys.executable,'generator_process.py',bpy.app.binary_path],cwd=os.path.dirname(os.path.realpath(__file__)),stdin=subprocess.PIPE,stdout=subprocess.PIPE) self.reader = self.process.stdout self.queue = [] self.args = None @@ -221,6 +222,10 @@ def send_exception(fatal = True, msg: str = None, trace: str = None): send_action(Action.EXCEPTION, fatal=fatal, msg=msg, trace=trace) try: + if sys.platform == 'win32': + from ctypes import WinDLL + WinDLL(os.path.join(os.path.dirname(sys.argv[1]),"python3.dll")) # fix for ImportError: DLL load failed while importing cv2: The specified module could not be found. + from absolute_path import absolute_path # Support Apple Silicon GPUs as much as possible. os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" From 08c9043068d83025ed8cc283d0ed854a3afffdbf Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 3 Oct 2022 20:14:31 -0400 Subject: [PATCH 5/6] Support Image Texture nodes --- operators/dream_texture.py | 2 +- operators/upscale.py | 26 +++++++++++++++++++------- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 11d0d189..f582e483 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -127,8 +127,8 @@ def view_step(step, width=None, height=None, shared_memory_name=None): 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() + step_image = bpy_image(f'Step {step + 1}/{scene.dream_textures_prompt.steps}', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) area.spaces.active.image = step_image if last_data_block is not None: bpy.data.images.remove(last_data_block) diff --git a/operators/upscale.py b/operators/upscale.py index b5759de1..619cb2db 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -45,6 +45,9 @@ def modal(self, context, event): 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) def save_temp_image(img, path=None): path = path if path is not None else tempfile.NamedTemporaryFile().name @@ -68,14 +71,18 @@ def save_temp_image(img, path=None): input_image = None input_image_path = None - for area in context.screen.areas: - if area.type == 'IMAGE_EDITOR': - if area.spaces.active.image is not None: - input_image = area.spaces.active.image - input_image_path = save_temp_image(input_image) + if active_node is not None and active_node.image is not None: + input_image = active_node.image + input_image_path = save_temp_image(input_image) + else: + for area in context.screen.areas: + if area.type == 'IMAGE_EDITOR': + if area.spaces.active.image is not None: + input_image = area.spaces.active.image + input_image_path = save_temp_image(input_image) if input_image is None: - self.report({"ERROR"}, "No open image in the Image Editor space") + self.report({"ERROR"}, "No open image in the Image Editor space, or selected Image Texture node.") return {"FINISHED"} def bpy_image(name, width, height, pixels): @@ -87,7 +94,12 @@ def bpy_image(name, width, height, pixels): def image_callback(shared_memory_name, seed, width, height): scene.dream_textures_info = "" shared_memory = SharedMemory(shared_memory_name) - bpy_image(seed + ' (Upscaled)', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) + image = bpy_image(seed + ' (Upscaled)', width, height, np.frombuffer(shared_memory.buf,dtype=np.float32)) + for area in context.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=""): From 2b5f9748580e30bc9aac76649c0fca14795d0a71 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 3 Oct 2022 21:15:26 -0400 Subject: [PATCH 6/6] Revisions from review --- __init__.py | 1 + generator_process.py | 8 +++----- operators/dream_texture.py | 3 +-- operators/upscale.py | 12 +++++++++++- ui/panel.py | 5 +++++ 5 files changed, 21 insertions(+), 8 deletions(-) diff --git a/__init__.py b/__init__.py index 5edd0076..498ea874 100644 --- a/__init__.py +++ b/__init__.py @@ -69,6 +69,7 @@ def register(): bpy.types.Scene.dream_textures_info = bpy.props.StringProperty(name="Info") 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) for cls in CLASSES: bpy.utils.register_class(cls) diff --git a/generator_process.py b/generator_process.py index f373c512..b5f46333 100644 --- a/generator_process.py +++ b/generator_process.py @@ -416,19 +416,17 @@ def update(self, n=1): from realesrgan.archs.srvgg_arch import SRVGGNetCompact # image = Image.open(args['input']) image = cv2.imread(args['input'], cv2.IMREAD_UNCHANGED) - model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu') + real_esrgan_model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu') netscale = 4 send_info("Loading Upsampler") upsampler = RealESRGANer( scale=netscale, model_path=REAL_ESRGAN_WEIGHTS_PATH, - dni_weight=1, - model=model, + model=real_esrgan_model, tile=0, tile_pad=10, pre_pad=0, - half=False, - gpu_id=None + half=not args['full_precision'] ) send_info("Enhancing Input") output, _ = upsampler.enhance(image, outscale=args['outscale']) diff --git a/operators/dream_texture.py b/operators/dream_texture.py index f582e483..6c87c04e 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -72,7 +72,6 @@ def handle_exception(fatal, msg, trace): 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: @@ -127,8 +126,8 @@ def view_step(step, width=None, height=None, shared_memory_name=None): for area in screen.areas: if area.type == 'IMAGE_EDITOR': shared_memory = SharedMemory(shared_memory_name) - shared_memory.close() 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) diff --git a/operators/upscale.py b/operators/upscale.py index 619cb2db..bf24d53e 100644 --- a/operators/upscale.py +++ b/operators/upscale.py @@ -47,7 +47,16 @@ 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) + 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 + + 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_info = bpy.props.StringProperty(name="Info", update=step_progress_update) def save_temp_image(img, path=None): path = path if path is not None else tempfile.NamedTemporaryFile().name @@ -116,6 +125,7 @@ def exception_callback(fatal, msg, trace): '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, } global generator_advance generator_advance = generator.upscale(args, image_callback, info_callback, exception_callback) diff --git a/ui/panel.py b/ui/panel.py index 2691989b..eae703e6 100644 --- a/ui/panel.py +++ b/ui/panel.py @@ -222,6 +222,11 @@ def draw(self, context): 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") + 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: