diff --git a/CHANGES.md b/CHANGES.md index 2e45c279..b53ac141 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -22,6 +22,15 @@ Our focus continues to remain on an easy installation experience, and an easy user-interface. While still remaining pretty powerful, in terms of features and speed. ### Detailed changelog +* 2.5.41 - 13 Jun 2023 - Fix multi-gpu bug with "low" VRAM usage mode while generating images. +* 2.5.41 - 12 Jun 2023 - Fix multi-gpu bug with CodeFormer. +* 2.5.41 - 6 Jun 2023 - Allow changing the strength of CodeFormer, and slightly improved styling of the CodeFormer options. +* 2.5.41 - 5 Jun 2023 - Allow sharing an Easy Diffusion instance via https://try.cloudflare.com/ . You can find this option at the bottom of the Settings tab. Thanks @JeLuf. +* 2.5.41 - 5 Jun 2023 - Show an option to download for tiled images. Shows a button on the generated image. Creates larger images by tiling them with the image generated by Easy Diffusion. Thanks @JeLuf. +* 2.5.41 - 5 Jun 2023 - (beta-only) Allow LoRA strengths between -2 and 2. Thanks @ogmaresca. +* 2.5.40 - 5 Jun 2023 - Reduce the VRAM usage of Latent Upscaling when using "balanced" VRAM usage mode. +* 2.5.40 - 5 Jun 2023 - Fix the "realesrgan" key error when using CodeFormer with more than 1 image in a batch. +* 2.5.40 - 3 Jun 2023 - Added CodeFormer as another option for fixing faces and eyes. CodeFormer tends to perform better than GFPGAN for many images. Thanks @patriceac for the implementation, and for contacting the CodeFormer team (who were supportive of it being integrated into Easy Diffusion). * 2.5.39 - 25 May 2023 - (beta-only) Seamless Tiling - make seamlessly tiled images, e.g. rock and grass textures. Thanks @JeLuf. * 2.5.38 - 24 May 2023 - Better reporting of errors, and show an explanation if the user cannot disable the "Use CPU" setting. * 2.5.38 - 23 May 2023 - Add Latent Upscaler as another option for upscaling images. Thanks @JeLuf for the implementation of the Latent Upscaler model. diff --git a/scripts/check_models.py b/scripts/check_models.py deleted file mode 100644 index 4b8d68c6..00000000 --- a/scripts/check_models.py +++ /dev/null @@ -1,101 +0,0 @@ -# this script runs inside the legacy "stable-diffusion" folder - -from sdkit.models import download_model, get_model_info_from_db -from sdkit.utils import hash_file_quick - -import os -import shutil -from glob import glob -import traceback - -models_base_dir = os.path.abspath(os.path.join("..", "models")) - -models_to_check = { - "stable-diffusion": [ - {"file_name": "sd-v1-4.ckpt", "model_id": "1.4"}, - ], - "gfpgan": [ - {"file_name": "GFPGANv1.4.pth", "model_id": "1.4"}, - ], - "realesrgan": [ - {"file_name": "RealESRGAN_x4plus.pth", "model_id": "x4plus"}, - {"file_name": "RealESRGAN_x4plus_anime_6B.pth", "model_id": "x4plus_anime_6"}, - ], - "vae": [ - {"file_name": "vae-ft-mse-840000-ema-pruned.ckpt", "model_id": "vae-ft-mse-840000-ema-pruned"}, - ], -} -MODEL_EXTENSIONS = { # copied from easydiffusion/model_manager.py - "stable-diffusion": [".ckpt", ".safetensors"], - "vae": [".vae.pt", ".ckpt", ".safetensors"], - "hypernetwork": [".pt", ".safetensors"], - "gfpgan": [".pth"], - "realesrgan": [".pth"], - "lora": [".ckpt", ".safetensors"], -} - - -def download_if_necessary(model_type: str, file_name: str, model_id: str): - model_path = os.path.join(models_base_dir, model_type, file_name) - expected_hash = get_model_info_from_db(model_type=model_type, model_id=model_id)["quick_hash"] - - other_models_exist = any_model_exists(model_type) - known_model_exists = os.path.exists(model_path) - known_model_is_corrupt = known_model_exists and hash_file_quick(model_path) != expected_hash - - if known_model_is_corrupt or (not other_models_exist and not known_model_exists): - print("> download", model_type, model_id) - download_model(model_type, model_id, download_base_dir=models_base_dir) - - -def init(): - migrate_legacy_model_location() - - for model_type, models in models_to_check.items(): - for model in models: - try: - download_if_necessary(model_type, model["file_name"], model["model_id"]) - except: - traceback.print_exc() - fail(model_type) - - print(model_type, "model(s) found.") - - -### utilities -def any_model_exists(model_type: str) -> bool: - extensions = MODEL_EXTENSIONS.get(model_type, []) - for ext in extensions: - if any(glob(f"{models_base_dir}/{model_type}/**/*{ext}", recursive=True)): - return True - - return False - - -def migrate_legacy_model_location(): - 'Move the models inside the legacy "stable-diffusion" folder, to their respective folders' - - for model_type, models in models_to_check.items(): - for model in models: - file_name = model["file_name"] - if os.path.exists(file_name): - dest_dir = os.path.join(models_base_dir, model_type) - os.makedirs(dest_dir, exist_ok=True) - shutil.move(file_name, os.path.join(dest_dir, file_name)) - - -def fail(model_name): - print( - f"""Error downloading the {model_name} model. Sorry about that, please try to: -1. Run this installer again. -2. If that doesn't fix it, please try to download the file manually. The address to download from, and the destination to save to are printed above this message. -3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB -4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues -Thanks!""" - ) - exit(1) - - -### start - -init() diff --git a/scripts/check_modules.py b/scripts/check_modules.py index 3686ca00..6275de45 100644 --- a/scripts/check_modules.py +++ b/scripts/check_modules.py @@ -18,13 +18,15 @@ os_name = platform.system() modules_to_check = { "torch": ("1.11.0", "1.13.1", "2.0.0"), "torchvision": ("0.12.0", "0.14.1", "0.15.1"), - "sdkit": "1.0.98", + "sdkit": "1.0.106", "stable-diffusion-sdkit": "2.1.4", "rich": "12.6.0", "uvicorn": "0.19.0", "fastapi": "0.85.1", + "pycloudflared": "0.2.0", # "xformers": "0.0.16", } +modules_to_log = ["torch", "torchvision", "sdkit", "stable-diffusion-sdkit"] def version(module_name: str) -> str: @@ -89,7 +91,8 @@ def init(): traceback.print_exc() fail(module_name) - print(f"{module_name}: {version(module_name)}") + if module_name in modules_to_log: + print(f"{module_name}: {version(module_name)}") ### utilities diff --git a/scripts/developer_console.sh b/scripts/developer_console.sh index 73972568..57846eeb 100755 --- a/scripts/developer_console.sh +++ b/scripts/developer_console.sh @@ -39,6 +39,8 @@ if [ "$0" == "bash" ]; then export PYTHONPATH="$(pwd)/stable-diffusion/env/lib/python3.8/site-packages" fi + export PYTHONNOUSERSITE=y + which python python --version diff --git a/scripts/on_env_start.bat b/scripts/on_env_start.bat index 44144cfa..bc92d0e9 100644 --- a/scripts/on_env_start.bat +++ b/scripts/on_env_start.bat @@ -67,7 +67,6 @@ if "%update_branch%"=="" ( @xcopy sd-ui-files\ui ui /s /i /Y /q @copy sd-ui-files\scripts\on_sd_start.bat scripts\ /Y @copy sd-ui-files\scripts\check_modules.py scripts\ /Y -@copy sd-ui-files\scripts\check_models.py scripts\ /Y @copy sd-ui-files\scripts\get_config.py scripts\ /Y @copy "sd-ui-files\scripts\Start Stable Diffusion UI.cmd" . /Y @copy "sd-ui-files\scripts\Developer Console.cmd" . /Y diff --git a/scripts/on_env_start.sh b/scripts/on_env_start.sh index 30465975..366b5dd1 100755 --- a/scripts/on_env_start.sh +++ b/scripts/on_env_start.sh @@ -50,7 +50,6 @@ cp -Rf sd-ui-files/ui . cp sd-ui-files/scripts/on_sd_start.sh scripts/ cp sd-ui-files/scripts/bootstrap.sh scripts/ cp sd-ui-files/scripts/check_modules.py scripts/ -cp sd-ui-files/scripts/check_models.py scripts/ cp sd-ui-files/scripts/get_config.py scripts/ cp sd-ui-files/scripts/start.sh . cp sd-ui-files/scripts/developer_console.sh . diff --git a/scripts/on_sd_start.bat b/scripts/on_sd_start.bat index ba205c9e..f92b9f6f 100644 --- a/scripts/on_sd_start.bat +++ b/scripts/on_sd_start.bat @@ -5,7 +5,6 @@ @copy sd-ui-files\scripts\on_env_start.bat scripts\ /Y @copy sd-ui-files\scripts\check_modules.py scripts\ /Y -@copy sd-ui-files\scripts\check_models.py scripts\ /Y @copy sd-ui-files\scripts\get_config.py scripts\ /Y if exist "%cd%\profile" ( @@ -79,13 +78,6 @@ call WHERE uvicorn > .tmp @echo conda_sd_ui_deps_installed >> ..\scripts\install_status.txt ) -@rem Download the required models -call python ..\scripts\check_models.py -if "%ERRORLEVEL%" NEQ "0" ( - pause - exit /b -) - @>nul findstr /m "sd_install_complete" ..\scripts\install_status.txt @if "%ERRORLEVEL%" NEQ "0" ( @echo sd_weights_downloaded >> ..\scripts\install_status.txt diff --git a/scripts/on_sd_start.sh b/scripts/on_sd_start.sh index 820c36ed..be5161d4 100755 --- a/scripts/on_sd_start.sh +++ b/scripts/on_sd_start.sh @@ -4,7 +4,6 @@ cp sd-ui-files/scripts/functions.sh scripts/ cp sd-ui-files/scripts/on_env_start.sh scripts/ cp sd-ui-files/scripts/bootstrap.sh scripts/ cp sd-ui-files/scripts/check_modules.py scripts/ -cp sd-ui-files/scripts/check_models.py scripts/ cp sd-ui-files/scripts/get_config.py scripts/ source ./scripts/functions.sh @@ -51,12 +50,6 @@ if ! command -v uvicorn &> /dev/null; then fail "UI packages not found!" fi -# Download the required models -if ! python ../scripts/check_models.py; then - read -p "Press any key to continue" - exit 1 -fi - if [ `grep -c sd_install_complete ../scripts/install_status.txt` -gt "0" ]; then echo sd_weights_downloaded >> ../scripts/install_status.txt echo sd_install_complete >> ../scripts/install_status.txt diff --git a/ui/easydiffusion/app.py b/ui/easydiffusion/app.py index 3064e151..38e3392c 100644 --- a/ui/easydiffusion/app.py +++ b/ui/easydiffusion/app.py @@ -90,8 +90,8 @@ def init(): os.makedirs(USER_SERVER_PLUGINS_DIR, exist_ok=True) # https://pytorch.org/docs/stable/storage.html - warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated') - + warnings.filterwarnings("ignore", category=UserWarning, message="TypedStorage is deprecated") + load_server_plugins() update_render_threads() @@ -221,12 +221,41 @@ def open_browser(): webbrowser.open(f"http://localhost:{port}") - Console().print(Panel( - "\n" + - "[white]Easy Diffusion is ready to serve requests.\n\n" + - "A new browser tab should have been opened by now.\n" + - f"If not, please open your web browser and navigate to [bold yellow underline]http://localhost:{port}/\n", - title="Easy Diffusion is ready", style="bold yellow on blue")) + Console().print( + Panel( + "\n" + + "[white]Easy Diffusion is ready to serve requests.\n\n" + + "A new browser tab should have been opened by now.\n" + + f"If not, please open your web browser and navigate to [bold yellow underline]http://localhost:{port}/\n", + title="Easy Diffusion is ready", + style="bold yellow on blue", + ) + ) + + +def fail_and_die(fail_type: str, data: str): + suggestions = [ + "Run this installer again.", + "If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB", + "If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues", + ] + + if fail_type == "model_download": + fail_label = f"Error downloading the {data} model" + suggestions.insert( + 1, + "If that doesn't fix it, please try to download the file manually. The address to download from, and the destination to save to are printed above this message.", + ) + else: + fail_label = "Error while installing Easy Diffusion" + + msg = [f"{fail_label}. Sorry about that, please try to:"] + for i, suggestion in enumerate(suggestions): + msg.append(f"{i+1}. {suggestion}") + msg.append("Thanks!") + + print("\n".join(msg)) + exit(1) def get_image_modifiers(): diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py index 0a1f1b5c..de2c10ac 100644 --- a/ui/easydiffusion/model_manager.py +++ b/ui/easydiffusion/model_manager.py @@ -1,10 +1,14 @@ import os +import shutil +from glob import glob +import traceback from easydiffusion import app from easydiffusion.types import TaskData from easydiffusion.utils import log from sdkit import Context -from sdkit.models import load_model, scan_model, unload_model +from sdkit.models import load_model, scan_model, unload_model, download_model, get_model_info_from_db +from sdkit.utils import hash_file_quick KNOWN_MODEL_TYPES = [ "stable-diffusion", @@ -13,6 +17,7 @@ KNOWN_MODEL_TYPES = [ "gfpgan", "realesrgan", "lora", + "codeformer", ] MODEL_EXTENSIONS = { "stable-diffusion": [".ckpt", ".safetensors"], @@ -21,14 +26,22 @@ MODEL_EXTENSIONS = { "gfpgan": [".pth"], "realesrgan": [".pth"], "lora": [".ckpt", ".safetensors"], + "codeformer": [".pth"], } DEFAULT_MODELS = { - "stable-diffusion": [ # needed to support the legacy installations - "custom-model", # only one custom model file was supported initially, creatively named 'custom-model' - "sd-v1-4", # Default fallback. + "stable-diffusion": [ + {"file_name": "sd-v1-4.ckpt", "model_id": "1.4"}, + ], + "gfpgan": [ + {"file_name": "GFPGANv1.4.pth", "model_id": "1.4"}, + ], + "realesrgan": [ + {"file_name": "RealESRGAN_x4plus.pth", "model_id": "x4plus"}, + {"file_name": "RealESRGAN_x4plus_anime_6B.pth", "model_id": "x4plus_anime_6"}, + ], + "vae": [ + {"file_name": "vae-ft-mse-840000-ema-pruned.ckpt", "model_id": "vae-ft-mse-840000-ema-pruned"}, ], - "gfpgan": ["GFPGANv1.3"], - "realesrgan": ["RealESRGAN_x4plus"], } MODELS_TO_LOAD_ON_START = ["stable-diffusion", "vae", "hypernetwork", "lora"] @@ -37,6 +50,8 @@ known_models = {} def init(): make_model_folders() + migrate_legacy_model_location() # if necessary + download_default_models_if_necessary() getModels() # run this once, to cache the picklescan results @@ -45,7 +60,7 @@ def load_default_models(context: Context): # init default model paths for model_type in MODELS_TO_LOAD_ON_START: - context.model_paths[model_type] = resolve_model_to_use(model_type=model_type) + context.model_paths[model_type] = resolve_model_to_use(model_type=model_type, fail_if_not_found=False) try: load_model( context, @@ -57,7 +72,12 @@ def load_default_models(context: Context): del context.model_load_errors[model_type] except Exception as e: log.error(f"[red]Error while loading {model_type} model: {context.model_paths[model_type]}[/red]") - log.exception(e) + if "DefaultCPUAllocator: not enough memory" in str(e): + log.error( + f"[red]Your PC is low on system RAM. Please add some virtual memory (or swap space) by following the instructions at this link: https://www.ibm.com/docs/en/opw/8.2.0?topic=tuning-optional-increasing-paging-file-size-windows-computers[/red]" + ) + else: + log.exception(e) del context.model_paths[model_type] context.model_load_errors[model_type] = str(e) # storing the entire Exception can lead to memory leaks @@ -70,12 +90,12 @@ def unload_all(context: Context): del context.model_load_errors[model_type] -def resolve_model_to_use(model_name: str = None, model_type: str = None): +def resolve_model_to_use(model_name: str = None, model_type: str = None, fail_if_not_found: bool = True): model_extensions = MODEL_EXTENSIONS.get(model_type, []) default_models = DEFAULT_MODELS.get(model_type, []) config = app.getConfig() - model_dirs = [os.path.join(app.MODELS_DIR, model_type), app.SD_DIR] + model_dir = os.path.join(app.MODELS_DIR, model_type) if not model_name: # When None try user configured model. # config = getConfig() if "model" in config and model_type in config["model"]: @@ -83,45 +103,42 @@ def resolve_model_to_use(model_name: str = None, model_type: str = None): if model_name: # Check models directory - models_dir_path = os.path.join(app.MODELS_DIR, model_type, model_name) + model_path = os.path.join(model_dir, model_name) + if os.path.exists(model_path): + return model_path for model_extension in model_extensions: - if os.path.exists(models_dir_path + model_extension): - return models_dir_path + model_extension + if os.path.exists(model_path + model_extension): + return model_path + model_extension if os.path.exists(model_name + model_extension): return os.path.abspath(model_name + model_extension) - # Default locations - if model_name in default_models: - default_model_path = os.path.join(app.SD_DIR, model_name) - for model_extension in model_extensions: - if os.path.exists(default_model_path + model_extension): - return default_model_path + model_extension - # Can't find requested model, check the default paths. - for default_model in default_models: - for model_dir in model_dirs: - default_model_path = os.path.join(model_dir, default_model) - for model_extension in model_extensions: - if os.path.exists(default_model_path + model_extension): - if model_name is not None: - log.warn( - f"Could not find the configured custom model {model_name}{model_extension}. Using the default one: {default_model_path}{model_extension}" - ) - return default_model_path + model_extension + if model_type == "stable-diffusion" and not fail_if_not_found: + for default_model in default_models: + default_model_path = os.path.join(model_dir, default_model["file_name"]) + if os.path.exists(default_model_path): + if model_name is not None: + log.warn( + f"Could not find the configured custom model {model_name}. Using the default one: {default_model_path}" + ) + return default_model_path - return None + if model_name and fail_if_not_found: + raise Exception(f"Could not find the desired model {model_name}! Is it present in the {model_dir} folder?") def reload_models_if_necessary(context: Context, task_data: TaskData): - use_upscale_lower = task_data.use_upscale.lower() if task_data.use_upscale else "" + face_fix_lower = task_data.use_face_correction.lower() if task_data.use_face_correction else "" + upscale_lower = task_data.use_upscale.lower() if task_data.use_upscale else "" model_paths_in_req = { "stable-diffusion": task_data.use_stable_diffusion_model, "vae": task_data.use_vae_model, "hypernetwork": task_data.use_hypernetwork_model, - "gfpgan": task_data.use_face_correction, - "realesrgan": task_data.use_upscale if "realesrgan" in use_upscale_lower else None, - "latent_upscaler": True if task_data.use_upscale == "latent_upscaler" else None, + "codeformer": task_data.use_face_correction if "codeformer" in face_fix_lower else None, + "gfpgan": task_data.use_face_correction if "gfpgan" in face_fix_lower else None, + "realesrgan": task_data.use_upscale if "realesrgan" in upscale_lower else None, + "latent_upscaler": True if "latent_upscaler" in upscale_lower else None, "nsfw_checker": True if task_data.block_nsfw else None, "lora": task_data.use_lora_model, } @@ -131,6 +148,13 @@ def reload_models_if_necessary(context: Context, task_data: TaskData): if context.model_paths.get(model_type) != path } + if task_data.codeformer_upscale_faces: + if "realesrgan" not in models_to_reload and "realesrgan" not in context.models: + default_realesrgan = DEFAULT_MODELS["realesrgan"][0]["file_name"] + models_to_reload["realesrgan"] = resolve_model_to_use(default_realesrgan, "realesrgan") + elif "realesrgan" in models_to_reload and models_to_reload["realesrgan"] is None: + del models_to_reload["realesrgan"] # don't unload realesrgan + if set_vram_optimizations(context) or set_clip_skip(context, task_data): # reload SD models_to_reload["stable-diffusion"] = model_paths_in_req["stable-diffusion"] @@ -157,7 +181,13 @@ def resolve_model_paths(task_data: TaskData): task_data.use_lora_model = resolve_model_to_use(task_data.use_lora_model, model_type="lora") if task_data.use_face_correction: - task_data.use_face_correction = resolve_model_to_use(task_data.use_face_correction, "gfpgan") + if "gfpgan" in task_data.use_face_correction.lower(): + model_type = "gfpgan" + elif "codeformer" in task_data.use_face_correction.lower(): + model_type = "codeformer" + download_if_necessary("codeformer", "codeformer.pth", "codeformer-0.1.0") + + task_data.use_face_correction = resolve_model_to_use(task_data.use_face_correction, model_type) if task_data.use_upscale and "realesrgan" in task_data.use_upscale.lower(): task_data.use_upscale = resolve_model_to_use(task_data.use_upscale, "realesrgan") @@ -167,7 +197,31 @@ def fail_if_models_did_not_load(context: Context): if model_type in context.model_load_errors: e = context.model_load_errors[model_type] raise Exception(f"Could not load the {model_type} model! Reason: " + e) - # concat 'e', don't use in format string (injection attack) + + +def download_default_models_if_necessary(): + for model_type, models in DEFAULT_MODELS.items(): + for model in models: + try: + download_if_necessary(model_type, model["file_name"], model["model_id"]) + except: + traceback.print_exc() + app.fail_and_die(fail_type="model_download", data=model_type) + + print(model_type, "model(s) found.") + + +def download_if_necessary(model_type: str, file_name: str, model_id: str): + model_path = os.path.join(app.MODELS_DIR, model_type, file_name) + expected_hash = get_model_info_from_db(model_type=model_type, model_id=model_id)["quick_hash"] + + other_models_exist = any_model_exists(model_type) + known_model_exists = os.path.exists(model_path) + known_model_is_corrupt = known_model_exists and hash_file_quick(model_path) != expected_hash + + if known_model_is_corrupt or (not other_models_exist and not known_model_exists): + print("> download", model_type, model_id) + download_model(model_type, model_id, download_base_dir=app.MODELS_DIR) def set_vram_optimizations(context: Context): @@ -181,6 +235,26 @@ def set_vram_optimizations(context: Context): return False +def migrate_legacy_model_location(): + 'Move the models inside the legacy "stable-diffusion" folder, to their respective folders' + + for model_type, models in DEFAULT_MODELS.items(): + for model in models: + file_name = model["file_name"] + legacy_path = os.path.join(app.SD_DIR, file_name) + if os.path.exists(legacy_path): + shutil.move(legacy_path, os.path.join(app.MODELS_DIR, model_type, file_name)) + + +def any_model_exists(model_type: str) -> bool: + extensions = MODEL_EXTENSIONS.get(model_type, []) + for ext in extensions: + if any(glob(f"{app.MODELS_DIR}/{model_type}/**/*{ext}", recursive=True)): + return True + + return False + + def set_clip_skip(context: Context, task_data: TaskData): clip_skip = task_data.clip_skip @@ -238,17 +312,12 @@ def is_malicious_model(file_path): def getModels(): models = { - "active": { - "stable-diffusion": "sd-v1-4", - "vae": "", - "hypernetwork": "", - "lora": "", - }, "options": { "stable-diffusion": ["sd-v1-4"], "vae": [], "hypernetwork": [], "lora": [], + "codeformer": ["codeformer"], }, } @@ -309,9 +378,4 @@ def getModels(): if models_scanned > 0: log.info(f"[green]Scanned {models_scanned} models. Nothing infected[/]") - # legacy - custom_weight_path = os.path.join(app.SD_DIR, "custom-model.ckpt") - if os.path.exists(custom_weight_path): - models["options"]["stable-diffusion"].append("custom-model") - return models diff --git a/ui/easydiffusion/renderer.py b/ui/easydiffusion/renderer.py index e2dae34f..a57dfc6c 100644 --- a/ui/easydiffusion/renderer.py +++ b/ui/easydiffusion/renderer.py @@ -7,10 +7,12 @@ from easydiffusion import device_manager from easydiffusion.types import GenerateImageRequest from easydiffusion.types import Image as ResponseImage from easydiffusion.types import Response, TaskData, UserInitiatedStop +from easydiffusion.model_manager import DEFAULT_MODELS, resolve_model_to_use from easydiffusion.utils import get_printable_request, log, save_images_to_disk from sdkit import Context from sdkit.filter import apply_filters from sdkit.generate import generate_images +from sdkit.models import load_model from sdkit.utils import ( diffusers_latent_samples_to_images, gc, @@ -34,6 +36,7 @@ def init(device): context.temp_images = {} context.partial_x_samples = None context.model_load_errors = {} + context.enable_codeformer = True from easydiffusion import app @@ -156,32 +159,51 @@ def filter_images(req: GenerateImageRequest, task_data: TaskData, images: list, if user_stopped: return images - filters_to_apply = [] - filter_params = {} if task_data.block_nsfw: - filters_to_apply.append("nsfw_checker") - if task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower(): - filters_to_apply.append("gfpgan") + images = apply_filters(context, "nsfw_checker", images) + + if task_data.use_face_correction and "codeformer" in task_data.use_face_correction.lower(): + default_realesrgan = DEFAULT_MODELS["realesrgan"][0]["file_name"] + prev_realesrgan_path = None + if task_data.codeformer_upscale_faces and default_realesrgan not in context.model_paths["realesrgan"]: + prev_realesrgan_path = context.model_paths["realesrgan"] + context.model_paths["realesrgan"] = resolve_model_to_use(default_realesrgan, "realesrgan") + load_model(context, "realesrgan") + + try: + images = apply_filters( + context, + "codeformer", + images, + upscale_faces=task_data.codeformer_upscale_faces, + codeformer_fidelity=task_data.codeformer_fidelity, + ) + finally: + if prev_realesrgan_path: + context.model_paths["realesrgan"] = prev_realesrgan_path + load_model(context, "realesrgan") + elif task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower(): + images = apply_filters(context, "gfpgan", images) + if task_data.use_upscale: if "realesrgan" in task_data.use_upscale.lower(): - filters_to_apply.append("realesrgan") + images = apply_filters(context, "realesrgan", images, scale=task_data.upscale_amount) elif task_data.use_upscale == "latent_upscaler": - filters_to_apply.append("latent_upscaler") + images = apply_filters( + context, + "latent_upscaler", + images, + scale=task_data.upscale_amount, + latent_upscaler_options={ + "prompt": req.prompt, + "negative_prompt": req.negative_prompt, + "seed": req.seed, + "num_inference_steps": task_data.latent_upscaler_steps, + "guidance_scale": 0, + }, + ) - filter_params["latent_upscaler_options"] = { - "prompt": req.prompt, - "negative_prompt": req.negative_prompt, - "seed": req.seed, - "num_inference_steps": task_data.latent_upscaler_steps, - "guidance_scale": 0, - } - - filter_params["scale"] = task_data.upscale_amount - - if len(filters_to_apply) == 0: - return images - - return apply_filters(context, filters_to_apply, images, **filter_params) + return images def construct_response(images: list, seeds: list, task_data: TaskData, base_seed: int): diff --git a/ui/easydiffusion/server.py b/ui/easydiffusion/server.py index a1aab6c0..d8940bb5 100644 --- a/ui/easydiffusion/server.py +++ b/ui/easydiffusion/server.py @@ -15,6 +15,7 @@ from fastapi import FastAPI, HTTPException from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, Extra from starlette.responses import FileResponse, JSONResponse, StreamingResponse +from pycloudflared import try_cloudflare log.info(f"started in {app.SD_DIR}") log.info(f"started at {datetime.datetime.now():%x %X}") @@ -113,6 +114,14 @@ def init(): def get_image(task_id: int, img_id: int): return get_image_internal(task_id, img_id) + @server_api.post("/tunnel/cloudflare/start") + def start_cloudflare_tunnel(req: dict): + return start_cloudflare_tunnel_internal(req) + + @server_api.post("/tunnel/cloudflare/stop") + def stop_cloudflare_tunnel(req: dict): + return stop_cloudflare_tunnel_internal(req) + @server_api.get("/") def read_root(): return FileResponse(os.path.join(app.SD_UI_DIR, "index.html"), headers=NOCACHE_HEADERS) @@ -211,6 +220,8 @@ def ping_internal(session_id: str = None): session = task_manager.get_cached_session(session_id, update_ttl=True) response["tasks"] = {id(t): t.status for t in session.tasks} response["devices"] = task_manager.get_devices() + if cloudflare.address != None: + response["cloudflare"] = cloudflare.address return JSONResponse(response, headers=NOCACHE_HEADERS) @@ -322,3 +333,47 @@ def get_image_internal(task_id: int, img_id: int): return StreamingResponse(img_data, media_type="image/jpeg") except KeyError as e: raise HTTPException(status_code=500, detail=str(e)) + +#---- Cloudflare Tunnel ---- +class CloudflareTunnel: + def __init__(self): + config = app.getConfig() + self.urls = None + self.port = config.get("net", {}).get("listen_port") + + def start(self): + if self.port: + self.urls = try_cloudflare(self.port) + + def stop(self): + if self.urls: + try_cloudflare.terminate(self.port) + self.urls = None + + @property + def address(self): + if self.urls: + return self.urls.tunnel + else: + return None + +cloudflare = CloudflareTunnel() + +def start_cloudflare_tunnel_internal(req: dict): + try: + cloudflare.start() + log.info(f"- Started cloudflare tunnel. Using address: {cloudflare.address}") + return JSONResponse({"address":cloudflare.address}) + except Exception as e: + log.error(str(e)) + log.error(traceback.format_exc()) + return HTTPException(status_code=500, detail=str(e)) + +def stop_cloudflare_tunnel_internal(req: dict): + try: + cloudflare.stop() + except Exception as e: + log.error(str(e)) + log.error(traceback.format_exc()) + return HTTPException(status_code=500, detail=str(e)) + diff --git a/ui/easydiffusion/types.py b/ui/easydiffusion/types.py index e4426714..abf8db29 100644 --- a/ui/easydiffusion/types.py +++ b/ui/easydiffusion/types.py @@ -23,7 +23,7 @@ class GenerateImageRequest(BaseModel): sampler_name: str = None # "ddim", "plms", "heun", "euler", "euler_a", "dpm2", "dpm2_a", "lms" hypernetwork_strength: float = 0 lora_alpha: float = 0 - tiling: str = "none" # "none", "x", "y", "xy" + tiling: str = "none" # "none", "x", "y", "xy" class TaskData(BaseModel): @@ -51,6 +51,8 @@ class TaskData(BaseModel): stream_image_progress: bool = False stream_image_progress_interval: int = 5 clip_skip: bool = False + codeformer_upscale_faces: bool = False + codeformer_fidelity: float = 0.5 class MergeRequest(BaseModel): diff --git a/ui/index.html b/ui/index.html index 21ec2550..0c4386de 100644 --- a/ui/index.html +++ b/ui/index.html @@ -30,7 +30,7 @@