diff --git a/CHANGELOG.md b/CHANGELOG.md index d31a0fbf5..5257b7de2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,13 +18,14 @@ instead of being used implicitly via quantization, thanks @CalamitousFelicitousness - removed: old `codeformer` and `gfpgan` face restorers, thanks @CalamitousFelicitousness - **UI** - - ui: **themes** add *CTD-NT64Light* and *CTD-NT64Dark*, thanks @resonantsky + - ui: **themes** add *CTD-NT64Light*, *CTD-NT64Medium* and *CTD-NT64Dark*, thanks @resonantsky - ui: **gallery** add option to auto-refresh gallery, thanks @awsr - **Internal** - refactor: to/from image/tensor logic, thanks @CalamitousFelicitousness - refactor: switch to `pyproject.toml` for tool configs - refactor: reorganize `cli` scripts - refactor: move tests to dedicated `/test/` + - refactor: all image handling to `modules/image/` - refactor: captioning part-2, thanks @CalamitousFelicitousness - refactor: remove face restoration, thanks @CalamitousFelicitousness - update `lint` rules, thanks @awsr diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 947696b04..02b0af3d9 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 947696b04a60e2aba761242fa97ba4c44bf1cdac +Subproject commit 02b0af3d9927ae09c8c1b8c9215a33c6c00e3fc2 diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py index f2f5d741e..efae5bbca 100644 --- a/modules/caption/joytag.py +++ b/modules/caption/joytag.py @@ -17,7 +17,8 @@ import torchvision import einops from einops.layers.torch import Rearrange import huggingface_hub -from modules import shared, devices, sd_models, images_sharpfin +from modules import shared, devices, sd_models +from modules.image import convert model = None @@ -1034,8 +1035,8 @@ def prepare_image(image: Image.Image, target_size: int) -> torch.Tensor: padded_image.paste(image, (pad_left, pad_top)) if max_dim != target_size: padded_image = padded_image.resize((target_size, target_size), Image.Resampling.LANCZOS) - image_tensor = images_sharpfin.to_tensor(padded_image) - image_tensor = images_sharpfin.normalize(image_tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]) + image_tensor = convert.to_tensor(padded_image) + image_tensor = convert.normalize(image_tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]) return image_tensor diff --git a/modules/image/convert.py b/modules/image/convert.py new file mode 100644 index 000000000..042436fd2 --- /dev/null +++ b/modules/image/convert.py @@ -0,0 +1,81 @@ +import sys +import torch +import numpy as np +from PIL import Image +from installer import log + + +def to_tensor(image: Image.Image | np.ndarray): + """PIL Image -> float32 CHW tensor [0,1]. Pure torch, no torchvision.""" + # fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access + if not isinstance(image, Image.Image): + pic = np.array(image, copy=True) + elif isinstance(image, np.ndarray): + pic = image.copy() + else: + raise TypeError(f"Expected PIL Image or np.ndarray, got {type(image)}") + if pic.ndim == 2: + pic = pic[:, :, np.newaxis] + tensor = torch.from_numpy(pic.transpose((2, 0, 1))).contiguous() + # log.debug(f'Convert: source={type(image)} target={tensor.shape} fn={fn}') + if tensor.dtype == torch.uint8: + return tensor.to(torch.float32).div_(255.0) + return tensor.to(torch.float32) + + +def to_pil(tensor: torch.Tensor | np.ndarray): + """Float CHW/HWC or BCHW/BHWC tensor [0,1] -> PIL Image. Pure torch, no torchvision.""" + if isinstance(tensor, torch.Tensor): + tensor = tensor.detach().cpu() + elif isinstance(tensor, np.ndarray): + tensor = torch.from_numpy(tensor) + else: + raise TypeError(f"Expected torch.Tensor, got {type(tensor)}") + try: + if tensor.dim() == 4: + if tensor.shape[-1] in (1, 3, 4) and tensor.shape[-1] < tensor.shape[-2]: # BHWC + tensor = tensor.permute(0, 3, 1, 2) + tensor = tensor[0] + elif tensor.dim() == 3: + if tensor.shape[-1] in (1, 3, 4) and tensor.shape[-1] < tensor.shape[-2] and tensor.shape[-1] < tensor.shape[-3]: # HWC + tensor = tensor.permute(2, 0, 1) + if tensor.dtype != torch.uint8: + tensor = (tensor.clamp(0, 1) * 255).round().to(torch.uint8) + ndarr = tensor.permute(1, 2, 0).numpy() + if ndarr.shape[2] == 1: + ndarr = ndarr[:, :, 0] + mode = 'L' + elif ndarr.shape[2] == 3: + mode = 'RGB' + else: + mode = 'RGBA' + image = Image.fromarray(ndarr, mode=mode) + except Exception as e: + image = Image.new('RGB', (tensor.shape[-1], tensor.shape[-2]), color=(152, 32, 48)) + fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access + log.error(f'Convert: source={type(tensor)} target={image} fn={fn} {e}') + return image + + +def pil_to_tensor(image): + """PIL Image -> uint8 CHW tensor (no float scaling). Replaces TF.pil_to_tensor.""" + if not isinstance(image, Image.Image): + raise TypeError(f"Expected PIL Image, got {type(image)}") + pic = np.array(image, copy=True) + if pic.ndim == 2: + pic = pic[:, :, np.newaxis] + return torch.from_numpy(pic.transpose((2, 0, 1))).contiguous() + + +def normalize(tensor, mean, std, inplace=False): + """Tensor normalization. Replaces TF.normalize.""" + if not inplace: + tensor = tensor.clone() + mean_t = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device) + std_t = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device) + if mean_t.ndim == 1: + mean_t = mean_t[:, None, None] + if std_t.ndim == 1: + std_t = std_t[:, None, None] + tensor.sub_(mean_t).div_(std_t) + return tensor diff --git a/modules/images_grid.py b/modules/image/grid.py similarity index 100% rename from modules/images_grid.py rename to modules/image/grid.py diff --git a/modules/image/metadata.py b/modules/image/metadata.py new file mode 100644 index 000000000..5b91c57c9 --- /dev/null +++ b/modules/image/metadata.py @@ -0,0 +1,189 @@ +import io +import re +import json +import piexif +from PIL import Image, ExifTags +from modules import shared, errors, sd_samplers +from modules.image.watermark import get_watermark + + +def safe_decode_string(s: bytes): + remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment + for encoding in ['utf_16_be', 'utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings + try: + s = remove_prefix(s, b'UNICODE') + s = remove_prefix(s, b'ASCII') + s = remove_prefix(s, b'\x00') + val = s.decode(encoding, errors="strict") + val = re.sub(r'[\x00-\x09]', '', val).strip() # remove remaining special characters + if len(val) == 0: # remove empty strings + val = None + return val + except Exception: + pass + return None + + +def parse_comfy_metadata(data: dict): + def parse_workflow(): + res = '' + try: + txt = data.get('workflow', {}) + dct = json.loads(txt) + nodes = len(dct.get('nodes', [])) + version = dct.get('extra', {}).get('frontendVersion', 'unknown') + if version is not None: + res = f" | Version: {version} | Nodes: {nodes}" + except Exception: + pass + return res + + def parse_prompt(): + res = '' + try: + txt = data.get('prompt', {}) + dct = json.loads(txt) + for val in dct.values(): + inp = val.get('inputs', {}) + if 'model' in inp: + model = inp.get('model', None) + if isinstance(model, str) and len(model) > 0: + res += f" | Model: {model} | Class: {val.get('class_type', '')}" + except Exception: + pass + return res + + workflow = parse_workflow() + prompt = parse_prompt() + if len(workflow) > 0 or len(prompt) > 0: + parsed = f'App: ComfyUI{workflow}{prompt}' + shared.log.info(f'Image metadata: {parsed}') + return parsed + return '' + + +def parse_invoke_metadata(data: dict): + def parse_metadtaa(): + res = '' + try: + txt = data.get('invokeai_metadata', {}) + dct = json.loads(txt) + if 'app_version' in dct: + version = dct['app_version'] + if isinstance(version, str) and len(version) > 0: + res += f" | Version: {version}" + except Exception: + pass + return res + + metadata = parse_metadtaa() + if len(metadata) > 0: + parsed = f'App: InvokeAI{metadata}' + shared.log.info(f'Image metadata: {parsed}') + return parsed + return '' + + +def parse_novelai_metadata(data: dict): + geninfo = '' + if data.get("Software", None) == "NovelAI": + try: + dct = json.loads(data["Comment"]) + sampler = sd_samplers.samplers_map.get(dct["sampler"], "Euler a") + geninfo = f'{data["Description"]} Negative prompt: {dct["uc"]} Steps: {dct["steps"]}, Sampler: {sampler}, CFG scale: {dct["scale"]}, Seed: {dct["seed"]}, Clip skip: 2, ENSD: 31337' + except Exception: + pass + return geninfo + + +def read_info_from_image(image: Image.Image, watermark: bool = False) -> tuple[str, dict]: + if image is None: + return '', {} + if isinstance(image, str): + try: + image = Image.open(image) + image.load() + except Exception: + return '', {} + items = image.info or {} + geninfo = items.pop('parameters', None) or items.pop('UserComment', None) or '' + if isinstance(geninfo, dict): + if 'UserComment' in geninfo: + geninfo = geninfo['UserComment'] # Info was nested + else: + geninfo = '' # Unknown format. Ignore contents + items['UserComment'] = geninfo + + if "exif" in items: + try: + exif = piexif.load(items["exif"]) + except Exception as e: + shared.log.error(f'Error loading EXIF data: {e}') + exif = {} + for _key, subkey in exif.items(): + if isinstance(subkey, dict): + for key, val in subkey.items(): + if isinstance(val, bytes): # decode bytestring + val = safe_decode_string(val) + if isinstance(val, tuple) and isinstance(val[0], int) and isinstance(val[1], int) and val[1] > 0: # convert camera ratios + val = round(val[0] / val[1], 2) + if val is not None and key in ExifTags.TAGS: # add known tags + if ExifTags.TAGS[key] == 'UserComment': # add geninfo from UserComment + geninfo = str(val) + items['parameters'] = val + else: + items[ExifTags.TAGS[key]] = val + elif val is not None and key in ExifTags.GPSTAGS: + items[ExifTags.GPSTAGS[key]] = val + if watermark: + wm = get_watermark(image) + if wm != '': + # geninfo += f' Watermark: {wm}' + items['watermark'] = wm + + for key, val in items.items(): + if isinstance(val, bytes): # decode bytestring + items[key] = safe_decode_string(val) + + geninfo += parse_comfy_metadata(items) + geninfo += parse_invoke_metadata(items) + geninfo += parse_novelai_metadata(items) + + for key in ['exif', 'ExifOffset', 'JpegIFOffset', 'JpegIFByteCount', 'ExifVersion', 'icc_profile', 'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'adobe', 'photoshop', 'loop', 'duration', 'dpi']: # remove unwanted tags + items.pop(key, None) + + try: + items['width'] = image.width + items['height'] = image.height + items['mode'] = image.mode + except Exception: + pass + + return geninfo, items + + +def image_data(data): + import gradio as gr + if data is None: + return gr.update(), None + err1 = None + err2 = None + try: + image = Image.open(io.BytesIO(data)) + image.load() + info, _ = read_info_from_image(image) + errors.log.debug(f'Decoded object: image={image} metadata={info}') + return info, None + except Exception as e: + err1 = e + try: + if len(data) > 1024 * 10: + errors.log.warning(f'Error decoding object: data too long: {len(data)}') + return gr.update(), None + info = data.decode('utf8') + errors.log.debug(f'Decoded object: data={len(data)} metadata={info}') + return info, None + except Exception as e: + err2 = e + errors.log.error(f'Error decoding object: {err1 or err2}') + return gr.update(), None diff --git a/modules/images_namegen.py b/modules/image/namegen.py similarity index 100% rename from modules/images_namegen.py rename to modules/image/namegen.py diff --git a/modules/images_resize.py b/modules/image/resize.py similarity index 97% rename from modules/images_resize.py rename to modules/image/resize.py index 67d14ac2c..029e2549e 100644 --- a/modules/images_resize.py +++ b/modules/image/resize.py @@ -4,7 +4,8 @@ import time import numpy as np import torch from PIL import Image -from modules import shared, upscaler, images_sharpfin +from modules import shared, upscaler +from modules.image import sharpfin def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width: int, height: int, upscaler_name: str=None, output_type: str='image', context: str=None): @@ -36,7 +37,7 @@ def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width: def resize(im: Union[Image.Image, torch.Tensor], w, h): w, h = int(w), int(h) if upscaler_name is None or upscaler_name == "None" or (hasattr(im, 'mode') and im.mode == 'L'): - return images_sharpfin.resize(im, (w, h), linearize=False) # force for mask + return sharpfin.resize(im, (w, h), linearize=False) # force for mask if isinstance(im, torch.Tensor): scale = max(w // 8 / im.shape[-1] , h // 8 / im.shape[-2]) else: @@ -53,7 +54,7 @@ def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width: shared.log.warning(f"Resize upscaler: invalid={upscaler_name} fallback={selected_upscaler.name}") shared.log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}") if isinstance(im, Image.Image) and (im.width != w or im.height != h): # probably downsample after upscaler created larger image - im = images_sharpfin.resize(im, (w, h)) + im = sharpfin.resize(im, (w, h)) return im def crop(im: Image.Image): diff --git a/modules/image/save.py b/modules/image/save.py new file mode 100644 index 000000000..8b592bd62 --- /dev/null +++ b/modules/image/save.py @@ -0,0 +1,215 @@ +import os +import sys +import queue +import datetime +import threading +import piexif.helper +from PIL import Image, PngImagePlugin +from modules import shared, script_callbacks, errors, paths +from modules.image.grid import check_grid_size +from modules.image.namegen import FilenameGenerator +from modules.image.watermark import set_watermark + + +debug = errors.log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None +debug_save = errors.log.trace if os.environ.get('SD_SAVE_DEBUG', None) is not None else lambda *args, **kwargs: None +try: + from pi_heif import register_heif_opener + register_heif_opener() +except Exception: + pass + + +def sanitize_filename_part(text, replace_spaces=True): + if text is None: + return None + if replace_spaces: + text = text.replace(' ', '_') + invalid_filename_chars = '#<>:"/\\|?*\n\r\t' + invalid_filename_prefix = ' ' + invalid_filename_postfix = ' .' + max_filename_part_length = 64 + text = text.translate({ord(x): '_' for x in invalid_filename_chars}) + text = text.lstrip(invalid_filename_prefix)[:max_filename_part_length] + text = text.rstrip(invalid_filename_postfix) + return text + + +def atomically_save_image(): + Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes + while True: + image, filename, extension, params, exifinfo, filename_txt, is_grid = save_queue.get() + jobid = shared.state.begin('Save image') + shared.state.image_history += 1 + if len(exifinfo) > 2: + with open(paths.params_path, "w", encoding="utf8") as file: + file.write(exifinfo) + fn = filename + extension + filename = filename.strip() + if extension[0] != '.': # add dot if missing + extension = '.' + extension + try: + image_format = Image.registered_extensions()[extension] + except Exception: + shared.log.warning(f'Save: unknown image format: {extension}') + image_format = 'JPEG' + exifinfo = (exifinfo or "") if shared.opts.image_metadata else "" + # additional metadata saved in files + if shared.opts.save_txt and len(exifinfo) > 0: + try: + with open(filename_txt, "w", encoding="utf8") as file: + file.write(f"{exifinfo}\n") + shared.log.info(f'Save: text="{filename_txt}" len={len(exifinfo)}') + except Exception as e: + shared.log.warning(f'Save failed: description={filename_txt} {e}') + + # actual save + if image_format == 'PNG': + pnginfo_data = PngImagePlugin.PngInfo() + for k, v in params.pnginfo.items(): + pnginfo_data.add_text(k, str(v)) + debug_save(f'Save pnginfo: {params.pnginfo.items()}') + save_args = { 'compress_level': 6, 'pnginfo': pnginfo_data if shared.opts.image_metadata else None } + elif image_format == 'JPEG': + if image.mode == 'RGBA': + shared.log.warning('Save: removing alpha channel') + image = image.convert("RGB") + elif image.mode == 'I;16': + image = image.point(lambda p: p * 0.0038910505836576).convert("L") + save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality } + if shared.opts.image_metadata: + debug_save(f'Save exif: {exifinfo}') + save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) + elif image_format == 'WEBP': + if image.mode == 'I;16': + image = image.point(lambda p: p * 0.0038910505836576).convert("RGB") + save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless } + if shared.opts.image_metadata: + debug_save(f'Save exif: {exifinfo}') + save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) + elif image_format == 'JXL': + if image.mode == 'I;16': + image = image.point(lambda p: p * 0.0038910505836576).convert("RGB") + elif image.mode not in {"RGB", "RGBA"}: + image = image.convert("RGBA") + save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless } + if shared.opts.image_metadata: + debug_save(f'Save exif: {exifinfo}') + save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) + else: + save_args = { 'quality': shared.opts.jpeg_quality } + try: + debug_save(f'Save args: {save_args}') + image.save(fn, format=image_format, **save_args) + except Exception as e: + shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}') + errors.display(e, 'Image save') + size = os.path.getsize(fn) if os.path.exists(fn) else 0 + what = 'grid' if is_grid else 'image' + shared.log.info(f'Save: {what}="{fn}" type={image_format} width={image.width} height={image.height} size={size}') + + if shared.opts.save_log_fn != '' and len(exifinfo) > 0: + fn = os.path.join(paths.data_path, shared.opts.save_log_fn) + if not fn.endswith('.json'): + fn += '.json' + entries = shared.readfile(fn, silent=True) + if not isinstance(entries, list): + entries = [] + idx = len(entries) + entry = { 'id': idx, 'filename': filename, 'time': datetime.datetime.now().isoformat(), 'info': exifinfo } + entries.append(entry) + shared.writefile(entries, fn, mode='w', silent=True) + shared.log.info(f'Save: json="{fn}" records={len(entries)}') + shared.state.outputs(filename) + shared.state.end(jobid) + save_queue.task_done() + + +save_queue: queue.Queue[tuple[Image.Image, str, str, script_callbacks.ImageSaveParams, str, str | None, bool]] = queue.Queue() +save_thread = threading.Thread(target=atomically_save_image, daemon=True) +save_thread.start() + + +def save_image(image, + path=None, + basename='', + seed=None, + prompt=None, + extension=shared.opts.samples_format, + info=None, + grid=False, + pnginfo_section_name='parameters', + p=None, + existing_info=None, + forced_filename=None, + suffix='', + save_to_dirs=None, + ): + fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access + debug_save(f'Save: fn={fn}') # pylint: disable=protected-access + if image is None: + shared.log.warning('Image is none') + return None, None, None + if isinstance(image, list): + if len(image) > 1: + shared.log.warning(f'Save: images={image} multiple images provided only the first one will be saved') + image = image[0] + if not check_grid_size([image]): + return None, None, None + if path is None or path == '': # set default path to avoid errors when functions are triggered manually or via api and param is not set + path = paths.resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_save) + namegen = FilenameGenerator(p, seed, prompt, image, grid=grid) + suffix = suffix if suffix is not None else '' + basename = '' if basename is None else basename + if save_to_dirs is not None and isinstance(save_to_dirs, str) and len(save_to_dirs) > 0: + dirname = save_to_dirs + path = os.path.join(path, dirname) + elif shared.opts.save_to_dirs: + dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]") + path = os.path.join(path, dirname) + if forced_filename is None: + if shared.opts.samples_filename_pattern and len(shared.opts.samples_filename_pattern) > 0: + file_decoration = shared.opts.samples_filename_pattern + else: + file_decoration = "[seq]-[prompt_words]" + file_decoration = namegen.apply(file_decoration) + file_decoration += suffix + if file_decoration.startswith(basename): + basename = '' + filename = os.path.join(path, f"{file_decoration}.{extension}") if basename == '' else os.path.join(path, f"{basename}-{file_decoration}.{extension}") + else: + forced_filename += suffix + if forced_filename.startswith(basename): + basename = '' + filename = os.path.join(path, f"{forced_filename}.{extension}") if basename == '' else os.path.join(path, f"{basename}-{forced_filename}.{extension}") + pnginfo = existing_info or {} + if info is None: + info = image.info.get(pnginfo_section_name, '') + if info is not None: + pnginfo[pnginfo_section_name] = info + + wm_text = getattr(p, 'watermark_text', shared.opts.image_watermark) + wm_image = getattr(p, 'watermark_image', shared.opts.image_watermark_image) + image = set_watermark(image, wm_text, wm_image) + + params = script_callbacks.ImageSaveParams(image, p, filename, pnginfo) + params.filename = namegen.sanitize(filename) + dirname = os.path.dirname(params.filename) + if dirname is not None and len(dirname) > 0: + os.makedirs(dirname, exist_ok=True) + params.filename = namegen.sequence(params.filename) + params.filename = namegen.sanitize(params.filename) + # callbacks + script_callbacks.before_image_saved_callback(params) + exifinfo = params.pnginfo.get('UserComment', '') + exifinfo = exifinfo + ', ' if len(exifinfo) > 0 else '' + exifinfo += params.pnginfo.get(pnginfo_section_name, '') + filename, extension = os.path.splitext(params.filename) + filename_txt = f"{filename}.txt" if shared.opts.save_txt and len(exifinfo) > 0 else None + save_queue.put((params.image, filename, extension, params, exifinfo, filename_txt, grid)) # actual save is executed in a thread that polls data from queue + save_queue.join() + if not hasattr(params.image, 'already_saved_as'): + debug(f'Image marked: "{params.filename}"') + params.image.already_saved_as = params.filename + script_callbacks.image_saved_callback(params) + return params.filename, filename_txt, exifinfo diff --git a/modules/images_sharpfin.py b/modules/image/sharpfin.py similarity index 76% rename from modules/images_sharpfin.py rename to modules/image/sharpfin.py index 67256fe11..9acffbba0 100644 --- a/modules/images_sharpfin.py +++ b/modules/image/sharpfin.py @@ -9,9 +9,9 @@ Non-CUDA devices fall back to PIL/torch.nn.functional automatically. import sys import torch -import numpy as np from PIL import Image from installer import log +from modules.image.convert import to_tensor, to_pil _sharpfin_checked = False @@ -220,79 +220,3 @@ def resize_tensor(tensor: torch.Tensor, target_size: tuple[int, int], *, kernel= if squeezed: result = result.squeeze(0) return result - - -def to_tensor(image: Image.Image | np.ndarray): - """PIL Image -> float32 CHW tensor [0,1]. Pure torch, no torchvision.""" - # fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access - if not isinstance(image, Image.Image): - pic = np.array(image, copy=True) - elif isinstance(image, np.ndarray): - pic = image.copy() - else: - raise TypeError(f"Expected PIL Image or np.ndarray, got {type(image)}") - if pic.ndim == 2: - pic = pic[:, :, np.newaxis] - tensor = torch.from_numpy(pic.transpose((2, 0, 1))).contiguous() - # log.debug(f'Convert: source={type(image)} target={tensor.shape} fn={fn}') - if tensor.dtype == torch.uint8: - return tensor.to(torch.float32).div_(255.0) - return tensor.to(torch.float32) - - -def to_pil(tensor: torch.Tensor | np.ndarray): - """Float CHW/HWC or BCHW/BHWC tensor [0,1] -> PIL Image. Pure torch, no torchvision.""" - if isinstance(tensor, torch.Tensor): - tensor = tensor.detach().cpu() - elif isinstance(tensor, np.ndarray): - tensor = torch.from_numpy(tensor) - else: - raise TypeError(f"Expected torch.Tensor, got {type(tensor)}") - try: - if tensor.dim() == 4: - if tensor.shape[-1] in (1, 3, 4) and tensor.shape[-1] < tensor.shape[-2]: # BHWC - tensor = tensor.permute(0, 3, 1, 2) - tensor = tensor[0] - elif tensor.dim() == 3: - if tensor.shape[-1] in (1, 3, 4) and tensor.shape[-1] < tensor.shape[-2] and tensor.shape[-1] < tensor.shape[-3]: # HWC - tensor = tensor.permute(2, 0, 1) - if tensor.dtype != torch.uint8: - tensor = (tensor.clamp(0, 1) * 255).round().to(torch.uint8) - ndarr = tensor.permute(1, 2, 0).numpy() - if ndarr.shape[2] == 1: - ndarr = ndarr[:, :, 0] - mode = 'L' - elif ndarr.shape[2] == 3: - mode = 'RGB' - else: - mode = 'RGBA' - image = Image.fromarray(ndarr, mode=mode) - except Exception as e: - image = Image.new('RGB', (tensor.shape[-1], tensor.shape[-2]), color=(152, 32, 48)) - fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access - log.error(f'Convert: source={type(tensor)} target={image} fn={fn} {e}') - return image - - -def pil_to_tensor(image): - """PIL Image -> uint8 CHW tensor (no float scaling). Replaces TF.pil_to_tensor.""" - if not isinstance(image, Image.Image): - raise TypeError(f"Expected PIL Image, got {type(image)}") - pic = np.array(image, copy=True) - if pic.ndim == 2: - pic = pic[:, :, np.newaxis] - return torch.from_numpy(pic.transpose((2, 0, 1))).contiguous() - - -def normalize(tensor, mean, std, inplace=False): - """Tensor normalization. Replaces TF.normalize.""" - if not inplace: - tensor = tensor.clone() - mean_t = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device) - std_t = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device) - if mean_t.ndim == 1: - mean_t = mean_t[:, None, None] - if std_t.ndim == 1: - std_t = std_t[:, None, None] - tensor.sub_(mean_t).div_(std_t) - return tensor diff --git a/modules/image/util.py b/modules/image/util.py new file mode 100644 index 000000000..6d7b03ce5 --- /dev/null +++ b/modules/image/util.py @@ -0,0 +1,20 @@ +from PIL import Image, ImageDraw +from modules.image.grid import get_font +from modules import shared + + +def draw_text(im, text: str = '', y_offset: int = 0): + d = ImageDraw.Draw(im) + fontsize = (im.width + im.height) // 50 + font = get_font(fontsize) + d.text((fontsize//2, fontsize//2 + y_offset), text, font=font, fill=shared.opts.font_color) + return im + + +def flatten(img, bgcolor): + """replaces transparency with bgcolor (example: "#ffffff"), returning an RGB mode image with no transparency""" + if img.mode == "RGBA": + background = Image.new('RGBA', img.size, bgcolor) + background.paste(img, mask=img) + img = background + return img.convert('RGB') diff --git a/modules/image/watermark.py b/modules/image/watermark.py new file mode 100644 index 000000000..016b2d524 --- /dev/null +++ b/modules/image/watermark.py @@ -0,0 +1,81 @@ +import random +import numpy as np +from PIL import Image +from modules import shared + + +def set_watermark(image, wm_text: str | None = None, wm_image: Image.Image | None = None): + if shared.opts.image_watermark_position != 'none' and wm_image is not None: # visible watermark + if isinstance(wm_image, str): + try: + wm_image = Image.open(wm_image) + except Exception as e: + shared.log.warning(f'Set image watermark: image={wm_image} {e}') + return image + if isinstance(wm_image, Image.Image): + if wm_image.mode != 'RGBA': + wm_image = wm_image.convert('RGBA') + if shared.opts.image_watermark_position == 'top/left': + position = (0, 0) + elif shared.opts.image_watermark_position == 'top/right': + position = (image.width - wm_image.width, 0) + elif shared.opts.image_watermark_position == 'bottom/left': + position = (0, image.height - wm_image.height) + elif shared.opts.image_watermark_position == 'bottom/right': + position = (image.width - wm_image.width, image.height - wm_image.height) + elif shared.opts.image_watermark_position == 'center': + position = ((image.width - wm_image.width) // 2, (image.height - wm_image.height) // 2) + else: + position = (random.randint(0, image.width - wm_image.width), random.randint(0, image.height - wm_image.height)) + try: + for x in range(wm_image.width): + for y in range(wm_image.height): + rgba = wm_image.getpixel((x, y)) + orig = image.getpixel((x+position[0], y+position[1])) + # alpha blend + a = rgba[3] / 255 + r = int(rgba[0] * a + orig[0] * (1 - a)) + g = int(rgba[1] * a + orig[1] * (1 - a)) + b = int(rgba[2] * a + orig[2] * (1 - a)) + if not a == 0: + image.putpixel((x+position[0], y+position[1]), (r, g, b)) + shared.log.debug(f'Set image watermark: image={wm_image} position={position}') + except Exception as e: + shared.log.warning(f'Set image watermark: image={wm_image} {e}') + + if shared.opts.image_watermark_enabled and wm_text is not None: # invisible watermark + from imwatermark import WatermarkEncoder + wm_type = 'bytes' + wm_method = 'dwtDctSvd' + wm_length = 32 + length = wm_length // 8 + info = image.info + data = np.asarray(image) + encoder = WatermarkEncoder() + text = f"{wm_text:<{length}}"[:length] + bytearr = text.encode(encoding='ascii', errors='ignore') + try: + encoder.set_watermark(wm_type, bytearr) + encoded = encoder.encode(data, wm_method) + image = Image.fromarray(encoded) + image.info = info + shared.log.debug(f'Set invisible watermark: {wm_text} method={wm_method} bits={wm_length}') + except Exception as e: + shared.log.warning(f'Set invisible watermark error: {wm_text} method={wm_method} bits={wm_length} {e}') + + return image + + +def get_watermark(image): + from imwatermark import WatermarkDecoder + wm_type = 'bytes' + wm_method = 'dwtDctSvd' + wm_length = 32 + data = np.asarray(image) + decoder = WatermarkDecoder(wm_type, wm_length) + try: + decoded = decoder.decode(data, wm_method) + wm = decoded.decode(encoding='ascii', errors='ignore') + except Exception: + wm = '' + return wm diff --git a/modules/images.py b/modules/images.py index 5e4c5c9dd..070beb9a4 100644 --- a/modules/images.py +++ b/modules/images.py @@ -1,498 +1,10 @@ -import io -import re -import os -import sys -import json -import queue -import random -import datetime -import threading -import numpy as np -import piexif -import piexif.helper -from PIL import Image, PngImagePlugin, ExifTags, ImageDraw -from modules import sd_samplers, shared, script_callbacks, errors, paths -from modules.images_grid import image_grid, get_grid_size, split_grid, combine_grid, check_grid_size, get_font, draw_grid_annotations, draw_prompt_matrix, GridAnnotation, Grid # pylint: disable=unused-import -from modules.images_resize import resize_image # pylint: disable=unused-import -from modules.images_namegen import FilenameGenerator, get_next_sequence_number # pylint: disable=unused-import +from modules.image.util import flatten, draw_text # pylint: disable=unused-import +from modules.image.save import save_image # pylint: disable=unused-import +from modules.image.convert import to_pil, to_tensor # pylint: disable=unused-import +from modules.image.metadata import read_info_from_image, image_data # pylint: disable=unused-import +from modules.image.resize import resize_image # pylint: disable=unused-import +from modules.image.sharpfin import resize # pylint: disable=unused-import +from modules.image.namegen import FilenameGenerator, get_next_sequence_number # pylint: disable=unused-import +from modules.image.watermark import set_watermark, get_watermark # pylint: disable=unused-import +from modules.image.grid import image_grid, get_grid_size, split_grid, combine_grid, check_grid_size, get_font, draw_grid_annotations, draw_prompt_matrix, GridAnnotation, Grid # pylint: disable=unused-import from modules.video import save_video # pylint: disable=unused-import - - -debug = errors.log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None -debug_save = errors.log.trace if os.environ.get('SD_SAVE_DEBUG', None) is not None else lambda *args, **kwargs: None -try: - from pi_heif import register_heif_opener - register_heif_opener() -except Exception: - pass - - -def sanitize_filename_part(text, replace_spaces=True): - if text is None: - return None - if replace_spaces: - text = text.replace(' ', '_') - invalid_filename_chars = '#<>:"/\\|?*\n\r\t' - invalid_filename_prefix = ' ' - invalid_filename_postfix = ' .' - max_filename_part_length = 64 - text = text.translate({ord(x): '_' for x in invalid_filename_chars}) - text = text.lstrip(invalid_filename_prefix)[:max_filename_part_length] - text = text.rstrip(invalid_filename_postfix) - return text - - -def atomically_save_image(): - Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes - while True: - image, filename, extension, params, exifinfo, filename_txt, is_grid = save_queue.get() - jobid = shared.state.begin('Save image') - shared.state.image_history += 1 - if len(exifinfo) > 2: - with open(paths.params_path, "w", encoding="utf8") as file: - file.write(exifinfo) - fn = filename + extension - filename = filename.strip() - if extension[0] != '.': # add dot if missing - extension = '.' + extension - try: - image_format = Image.registered_extensions()[extension] - except Exception: - shared.log.warning(f'Save: unknown image format: {extension}') - image_format = 'JPEG' - exifinfo = (exifinfo or "") if shared.opts.image_metadata else "" - # additional metadata saved in files - if shared.opts.save_txt and len(exifinfo) > 0: - try: - with open(filename_txt, "w", encoding="utf8") as file: - file.write(f"{exifinfo}\n") - shared.log.info(f'Save: text="{filename_txt}" len={len(exifinfo)}') - except Exception as e: - shared.log.warning(f'Save failed: description={filename_txt} {e}') - - # actual save - if image_format == 'PNG': - pnginfo_data = PngImagePlugin.PngInfo() - for k, v in params.pnginfo.items(): - pnginfo_data.add_text(k, str(v)) - debug_save(f'Save pnginfo: {params.pnginfo.items()}') - save_args = { 'compress_level': 6, 'pnginfo': pnginfo_data if shared.opts.image_metadata else None } - elif image_format == 'JPEG': - if image.mode == 'RGBA': - shared.log.warning('Save: removing alpha channel') - image = image.convert("RGB") - elif image.mode == 'I;16': - image = image.point(lambda p: p * 0.0038910505836576).convert("L") - save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality } - if shared.opts.image_metadata: - debug_save(f'Save exif: {exifinfo}') - save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) - elif image_format == 'WEBP': - if image.mode == 'I;16': - image = image.point(lambda p: p * 0.0038910505836576).convert("RGB") - save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless } - if shared.opts.image_metadata: - debug_save(f'Save exif: {exifinfo}') - save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) - elif image_format == 'JXL': - if image.mode == 'I;16': - image = image.point(lambda p: p * 0.0038910505836576).convert("RGB") - elif image.mode not in {"RGB", "RGBA"}: - image = image.convert("RGBA") - save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless } - if shared.opts.image_metadata: - debug_save(f'Save exif: {exifinfo}') - save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) - else: - save_args = { 'quality': shared.opts.jpeg_quality } - try: - debug_save(f'Save args: {save_args}') - image.save(fn, format=image_format, **save_args) - except Exception as e: - shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}') - errors.display(e, 'Image save') - size = os.path.getsize(fn) if os.path.exists(fn) else 0 - what = 'grid' if is_grid else 'image' - shared.log.info(f'Save: {what}="{fn}" type={image_format} width={image.width} height={image.height} size={size}') - - if shared.opts.save_log_fn != '' and len(exifinfo) > 0: - fn = os.path.join(paths.data_path, shared.opts.save_log_fn) - if not fn.endswith('.json'): - fn += '.json' - entries = shared.readfile(fn, silent=True) - if not isinstance(entries, list): - entries = [] - idx = len(entries) - entry = { 'id': idx, 'filename': filename, 'time': datetime.datetime.now().isoformat(), 'info': exifinfo } - entries.append(entry) - shared.writefile(entries, fn, mode='w', silent=True) - shared.log.info(f'Save: json="{fn}" records={len(entries)}') - shared.state.outputs(filename) - shared.state.end(jobid) - save_queue.task_done() - - -save_queue: queue.Queue[tuple[Image.Image, str, str, script_callbacks.ImageSaveParams, str, str | None, bool]] = queue.Queue() -save_thread = threading.Thread(target=atomically_save_image, daemon=True) -save_thread.start() - - -def save_image(image, - path=None, - basename='', - seed=None, - prompt=None, - extension=shared.opts.samples_format, - info=None, - grid=False, - pnginfo_section_name='parameters', - p=None, - existing_info=None, - forced_filename=None, - suffix='', - save_to_dirs=None, - ): - fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access - debug_save(f'Save: fn={fn}') # pylint: disable=protected-access - if image is None: - shared.log.warning('Image is none') - return None, None, None - if isinstance(image, list): - if len(image) > 1: - shared.log.warning(f'Save: images={image} multiple images provided only the first one will be saved') - image = image[0] - if not check_grid_size([image]): - return None, None, None - if path is None or path == '': # set default path to avoid errors when functions are triggered manually or via api and param is not set - path = paths.resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_save) - namegen = FilenameGenerator(p, seed, prompt, image, grid=grid) - suffix = suffix if suffix is not None else '' - basename = '' if basename is None else basename - if save_to_dirs is not None and isinstance(save_to_dirs, str) and len(save_to_dirs) > 0: - dirname = save_to_dirs - path = os.path.join(path, dirname) - elif shared.opts.save_to_dirs: - dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]") - path = os.path.join(path, dirname) - if forced_filename is None: - if shared.opts.samples_filename_pattern and len(shared.opts.samples_filename_pattern) > 0: - file_decoration = shared.opts.samples_filename_pattern - else: - file_decoration = "[seq]-[prompt_words]" - file_decoration = namegen.apply(file_decoration) - file_decoration += suffix - if file_decoration.startswith(basename): - basename = '' - filename = os.path.join(path, f"{file_decoration}.{extension}") if basename == '' else os.path.join(path, f"{basename}-{file_decoration}.{extension}") - else: - forced_filename += suffix - if forced_filename.startswith(basename): - basename = '' - filename = os.path.join(path, f"{forced_filename}.{extension}") if basename == '' else os.path.join(path, f"{basename}-{forced_filename}.{extension}") - pnginfo = existing_info or {} - if info is None: - info = image.info.get(pnginfo_section_name, '') - if info is not None: - pnginfo[pnginfo_section_name] = info - - wm_text = getattr(p, 'watermark_text', shared.opts.image_watermark) - wm_image = getattr(p, 'watermark_image', shared.opts.image_watermark_image) - image = set_watermark(image, wm_text, wm_image) - - params = script_callbacks.ImageSaveParams(image, p, filename, pnginfo) - params.filename = namegen.sanitize(filename) - dirname = os.path.dirname(params.filename) - if dirname is not None and len(dirname) > 0: - os.makedirs(dirname, exist_ok=True) - params.filename = namegen.sequence(params.filename) - params.filename = namegen.sanitize(params.filename) - # callbacks - script_callbacks.before_image_saved_callback(params) - exifinfo = params.pnginfo.get('UserComment', '') - exifinfo = exifinfo + ', ' if len(exifinfo) > 0 else '' - exifinfo += params.pnginfo.get(pnginfo_section_name, '') - filename, extension = os.path.splitext(params.filename) - filename_txt = f"{filename}.txt" if shared.opts.save_txt and len(exifinfo) > 0 else None - save_queue.put((params.image, filename, extension, params, exifinfo, filename_txt, grid)) # actual save is executed in a thread that polls data from queue - save_queue.join() - if not hasattr(params.image, 'already_saved_as'): - debug(f'Image marked: "{params.filename}"') - params.image.already_saved_as = params.filename - script_callbacks.image_saved_callback(params) - return params.filename, filename_txt, exifinfo - - -def safe_decode_string(s: bytes): - remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment - for encoding in ['utf_16_be', 'utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings - try: - s = remove_prefix(s, b'UNICODE') - s = remove_prefix(s, b'ASCII') - s = remove_prefix(s, b'\x00') - val = s.decode(encoding, errors="strict") - val = re.sub(r'[\x00-\x09]', '', val).strip() # remove remaining special characters - if len(val) == 0: # remove empty strings - val = None - return val - except Exception: - pass - return None - - -def parse_comfy_metadata(data: dict): - def parse_workflow(): - res = '' - try: - txt = data.get('workflow', {}) - dct = json.loads(txt) - nodes = len(dct.get('nodes', [])) - version = dct.get('extra', {}).get('frontendVersion', 'unknown') - if version is not None: - res = f" | Version: {version} | Nodes: {nodes}" - except Exception: - pass - return res - - def parse_prompt(): - res = '' - try: - txt = data.get('prompt', {}) - dct = json.loads(txt) - for val in dct.values(): - inp = val.get('inputs', {}) - if 'model' in inp: - model = inp.get('model', None) - if isinstance(model, str) and len(model) > 0: - res += f" | Model: {model} | Class: {val.get('class_type', '')}" - except Exception: - pass - return res - - workflow = parse_workflow() - prompt = parse_prompt() - if len(workflow) > 0 or len(prompt) > 0: - parsed = f'App: ComfyUI{workflow}{prompt}' - shared.log.info(f'Image metadata: {parsed}') - return parsed - return '' - - -def parse_invoke_metadata(data: dict): - def parse_metadtaa(): - res = '' - try: - txt = data.get('invokeai_metadata', {}) - dct = json.loads(txt) - if 'app_version' in dct: - version = dct['app_version'] - if isinstance(version, str) and len(version) > 0: - res += f" | Version: {version}" - except Exception: - pass - return res - - metadata = parse_metadtaa() - if len(metadata) > 0: - parsed = f'App: InvokeAI{metadata}' - shared.log.info(f'Image metadata: {parsed}') - return parsed - return '' - - -def parse_novelai_metadata(data: dict): - geninfo = '' - if data.get("Software", None) == "NovelAI": - try: - dct = json.loads(data["Comment"]) - sampler = sd_samplers.samplers_map.get(dct["sampler"], "Euler a") - geninfo = f'{data["Description"]} Negative prompt: {dct["uc"]} Steps: {dct["steps"]}, Sampler: {sampler}, CFG scale: {dct["scale"]}, Seed: {dct["seed"]}, Clip skip: 2, ENSD: 31337' - except Exception: - pass - return geninfo - - -def read_info_from_image(image: Image.Image, watermark: bool = False) -> tuple[str, dict]: - if image is None: - return '', {} - if isinstance(image, str): - try: - image = Image.open(image) - image.load() - except Exception: - return '', {} - items = image.info or {} - geninfo = items.pop('parameters', None) or items.pop('UserComment', None) or '' - if isinstance(geninfo, dict): - if 'UserComment' in geninfo: - geninfo = geninfo['UserComment'] # Info was nested - else: - geninfo = '' # Unknown format. Ignore contents - items['UserComment'] = geninfo - - if "exif" in items: - try: - exif = piexif.load(items["exif"]) - except Exception as e: - shared.log.error(f'Error loading EXIF data: {e}') - exif = {} - for _key, subkey in exif.items(): - if isinstance(subkey, dict): - for key, val in subkey.items(): - if isinstance(val, bytes): # decode bytestring - val = safe_decode_string(val) - if isinstance(val, tuple) and isinstance(val[0], int) and isinstance(val[1], int) and val[1] > 0: # convert camera ratios - val = round(val[0] / val[1], 2) - if val is not None and key in ExifTags.TAGS: # add known tags - if ExifTags.TAGS[key] == 'UserComment': # add geninfo from UserComment - geninfo = str(val) - items['parameters'] = val - else: - items[ExifTags.TAGS[key]] = val - elif val is not None and key in ExifTags.GPSTAGS: - items[ExifTags.GPSTAGS[key]] = val - if watermark: - wm = get_watermark(image) - if wm != '': - # geninfo += f' Watermark: {wm}' - items['watermark'] = wm - - for key, val in items.items(): - if isinstance(val, bytes): # decode bytestring - items[key] = safe_decode_string(val) - - geninfo += parse_comfy_metadata(items) - geninfo += parse_invoke_metadata(items) - geninfo += parse_novelai_metadata(items) - - for key in ['exif', 'ExifOffset', 'JpegIFOffset', 'JpegIFByteCount', 'ExifVersion', 'icc_profile', 'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'adobe', 'photoshop', 'loop', 'duration', 'dpi']: # remove unwanted tags - items.pop(key, None) - - try: - items['width'] = image.width - items['height'] = image.height - items['mode'] = image.mode - except Exception: - pass - - return geninfo, items - - -def image_data(data): - import gradio as gr - if data is None: - return gr.update(), None - err1 = None - err2 = None - try: - image = Image.open(io.BytesIO(data)) - image.load() - info, _ = read_info_from_image(image) - errors.log.debug(f'Decoded object: image={image} metadata={info}') - return info, None - except Exception as e: - err1 = e - try: - if len(data) > 1024 * 10: - errors.log.warning(f'Error decoding object: data too long: {len(data)}') - return gr.update(), None - info = data.decode('utf8') - errors.log.debug(f'Decoded object: data={len(data)} metadata={info}') - return info, None - except Exception as e: - err2 = e - errors.log.error(f'Error decoding object: {err1 or err2}') - return gr.update(), None - - -def flatten(img, bgcolor): - """replaces transparency with bgcolor (example: "#ffffff"), returning an RGB mode image with no transparency""" - if img.mode == "RGBA": - background = Image.new('RGBA', img.size, bgcolor) - background.paste(img, mask=img) - img = background - return img.convert('RGB') - - -def draw_overlay(im, text: str = '', y_offset: int = 0): - d = ImageDraw.Draw(im) - fontsize = (im.width + im.height) // 50 - font = get_font(fontsize) - d.text((fontsize//2, fontsize//2 + y_offset), text, font=font, fill=shared.opts.font_color) - return im - - -def set_watermark(image, wm_text: str | None = None, wm_image: Image.Image | None = None): - if shared.opts.image_watermark_position != 'none' and wm_image is not None: # visible watermark - if isinstance(wm_image, str): - try: - wm_image = Image.open(wm_image) - except Exception as e: - shared.log.warning(f'Set image watermark: image={wm_image} {e}') - return image - if isinstance(wm_image, Image.Image): - if wm_image.mode != 'RGBA': - wm_image = wm_image.convert('RGBA') - if shared.opts.image_watermark_position == 'top/left': - position = (0, 0) - elif shared.opts.image_watermark_position == 'top/right': - position = (image.width - wm_image.width, 0) - elif shared.opts.image_watermark_position == 'bottom/left': - position = (0, image.height - wm_image.height) - elif shared.opts.image_watermark_position == 'bottom/right': - position = (image.width - wm_image.width, image.height - wm_image.height) - elif shared.opts.image_watermark_position == 'center': - position = ((image.width - wm_image.width) // 2, (image.height - wm_image.height) // 2) - else: - position = (random.randint(0, image.width - wm_image.width), random.randint(0, image.height - wm_image.height)) - try: - for x in range(wm_image.width): - for y in range(wm_image.height): - rgba = wm_image.getpixel((x, y)) - orig = image.getpixel((x+position[0], y+position[1])) - # alpha blend - a = rgba[3] / 255 - r = int(rgba[0] * a + orig[0] * (1 - a)) - g = int(rgba[1] * a + orig[1] * (1 - a)) - b = int(rgba[2] * a + orig[2] * (1 - a)) - if not a == 0: - image.putpixel((x+position[0], y+position[1]), (r, g, b)) - shared.log.debug(f'Set image watermark: image={wm_image} position={position}') - except Exception as e: - shared.log.warning(f'Set image watermark: image={wm_image} {e}') - - if shared.opts.image_watermark_enabled and wm_text is not None: # invisible watermark - from imwatermark import WatermarkEncoder - wm_type = 'bytes' - wm_method = 'dwtDctSvd' - wm_length = 32 - length = wm_length // 8 - info = image.info - data = np.asarray(image) - encoder = WatermarkEncoder() - text = f"{wm_text:<{length}}"[:length] - bytearr = text.encode(encoding='ascii', errors='ignore') - try: - encoder.set_watermark(wm_type, bytearr) - encoded = encoder.encode(data, wm_method) - image = Image.fromarray(encoded) - image.info = info - shared.log.debug(f'Set invisible watermark: {wm_text} method={wm_method} bits={wm_length}') - except Exception as e: - shared.log.warning(f'Set invisible watermark error: {wm_text} method={wm_method} bits={wm_length} {e}') - - return image - - -def get_watermark(image): - from imwatermark import WatermarkDecoder - wm_type = 'bytes' - wm_method = 'dwtDctSvd' - wm_length = 32 - data = np.asarray(image) - decoder = WatermarkDecoder(wm_type, wm_length) - try: - decoded = decoder.decode(data, wm_method) - wm = decoded.decode(encoding='ascii', errors='ignore') - except Exception: - wm = '' - return wm diff --git a/modules/lama.py b/modules/lama.py index 32029da85..36e27b24c 100644 --- a/modules/lama.py +++ b/modules/lama.py @@ -5,7 +5,8 @@ import torch import numpy as np from torch.hub import download_url_to_file, get_dir from PIL import Image -from modules import devices, images_sharpfin +from modules import devices +from modules.image import convert from installer import log @@ -96,5 +97,5 @@ class SimpleLama: image, mask = prepare_img_and_mask(image, mask, self.device) with devices.inference_context(): inpainted = self.model(image, mask) - cur_res = images_sharpfin.to_pil(inpainted[0]) + cur_res = convert.to_pil(inpainted[0]) return cur_res diff --git a/modules/postprocess/pixelart.py b/modules/postprocess/pixelart.py index f2b27e02b..3295a0a21 100644 --- a/modules/postprocess/pixelart.py +++ b/modules/postprocess/pixelart.py @@ -9,8 +9,8 @@ from diffusers.image_processor import PipelineImageInput from diffusers.configuration_utils import ConfigMixin, register_to_config from transformers import ImageProcessingMixin -from modules import images_sharpfin from modules import devices +from modules.image import sharpfin @devices.inference_context() @@ -64,9 +64,9 @@ def edge_detect_for_pixelart(image: PipelineImageInput, image_weight: float = 1. greyscale_reshaped = greyscale_reshaped.reshape(batch_size, block_size_sq, block_height, block_width) greyscale_range = greyscale_reshaped.amax(dim=1, keepdim=True).sub_(greyscale_reshaped.amin(dim=1, keepdim=True)) - range_weight = images_sharpfin.resize_tensor(greyscale_range, (height, width), linearize=False) + range_weight = sharpfin.resize_tensor(greyscale_range, (height, width), linearize=False) range_weight = range_weight.div_(range_weight.max()) - weight_map = images_sharpfin.resize_tensor((greyscale > greyscale.median()).to(dtype=torch.float32), (height, width), linearize=False) + weight_map = sharpfin.resize_tensor((greyscale > greyscale.median()).to(dtype=torch.float32), (height, width), linearize=False) weight_map = weight_map.unsqueeze(0).add_(range_weight).mul_(image_weight / 2) new_image = new_image.mul_(weight_map).addcmul_(min_pool, (1-weight_map)) @@ -158,7 +158,7 @@ def encode_jpeg_tensor(img: torch.FloatTensor, block_size: int=16, cbcr_downscal img = img[:, :, :(img.shape[-2]//block_size)*block_size, :(img.shape[-1]//block_size)*block_size] # crop to a multiply of block_size cbcr_block_size = block_size//cbcr_downscale _, _, height, width = img.shape - down_img = images_sharpfin.resize_tensor(img[:, 1:,:,:], (height//cbcr_downscale, width//cbcr_downscale), linearize=False) + down_img = sharpfin.resize_tensor(img[:, 1:,:,:], (height//cbcr_downscale, width//cbcr_downscale), linearize=False) y = encode_single_channel_dct_2d(img[:, 0, :,:], block_size=block_size, norm=norm) cb = encode_single_channel_dct_2d(down_img[:, 0, :,:], block_size=cbcr_block_size, norm=norm) cr = encode_single_channel_dct_2d(down_img[:, 1, :,:], block_size=cbcr_block_size, norm=norm) @@ -176,8 +176,8 @@ def decode_jpeg_tensor(jpeg_img: torch.FloatTensor, block_size: int=16, cbcr_dow y = decode_single_channel_dct_2d(y, norm=norm) cb = decode_single_channel_dct_2d(cb, norm=norm) cr = decode_single_channel_dct_2d(cr, norm=norm) - cb = images_sharpfin.resize_tensor(cb, (h_blocks*block_size, w_blocks*block_size), linearize=False) - cr = images_sharpfin.resize_tensor(cr, (h_blocks*block_size, w_blocks*block_size), linearize=False) + cb = sharpfin.resize_tensor(cb, (h_blocks*block_size, w_blocks*block_size), linearize=False) + cr = sharpfin.resize_tensor(cr, (h_blocks*block_size, w_blocks*block_size), linearize=False) return torch.stack([y,cb,cr], dim=1) diff --git a/modules/postprocess/seedvr_model.py b/modules/postprocess/seedvr_model.py index b78c174bc..6492f0bb2 100644 --- a/modules/postprocess/seedvr_model.py +++ b/modules/postprocess/seedvr_model.py @@ -3,9 +3,10 @@ import random import numpy as np import torch from PIL import Image -from modules import devices, images_sharpfin +from modules import devices from modules.shared import opts, log from modules.upscaler import Upscaler, UpscalerData +from modules.image import convert MODELS_MAP = { @@ -13,7 +14,6 @@ MODELS_MAP = { "SeedVR2 7B": "seedvr2_ema_7b_fp16.safetensors", "SeedVR2 7B Sharp": "seedvr2_ema_7b_sharp_fp16.safetensors", } -to_pil = images_sharpfin.to_pil class UpscalerSeedVR(Upscaler): @@ -158,7 +158,7 @@ class UpscalerSeedVR(Upscaler): ) t1 = time.time() log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={opts.seedvt_cfg_scale} seed={seed} time={t1 - t0:.2f}') - img = to_pil(result_tensor.squeeze()) + img = convert.to_pil(result_tensor.squeeze()) if opts.upscaler_unload: self.model.dit = None diff --git a/modules/processing_class.py b/modules/processing_class.py index a6f9fa5dd..26d942efa 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -8,6 +8,7 @@ import numpy as np from PIL import Image, ImageOps from modules import shared, images, scripts_manager, masking, sd_models, sd_vae, processing_helpers from modules.paths import resolve_output_path +from modules.image.util import flatten debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -538,7 +539,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): self.init_img_height = getattr(self, 'init_img_height', img.height) # pylint: disable=attribute-defined-outside-init if shared.opts.save_init_img: images.save_image(img, path=resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_init_images), basename=None, forced_filename=self.init_img_hash, suffix="-init-image") - image = images.flatten(img, shared.opts.img2img_background_color) + image = flatten(img, shared.opts.img2img_background_color) if crop_region is None and self.resize_mode > 0: image = images.resize_image(self.resize_mode, image, self.width, self.height, upscaler_name=self.resize_name, context=self.resize_context) self.width = image.width diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index e559f3ce0..7f0525780 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -4,11 +4,12 @@ import time import numpy as np import torch from PIL import Image -from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks, attention, images_sharpfin +from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks, attention from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name from modules.processing_args import set_pipeline_args from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed from modules.lora import lora_common +from modules.image import convert debug = os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None @@ -269,9 +270,9 @@ def process_hires(p: processing.StableDiffusionProcessing, output): sd_hijack_hypertile.hypertile_set(p, hr=True) elif torch.is_tensor(output.images) and output.images.shape[-1] == 3: # nhwc if output.images.dim() == 3: - output.images = images_sharpfin.to_pil(output.images) + output.images = convert.to_pil(output.images) elif output.images.dim() == 4: - output.images = [images_sharpfin.to_pil(output.images[i]) for i in range(output.images.shape[0])] + output.images = [convert.to_pil(output.images[i]) for i in range(output.images.shape[0])] strength = p.hr_denoising_strength if p.hr_denoising_strength > 0 else p.denoising_strength if (p.hr_upscaler is not None) and (p.hr_upscaler.lower().startswith('latent') or p.hr_force) and strength > 0: @@ -571,7 +572,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): if hasattr(shared.sd_model, 'unet') and hasattr(shared.sd_model.unet, 'config') and hasattr(shared.sd_model.unet.config, 'in_channels') and shared.sd_model.unet.config.in_channels == 9 and not is_control: shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING) # force pipeline if len(getattr(p, 'init_images', [])) == 0: - p.init_images = [images_sharpfin.to_pil(torch.rand((3, getattr(p, 'height', 512), getattr(p, 'width', 512))))] + p.init_images = [convert.to_pil(torch.rand((3, getattr(p, 'height', 512), getattr(p, 'width', 512))))] if not p.prompts: p.prompts = p.all_prompts[p.iteration * p.batch_size:(p.iteration+1) * p.batch_size] if not p.negative_prompts: diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 77da7ed57..53f8c3751 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -334,7 +334,7 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable jobid = shared.state.begin('VAE Encode') - from modules import images_sharpfin + from modules.image import convert if shared.state.interrupted or shared.state.skipped: return [] if not hasattr(model, 'vae') and hasattr(model, 'pipe'): @@ -342,7 +342,7 @@ def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable if not hasattr(model, 'vae'): shared.log.error('VAE not found in model') return [] - tensor = images_sharpfin.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae) + tensor = convert.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae) if vae_type == 'Tiny': latents = taesd_vae_encode(image=tensor) elif vae_type == 'Full' and hasattr(model, 'vae'): diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 6e7358be8..ee58c4057 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -3,8 +3,9 @@ import threading from collections import namedtuple import torch from PIL import Image -from modules import shared, devices, processing, images, sd_samplers, timer, images_sharpfin +from modules import shared, devices, processing, images, sd_samplers, timer from modules.vae import sd_vae_approx, sd_vae_taesd, sd_vae_stablecascade +from modules.image import convert SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options']) @@ -83,7 +84,7 @@ def single_sample_to_image(sample, approximation=None): x_sample = (255.0 * x_sample).to(torch.uint8) if len(x_sample.shape) == 4: x_sample = x_sample[0] - image = images_sharpfin.to_pil(x_sample) + image = convert.to_pil(x_sample) except Exception as e: warn_once(f'Preview: {e}') image = Image.new(mode="RGB", size=(512, 512)) diff --git a/modules/upscaler.py b/modules/upscaler.py index 6a7d5e275..da7c3e6bf 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -109,8 +109,8 @@ class Upscaler: if img.width >= dest_w and img.height >= dest_h: break if img.width != dest_w or img.height != dest_h: - from modules import images_sharpfin - img = images_sharpfin.resize(img, (int(dest_w), int(dest_h))) + from modules.image import sharpfin + img = sharpfin.resize(img, (int(dest_w), int(dest_h))) shared.state.end(jobid) return img diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py index 5e5e606f3..fbef4be6c 100644 --- a/modules/upscaler_simple.py +++ b/modules/upscaler_simple.py @@ -47,11 +47,11 @@ class UpscalerResize(Upscaler): elif selected_model == "Resize Box": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BOX) elif selected_model == "Resize Sharpfin MKS2021": - from modules import images_sharpfin - return images_sharpfin.resize(img, (int(img.width * self.scale), int(img.height * self.scale)), kernel="Sharpfin MKS2021") + from modules.image import sharpfin + return sharpfin.resize(img, (int(img.width * self.scale), int(img.height * self.scale)), kernel="Sharpfin MKS2021") elif selected_model == "Resize Sharpfin Lanczos3": - from modules import images_sharpfin - return images_sharpfin.resize(img, (int(img.width * self.scale), int(img.height * self.scale)), kernel="Sharpfin Lanczos3") + from modules.image import sharpfin + return sharpfin.resize(img, (int(img.width * self.scale), int(img.height * self.scale)), kernel="Sharpfin Lanczos3") else: return img diff --git a/modules/upscaler_spandrel.py b/modules/upscaler_spandrel.py index 20e69d558..e8f95ec1c 100644 --- a/modules/upscaler_spandrel.py +++ b/modules/upscaler_spandrel.py @@ -25,15 +25,15 @@ class UpscalerSpandrel(Upscaler): self.scalers.append(scaler) def process(self, img: Image.Image) -> Image.Image: - from modules import images_sharpfin - tensor = images_sharpfin.to_tensor(img).unsqueeze(0).to(devices.device) + from modules.image import convert + tensor = convert.to_tensor(img).unsqueeze(0).to(devices.device) img = img.convert('RGB') t0 = time.time() with devices.inference_context(): tensor = self.model(tensor) tensor = tensor.clamp(0, 1).squeeze(0).cpu() t1 = time.time() - upscaled = images_sharpfin.to_pil(tensor) + upscaled = convert.to_pil(tensor) log.debug(f'Upscale: name="{self.selected}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}') return upscaled diff --git a/modules/upscaler_vae.py b/modules/upscaler_vae.py index fb86702a1..871118c42 100644 --- a/modules/upscaler_vae.py +++ b/modules/upscaler_vae.py @@ -18,7 +18,8 @@ class UpscalerAsymmetricVAE(Upscaler): if selected_model is None: return img import diffusers - from modules import shared, devices, images_sharpfin + from modules import shared, devices + from modules.image import sharpfin, convert if self.vae is None or (selected_model != self.selected): if 'v1' in selected_model: repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler' @@ -31,11 +32,11 @@ class UpscalerAsymmetricVAE(Upscaler): self.selected = selected_model shared.log.debug(f'Upscaler load: selected="{self.selected}" vae="{repo_id}"') t0 = time.time() - img = images_sharpfin.resize(img, (8 * (img.width // 8), 8 * (img.height // 8))).convert('RGB') - tensor = images_sharpfin.to_tensor(img).unsqueeze(0).to(device=devices.device, dtype=devices.dtype) + img = sharpfin.resize(img, (8 * (img.width // 8), 8 * (img.height // 8))).convert('RGB') + tensor = convert.to_tensor(img).unsqueeze(0).to(device=devices.device, dtype=devices.dtype) self.vae = self.vae.to(device=devices.device) tensor = self.vae(tensor).sample - upscaled = images_sharpfin.to_pil(tensor.squeeze().clamp(0.0, 1.0).float().cpu()) + upscaled = convert.to_pil(tensor.squeeze().clamp(0.0, 1.0).float().cpu()) self.vae = self.vae.to(device=devices.cpu) t1 = time.time() shared.log.debug(f'Upscale: name="{self.selected}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}') @@ -58,7 +59,8 @@ class UpscalerWanUpscale(Upscaler): return img import torch.nn.functional as FN import diffusers - from modules import shared, devices, images_sharpfin + from modules import shared, devices + from modules.image import convert if (self.vae_encode is None) or (self.vae_decode is None) or (selected_model != self.selected): repo_encode = 'Qwen/Qwen-Image-Edit-2509' subfolder_encode = 'vae' @@ -77,7 +79,7 @@ class UpscalerWanUpscale(Upscaler): t0 = time.time() self.vae_encode = self.vae_encode.to(device=devices.device) - tensor = images_sharpfin.to_tensor(img).unsqueeze(0).unsqueeze(2).to(device=devices.device, dtype=devices.dtype) + tensor = convert.to_tensor(img).unsqueeze(0).unsqueeze(2).to(device=devices.device, dtype=devices.dtype) tensor = self.vae_encode.encode(tensor).latent_dist.mode() self.vae_encode.to(device=devices.cpu) @@ -86,7 +88,7 @@ class UpscalerWanUpscale(Upscaler): tensor = FN.pixel_shuffle(tensor.movedim(2, 1), upscale_factor=2).movedim(1, 2) # pixel shuffle needs [..., C, H, W] format self.vae_decode.to(device=devices.cpu) - upscaled = images_sharpfin.to_pil(tensor.squeeze().clamp(0.0, 1.0).float().cpu()) + upscaled = convert.to_pil(tensor.squeeze().clamp(0.0, 1.0).float().cpu()) t1 = time.time() shared.log.debug(f'Upscale: name="{self.selected}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}') return upscaled diff --git a/modules/video.py b/modules/video.py index b95ae7aa1..17d51eb38 100644 --- a/modules/video.py +++ b/modules/video.py @@ -3,7 +3,7 @@ import threading import numpy as np from PIL import Image from modules import shared, errors -from modules.images_namegen import FilenameGenerator # pylint: disable=unused-import +from modules.image.namegen import FilenameGenerator # pylint: disable=unused-import def interpolate_frames(images, count: int = 0, scale: float = 1.0, pad: int = 1, change: float = 0.3): diff --git a/modules/video_models/video_save.py b/modules/video_models/video_save.py index ac4b9b725..e137ca966 100644 --- a/modules/video_models/video_save.py +++ b/modules/video_models/video_save.py @@ -11,7 +11,7 @@ from modules.video_models.video_utils import check_av def get_video_filename(p:processing.StableDiffusionProcessingVideo): - from modules.images_namegen import FilenameGenerator + from modules.image.namegen import FilenameGenerator namegen = FilenameGenerator(p, seed=p.seed if p is not None else 0, prompt=p.prompt if p is not None else '') filename = namegen.apply(shared.opts.samples_filename_pattern if shared.opts.samples_filename_pattern and len(shared.opts.samples_filename_pattern) > 0 else "[seq]-[prompt_words]") if shared.opts.save_to_dirs: diff --git a/pipelines/f_lite/pipeline.py b/pipelines/f_lite/pipeline.py index 3cd8b5b09..a760c9617 100644 --- a/pipelines/f_lite/pipeline.py +++ b/pipelines/f_lite/pipeline.py @@ -293,9 +293,9 @@ class FLitePipeline(DiffusionPipeline): raise # 8. Post-process images - from modules import images_sharpfin + from modules.image import convert images = (decoded_images / 2 + 0.5).clamp(0, 1) - pil_images = [images_sharpfin.to_pil(img) for img in images] + pil_images = [convert.to_pil(img) for img in images] return FLitePipelineOutput( images=pil_images, diff --git a/pipelines/model_stablecascade.py b/pipelines/model_stablecascade.py index 5a64e4ae0..7f1df501b 100644 --- a/pipelines/model_stablecascade.py +++ b/pipelines/model_stablecascade.py @@ -332,8 +332,8 @@ class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline): if output_type == "np": images = images.permute(0, 2, 3, 1).cpu().float().numpy() # float() as bfloat16-> numpy doesnt work elif output_type == "pil": - from modules import images_sharpfin - images = [images_sharpfin.to_pil(images[i]) for i in range(images.shape[0])] + from modules.image import convert + images = [convert.to_pil(images[i]) for i in range(images.shape[0])] shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) else: images = latents diff --git a/scripts/ctrlx/media.py b/scripts/ctrlx/media.py index dea5a0b76..b50778630 100644 --- a/scripts/ctrlx/media.py +++ b/scripts/ctrlx/media.py @@ -1,7 +1,7 @@ import numpy as np import torch -from modules import images_sharpfin import PIL +from modules.image import convert JPEG_QUALITY = 95 @@ -13,7 +13,7 @@ def preprocess(image, processor, **kwargs): elif isinstance(image, np.ndarray): image = PIL.Image.fromarray(image) elif isinstance(image, torch.Tensor): - image = images_sharpfin.to_pil(image) + image = convert.to_pil(image) else: raise TypeError(f"Image must be of type PIL.Image, np.ndarray, or torch.Tensor, got {type(image)} instead.") diff --git a/scripts/differential_diffusion.py b/scripts/differential_diffusion.py index fd772a600..d5bebb1ed 100644 --- a/scripts/differential_diffusion.py +++ b/scripts/differential_diffusion.py @@ -858,8 +858,8 @@ class StableDiffusionXLDiffImg2ImgPipeline(DiffusionPipeline, FromSingleFileMixi # 4. Preprocess image #image = self.image_processor.preprocess(image) #ideally we would have preprocess the image with diffusers, but for this POC we won't --- it throws a deprecated warning - from modules import images_sharpfin - map = images_sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in original_image.shape[2:]), linearize=False) + from modules.image import sharpfin + map = sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in original_image.shape[2:]), linearize=False) # 5. Prepare timesteps def denoising_value_valid(dnv): return type(denoising_end) == float and 0 < dnv < 1 @@ -1758,8 +1758,8 @@ class StableDiffusionDiffImg2ImgPipeline(DiffusionPipeline): # 7. Prepare extra step kwargs. extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - from modules import images_sharpfin - map = images_sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in image.shape[2:]), linearize=False) + from modules.image import sharpfin + map = sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in image.shape[2:]), linearize=False) # 8. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order @@ -1834,7 +1834,8 @@ class StableDiffusionDiffImg2ImgPipeline(DiffusionPipeline): import gradio as gr import diffusers from PIL import Image, ImageEnhance, ImageOps # pylint: disable=reimported -from modules import errors, shared, devices, scripts_manager, processing, sd_models, images, images_sharpfin +from modules import errors, shared, devices, scripts_manager, processing, sd_models, images +from modules.image import convert detector = None @@ -1888,9 +1889,9 @@ class Script(scripts_manager.Script): else: return None, None, None image_mask = image_map.copy() - image_map = images_sharpfin.to_tensor(image_map) + image_map = convert.to_tensor(image_map) image_map = image_map.to(devices.device) - image_init = 2 * images_sharpfin.to_tensor(image_init) - 1 + image_init = 2 * convert.to_tensor(image_init) - 1 image_init = image_init.unsqueeze(0) image_init = image_init.to(devices.device) return image_init, image_map, image_mask diff --git a/scripts/lbm_ext.py b/scripts/lbm_ext.py index c0f3f7428..9a5b09d04 100644 --- a/scripts/lbm_ext.py +++ b/scripts/lbm_ext.py @@ -1,7 +1,7 @@ from copy import deepcopy from PIL import Image import gradio as gr -from modules import scripts_manager, processing, shared, devices, sd_models +from modules import scripts_manager, processing, shared, devices birefnet = None @@ -84,8 +84,8 @@ class Script(scripts_manager.Script): from installer import install install('lpips') - from modules import images_sharpfin - from scripts.lbm import get_model, extract_object, resize_and_center_crop # pylint: disable=no-name-in-module + from modules.image import convert + from scripts.lbm import extract_object, resize_and_center_crop # pylint: disable=no-name-in-module ori_h_bg, ori_w_bg = fg_image.size ar_bg = ori_h_bg / ori_w_bg @@ -110,7 +110,7 @@ class Script(scripts_manager.Script): if lbm_method == 'Simple': output_image = img_pasted else: - img_pasted_tensor = images_sharpfin.to_tensor(img_pasted).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) * 2 - 1 + img_pasted_tensor = convert.to_tensor(img_pasted).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) * 2 - 1 batch = { "source_image": img_pasted_tensor } z_source = model.vae.encode(batch[model.source_key]) output_image = model.sample( @@ -120,7 +120,7 @@ class Script(scripts_manager.Script): max_samples=1, ) output_image = (output_image[0].clamp(-1, 1).float().cpu() + 1) / 2 - output_image = images_sharpfin.to_pil(output_image) + output_image = convert.to_pil(output_image) if lbm_composite: output_image = Image.composite(output_image, bg_image, fg_mask) diff --git a/scripts/pixelsmith_ext.py b/scripts/pixelsmith_ext.py index a87570d0f..55d0b8d80 100644 --- a/scripts/pixelsmith_ext.py +++ b/scripts/pixelsmith_ext.py @@ -26,13 +26,13 @@ class Script(scripts_manager.Script): def encode(self, p: processing.StableDiffusionProcessing, image: Image.Image): if image is None: return None - from modules import images_sharpfin + from modules.image import convert if p.width is None or p.width == 0: p.width = int(8 * (image.width * p.scale_by // 8)) if p.height is None or p.height == 0: p.height = int(8 * (image.height * p.scale_by // 8)) image = images.resize_image(p.resize_mode, image, p.width, p.height, upscaler_name=p.resize_name, context=p.resize_context) - tensor = images_sharpfin.to_tensor(image).unsqueeze(0).to(device=devices.device, dtype=devices.dtype) + tensor = convert.to_tensor(image).unsqueeze(0).to(device=devices.device, dtype=devices.dtype) tensor = 2.0 * tensor - 1.0 with devices.inference_context(): latent = shared.sd_model.vae.tiled_encode(tensor) diff --git a/scripts/sd_upscale.py b/scripts/sd_upscale.py index 1e1e97229..ee944f54a 100644 --- a/scripts/sd_upscale.py +++ b/scripts/sd_upscale.py @@ -4,6 +4,7 @@ from PIL import Image from modules import processing, shared, images, devices, scripts_manager from modules.processing import get_processed from modules.shared import opts, state, log +from modules.image.util import flatten class Script(scripts_manager.Script): @@ -32,7 +33,7 @@ class Script(scripts_manager.Script): if init_img is None: return None - init_img = images.flatten(init_img, opts.img2img_background_color) + init_img = flatten(init_img, opts.img2img_background_color) if isinstance(upscaler_index, str): upscaler_index = [x.name.lower() for x in shared.sd_upscalers].index(upscaler_index.lower()) diff --git a/scripts/softfill.py b/scripts/softfill.py index 51a286dc7..e4c84db31 100644 --- a/scripts/softfill.py +++ b/scripts/softfill.py @@ -1322,8 +1322,8 @@ class StableDiffusionXLSoftFillPipeline( image.save("noised_image.png") image = transforms.CenterCrop((image.size[1] // 64 * 64, image.size[0] // 64 * 64))(image) - from modules import images_sharpfin - image = images_sharpfin.to_tensor(image) + from modules.image import convert + image = convert.to_tensor(image) image = image * 2 - 1 # Normalize to [-1, 1] return image.unsqueeze(0) @@ -1334,8 +1334,8 @@ class StableDiffusionXLSoftFillPipeline( """ map = map.convert("L") map = transforms.CenterCrop((map.size[1] // 64 * 64, map.size[0] // 64 * 64))(map) - from modules import images_sharpfin - map = images_sharpfin.to_tensor(map) + from modules.image import convert + map = convert.to_tensor(map) map = (map - 0.05) / (0.95 - 0.05) map = torch.clamp(map, 0.0, 1.0) return 1.0 - map @@ -1350,8 +1350,8 @@ class StableDiffusionXLSoftFillPipeline( # Prepare mask as rescaled tensor map map = preprocess_map(mask).to(device) - from modules import images_sharpfin - map = images_sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in original_image_tensor.shape[2:]), linearize=False) + from modules.image import sharpfin + map = sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in original_image_tensor.shape[2:]), linearize=False) # Generate latent tensor with noise original_with_noise = self.prepare_latents( diff --git a/scripts/xyz/xyz_grid_draw.py b/scripts/xyz/xyz_grid_draw.py index 646a3bd52..2338d089d 100644 --- a/scripts/xyz/xyz_grid_draw.py +++ b/scripts/xyz/xyz_grid_draw.py @@ -2,6 +2,7 @@ import time from copy import copy from PIL import Image from modules import shared, images, processing +from modules.image.util import draw_text def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, no_grid: False, include_time: False, include_text: False): # pylint: disable=unused-argument @@ -50,7 +51,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend if include_time: overlay_text += f'Time: {elapsed:.2f}' if len(overlay_text) > 0: - processed_result.images[idx] = images.draw_overlay(processed_result.images[idx], overlay_text) + processed_result.images[idx] = draw_text(processed_result.images[idx], overlay_text) processed_result.all_prompts[idx] = processed.prompt processed_result.all_seeds[idx] = processed.seed processed_result.infotexts[idx] = processed.infotexts[0]