import mimetypes import transformers import json import os import traceback import eventlet import glob import shlex import math import shutil import sys sys.path.append(".") from argparse import ArgumentTypeError from modules.create_cmd_parser import create_cmd_parser parser = create_cmd_parser() opt = parser.parse_args() from flask_socketio import SocketIO from flask import Flask, send_from_directory, url_for, jsonify from pathlib import Path from PIL import Image from pytorch_lightning import logging from threading import Event from uuid import uuid4 from send2trash import send2trash from ldm.generate import Generate from ldm.dream.restoration import Restoration from ldm.dream.pngwriter import PngWriter, retrieve_metadata from ldm.dream.args import APP_ID, APP_VERSION, calculate_init_img_hash from ldm.dream.conditioning import split_weighted_subprompts from modules.parameters import parameters_to_command """ USER CONFIG """ if opt.cors and "*" in opt.cors: raise ArgumentTypeError('"*" is not an allowed CORS origin') output_dir = "outputs/" # Base output directory for images host = opt.host # Web & socket.io host port = opt.port # Web & socket.io port verbose = opt.verbose # enables copious socket.io logging precision = opt.precision free_gpu_mem = opt.free_gpu_mem embedding_path = opt.embedding_path additional_allowed_origins = ( opt.cors if opt.cors else [] ) # additional CORS allowed origins model = "stable-diffusion-1.4" """ END USER CONFIG """ print("* Initializing, be patient...\n") """ SERVER SETUP """ # fix missing mimetypes on windows due to registry wonkiness mimetypes.add_type("application/javascript", ".js") mimetypes.add_type("text/css", ".css") app = Flask(__name__, static_url_path="", static_folder="../frontend/dist/") app.config["OUTPUTS_FOLDER"] = "../outputs" @app.route("/outputs/") def outputs(filename): return send_from_directory(app.config["OUTPUTS_FOLDER"], filename) @app.route("/", defaults={"path": ""}) def serve(path): return send_from_directory(app.static_folder, "index.html") logger = True if verbose else False engineio_logger = True if verbose else False # default 1,000,000, needs to be higher for socketio to accept larger images max_http_buffer_size = 10000000 cors_allowed_origins = [f"http://{host}:{port}"] + additional_allowed_origins socketio = SocketIO( app, logger=logger, engineio_logger=engineio_logger, max_http_buffer_size=max_http_buffer_size, cors_allowed_origins=cors_allowed_origins, ping_interval=(50, 50), ping_timeout=60, ) """ END SERVER SETUP """ """ APP SETUP """ class CanceledException(Exception): pass try: gfpgan, codeformer, esrgan = None, None, None from ldm.dream.restoration.base import Restoration restoration = Restoration() gfpgan, codeformer = restoration.load_face_restore_models() esrgan = restoration.load_esrgan() # coreformer.process(self, image, strength, device, seed=None, fidelity=0.75) except (ModuleNotFoundError, ImportError): print(traceback.format_exc(), file=sys.stderr) print(">> You may need to install the ESRGAN and/or GFPGAN modules") canceled = Event() # reduce logging outputs to error transformers.logging.set_verbosity_error() logging.getLogger("pytorch_lightning").setLevel(logging.ERROR) # Initialize and load model generate = Generate( model, precision=precision, embedding_path=embedding_path, ) generate.free_gpu_mem = free_gpu_mem generate.load_model() # location for "finished" images result_path = os.path.join(output_dir, "img-samples/") # temporary path for intermediates intermediate_path = os.path.join(result_path, "intermediates/") # path for user-uploaded init images and masks init_image_path = os.path.join(result_path, "init-images/") mask_image_path = os.path.join(result_path, "mask-images/") # txt log log_path = os.path.join(result_path, "dream_log.txt") # make all output paths [ os.makedirs(path, exist_ok=True) for path in [result_path, intermediate_path, init_image_path, mask_image_path] ] """ END APP SETUP """ """ SOCKET.IO LISTENERS """ @socketio.on("requestSystemConfig") def handle_request_capabilities(): print(f">> System config requested") config = get_system_config() socketio.emit("systemConfig", config) @socketio.on("requestImages") def handle_request_images(page=1, offset=0, last_mtime=None): chunk_size = 50 if last_mtime: print(f">> Latest images requested") else: print( f">> Page {page} of images requested (page size {chunk_size} offset {offset})" ) paths = glob.glob(os.path.join(result_path, "*.png")) sorted_paths = sorted(paths, key=lambda x: os.path.getmtime(x), reverse=True) if last_mtime: image_paths = filter(lambda x: os.path.getmtime(x) > last_mtime, sorted_paths) else: image_paths = sorted_paths[ slice(chunk_size * (page - 1) + offset, chunk_size * page + offset) ] page = page + 1 image_array = [] for path in image_paths: metadata = retrieve_metadata(path) image_array.append( { "url": path, "mtime": os.path.getmtime(path), "metadata": metadata["sd-metadata"], } ) socketio.emit( "galleryImages", { "images": image_array, "nextPage": page, "offset": offset, "onlyNewImages": True if last_mtime else False, }, ) @socketio.on("generateImage") def handle_generate_image_event( generation_parameters, esrgan_parameters, gfpgan_parameters ): print( f">> Image generation requested: {generation_parameters}\nESRGAN parameters: {esrgan_parameters}\nGFPGAN parameters: {gfpgan_parameters}" ) generate_images(generation_parameters, esrgan_parameters, gfpgan_parameters) @socketio.on("runESRGAN") def handle_run_esrgan_event(original_image, esrgan_parameters): print( f'>> ESRGAN upscale requested for "{original_image["url"]}": {esrgan_parameters}' ) progress = { "currentStep": 1, "totalSteps": 1, "currentIteration": 1, "totalIterations": 1, "currentStatus": "Preparing", "isProcessing": True, "currentStatusHasSteps": False, } socketio.emit("progressUpdate", progress) eventlet.sleep(0) image = Image.open(original_image["url"]) seed = ( original_image["metadata"]["seed"] if "seed" in original_image["metadata"] else "unknown_seed" ) progress["currentStatus"] = "Upscaling" socketio.emit("progressUpdate", progress) eventlet.sleep(0) image = esrgan.process( image=image, upsampler_scale=esrgan_parameters["upscale"][0], strength=esrgan_parameters["upscale"][1], seed=seed, ) progress["currentStatus"] = "Saving image" socketio.emit("progressUpdate", progress) eventlet.sleep(0) esrgan_parameters["seed"] = seed metadata = parameters_to_post_processed_image_metadata( parameters=esrgan_parameters, original_image_path=original_image["url"], type="esrgan", ) command = parameters_to_command(esrgan_parameters) path = save_image(image, command, metadata, result_path, postprocessing="esrgan") write_log_message(f'[Upscaled] "{original_image["url"]}" > "{path}": {command}') progress["currentStatus"] = "Finished" progress["currentStep"] = 0 progress["totalSteps"] = 0 progress["currentIteration"] = 0 progress["totalIterations"] = 0 progress["isProcessing"] = False socketio.emit("progressUpdate", progress) eventlet.sleep(0) socketio.emit( "esrganResult", { "url": os.path.relpath(path), "mtime": os.path.getmtime(path), "metadata": metadata, }, ) @socketio.on("runGFPGAN") def handle_run_gfpgan_event(original_image, gfpgan_parameters): print( f'>> GFPGAN face fix requested for "{original_image["url"]}": {gfpgan_parameters}' ) progress = { "currentStep": 1, "totalSteps": 1, "currentIteration": 1, "totalIterations": 1, "currentStatus": "Preparing", "isProcessing": True, "currentStatusHasSteps": False, } socketio.emit("progressUpdate", progress) eventlet.sleep(0) image = Image.open(original_image["url"]) seed = ( original_image["metadata"]["seed"] if "seed" in original_image["metadata"] else "unknown_seed" ) progress["currentStatus"] = "Fixing faces" socketio.emit("progressUpdate", progress) eventlet.sleep(0) image = gfpgan.process( image=image, strength=gfpgan_parameters["gfpgan_strength"], seed=seed ) progress["currentStatus"] = "Saving image" socketio.emit("progressUpdate", progress) eventlet.sleep(0) gfpgan_parameters["seed"] = seed metadata = parameters_to_post_processed_image_metadata( parameters=gfpgan_parameters, original_image_path=original_image["url"], type="gfpgan", ) command = parameters_to_command(gfpgan_parameters) path = save_image(image, command, metadata, result_path, postprocessing="gfpgan") write_log_message(f'[Fixed faces] "{original_image["url"]}" > "{path}": {command}') progress["currentStatus"] = "Finished" progress["currentStep"] = 0 progress["totalSteps"] = 0 progress["currentIteration"] = 0 progress["totalIterations"] = 0 progress["isProcessing"] = False socketio.emit("progressUpdate", progress) eventlet.sleep(0) socketio.emit( "gfpganResult", { "url": os.path.relpath(path), "mtime": os.path.mtime(path), "metadata": metadata, }, ) @socketio.on("cancel") def handle_cancel(): print(f">> Cancel processing requested") canceled.set() socketio.emit("processingCanceled") # TODO: I think this needs a safety mechanism. @socketio.on("deleteImage") def handle_delete_image(path, uuid): print(f'>> Delete requested "{path}"') send2trash(path) socketio.emit("imageDeleted", {"url": path, "uuid": uuid}) # TODO: I think this needs a safety mechanism. @socketio.on("uploadInitialImage") def handle_upload_initial_image(bytes, name): print(f'>> Init image upload requested "{name}"') uuid = uuid4().hex split = os.path.splitext(name) name = f"{split[0]}.{uuid}{split[1]}" file_path = os.path.join(init_image_path, name) os.makedirs(os.path.dirname(file_path), exist_ok=True) newFile = open(file_path, "wb") newFile.write(bytes) socketio.emit("initialImageUploaded", {"url": file_path, "uuid": ""}) # TODO: I think this needs a safety mechanism. @socketio.on("uploadMaskImage") def handle_upload_mask_image(bytes, name): print(f'>> Mask image upload requested "{name}"') uuid = uuid4().hex split = os.path.splitext(name) name = f"{split[0]}.{uuid}{split[1]}" file_path = os.path.join(mask_image_path, name) os.makedirs(os.path.dirname(file_path), exist_ok=True) newFile = open(file_path, "wb") newFile.write(bytes) socketio.emit("maskImageUploaded", {"url": file_path, "uuid": ""}) """ END SOCKET.IO LISTENERS """ """ ADDITIONAL FUNCTIONS """ def get_system_config(): return { "model": "stable diffusion", "model_id": model, "model_hash": generate.model_hash, "app_id": APP_ID, "app_version": APP_VERSION, } def parameters_to_post_processed_image_metadata(parameters, original_image_path, type): # top-level metadata minus `image` or `images` metadata = get_system_config() orig_hash = calculate_init_img_hash(original_image_path) image = {"orig_path": original_image_path, "orig_hash": orig_hash} if type == "esrgan": image["type"] = "esrgan" image["scale"] = parameters["upscale"][0] image["strength"] = parameters["upscale"][1] elif type == "gfpgan": image["type"] = "gfpgan" image["strength"] = parameters["gfpgan_strength"] else: raise TypeError(f"Invalid type: {type}") metadata["image"] = image return metadata def parameters_to_generated_image_metadata(parameters): # top-level metadata minus `image` or `images` metadata = get_system_config() # remove any image keys not mentioned in RFC #266 rfc266_img_fields = [ "type", "postprocessing", "sampler", "prompt", "seed", "variations", "steps", "cfg_scale", "step_number", "width", "height", "extra", "seamless", ] rfc_dict = {} for item in parameters.items(): key, value = item if key in rfc266_img_fields: rfc_dict[key] = value postprocessing = [] # 'postprocessing' is either null or an if "gfpgan_strength" in parameters: postprocessing.append( {"type": "gfpgan", "strength": float(parameters["gfpgan_strength"])} ) if "upscale" in parameters: postprocessing.append( { "type": "esrgan", "scale": int(parameters["upscale"][0]), "strength": float(parameters["upscale"][1]), } ) rfc_dict["postprocessing"] = postprocessing if len(postprocessing) > 0 else None # semantic drift rfc_dict["sampler"] = parameters["sampler_name"] # display weighted subprompts (liable to change) subprompts = split_weighted_subprompts(parameters["prompt"]) subprompts = [{"prompt": x[0], "weight": x[1]} for x in subprompts] rfc_dict["prompt"] = subprompts # 'variations' should always exist and be an array, empty or consisting of {'seed': seed, 'weight': weight} pairs variations = [] if "with_variations" in parameters: variations = [ {"seed": x[0], "weight": x[1]} for x in parameters["with_variations"] ] rfc_dict["variations"] = variations if "init_img" in parameters: rfc_dict["type"] = "img2img" rfc_dict["strength"] = parameters["strength"] rfc_dict["fit"] = parameters["fit"] # TODO: Noncompliant rfc_dict["orig_hash"] = calculate_init_img_hash(parameters["init_img"]) rfc_dict["init_image_path"] = parameters["init_img"] # TODO: Noncompliant rfc_dict["sampler"] = "ddim" # TODO: FIX ME WHEN IMG2IMG SUPPORTS ALL SAMPLERS if "init_mask" in parameters: rfc_dict["mask_hash"] = calculate_init_img_hash( parameters["init_mask"] ) # TODO: Noncompliant rfc_dict["mask_image_path"] = parameters["init_mask"] # TODO: Noncompliant else: rfc_dict["type"] = "txt2img" metadata["image"] = rfc_dict return metadata def make_unique_init_image_filename(name): uuid = uuid4().hex split = os.path.splitext(name) name = f"{split[0]}.{uuid}{split[1]}" return name def write_log_message(message, log_path=log_path): """Logs the filename and parameters used to generate or process that image to log file""" message = f"{message}\n" with open(log_path, "a", encoding="utf-8") as file: file.writelines(message) def save_image( image, command, metadata, output_dir, step_index=None, postprocessing=False ): pngwriter = PngWriter(output_dir) prefix = pngwriter.unique_prefix() seed = "unknown_seed" if "image" in metadata: if "seed" in metadata["image"]: seed = metadata["image"]["seed"] filename = f"{prefix}.{seed}" if step_index: filename += f".{step_index}" if postprocessing: filename += f".postprocessed" filename += ".png" path = pngwriter.save_image_and_prompt_to_png( image=image, dream_prompt=command, metadata=metadata, name=filename ) return path def calculate_real_steps(steps, strength, has_init_image): return math.floor(strength * steps) if has_init_image else steps def generate_images(generation_parameters, esrgan_parameters, gfpgan_parameters): canceled.clear() step_index = 1 prior_variations = ( generation_parameters["with_variations"] if "with_variations" in generation_parameters else [] ) """ If a result image is used as an init image, and then deleted, we will want to be able to use it as an init image in the future. Need to copy it. If the init/mask image doesn't exist in the init_image_path/mask_image_path, make a unique filename for it and copy it there. """ if "init_img" in generation_parameters: filename = os.path.basename(generation_parameters["init_img"]) if not os.path.exists(os.path.join(init_image_path, filename)): unique_filename = make_unique_init_image_filename(filename) new_path = os.path.join(init_image_path, unique_filename) shutil.copy(generation_parameters["init_img"], new_path) generation_parameters["init_img"] = new_path if "init_mask" in generation_parameters: filename = os.path.basename(generation_parameters["init_mask"]) if not os.path.exists(os.path.join(mask_image_path, filename)): unique_filename = make_unique_init_image_filename(filename) new_path = os.path.join(init_image_path, unique_filename) shutil.copy(generation_parameters["init_img"], new_path) generation_parameters["init_mask"] = new_path totalSteps = calculate_real_steps( steps=generation_parameters["steps"], strength=generation_parameters["strength"] if "strength" in generation_parameters else None, has_init_image="init_img" in generation_parameters, ) progress = { "currentStep": 1, "totalSteps": totalSteps, "currentIteration": 1, "totalIterations": generation_parameters["iterations"], "currentStatus": "Preparing", "isProcessing": True, "currentStatusHasSteps": False, } socketio.emit("progressUpdate", progress) eventlet.sleep(0) def image_progress(sample, step): if canceled.is_set(): raise CanceledException nonlocal step_index nonlocal generation_parameters nonlocal progress progress["currentStep"] = step + 1 progress["currentStatus"] = "Generating" progress["currentStatusHasSteps"] = True if ( generation_parameters["progress_images"] and step % 5 == 0 and step < generation_parameters["steps"] - 1 ): image = generate.sample_to_image(sample) metadata = parameters_to_generated_image_metadata(generation_parameters) command = parameters_to_command(generation_parameters) path = save_image(image, command, metadata, intermediate_path, step_index=step_index, postprocessing=False) step_index += 1 socketio.emit( "intermediateResult", { "url": os.path.relpath(path), "mtime": os.path.getmtime(path), "metadata": metadata, }, ) socketio.emit("progressUpdate", progress) eventlet.sleep(0) def image_done(image, seed, first_seed): nonlocal generation_parameters nonlocal esrgan_parameters nonlocal gfpgan_parameters nonlocal progress step_index = 1 nonlocal prior_variations progress["currentStatus"] = "Generation complete" socketio.emit("progressUpdate", progress) eventlet.sleep(0) all_parameters = generation_parameters postprocessing = False if ( "variation_amount" in all_parameters and all_parameters["variation_amount"] > 0 ): first_seed = first_seed or seed this_variation = [[seed, all_parameters["variation_amount"]]] all_parameters["with_variations"] = prior_variations + this_variation all_parameters["seed"] = first_seed elif ("with_variations" in all_parameters): all_parameters["seed"] = first_seed else: all_parameters["seed"] = seed if esrgan_parameters: progress["currentStatus"] = "Upscaling" progress["currentStatusHasSteps"] = False socketio.emit("progressUpdate", progress) eventlet.sleep(0) image = esrgan.process( image=image, upsampler_scale=esrgan_parameters["level"], strength=esrgan_parameters["strength"], seed=seed, ) postprocessing = True all_parameters["upscale"] = [ esrgan_parameters["level"], esrgan_parameters["strength"], ] if gfpgan_parameters: progress["currentStatus"] = "Fixing faces" progress["currentStatusHasSteps"] = False socketio.emit("progressUpdate", progress) eventlet.sleep(0) image = gfpgan.process( image=image, strength=gfpgan_parameters["strength"], seed=seed ) postprocessing = True all_parameters["gfpgan_strength"] = gfpgan_parameters["strength"] progress["currentStatus"] = "Saving image" socketio.emit("progressUpdate", progress) eventlet.sleep(0) metadata = parameters_to_generated_image_metadata(all_parameters) command = parameters_to_command(all_parameters) path = save_image( image, command, metadata, result_path, postprocessing=postprocessing ) print(f'>> Image generated: "{path}"') write_log_message(f'[Generated] "{path}": {command}') if progress["totalIterations"] > progress["currentIteration"]: progress["currentStep"] = 1 progress["currentIteration"] += 1 progress["currentStatus"] = "Iteration finished" progress["currentStatusHasSteps"] = False else: progress["currentStep"] = 0 progress["totalSteps"] = 0 progress["currentIteration"] = 0 progress["totalIterations"] = 0 progress["currentStatus"] = "Finished" progress["isProcessing"] = False socketio.emit("progressUpdate", progress) eventlet.sleep(0) socketio.emit( "generationResult", { "url": os.path.relpath(path), "mtime": os.path.getmtime(path), "metadata": metadata, }, ) eventlet.sleep(0) try: generate.prompt2image( **generation_parameters, step_callback=image_progress, image_callback=image_done, ) except KeyboardInterrupt: raise except CanceledException: pass except Exception as e: socketio.emit("error", {"message": (str(e))}) print("\n") traceback.print_exc() print("\n") """ END ADDITIONAL FUNCTIONS """ if __name__ == "__main__": print(f">> Starting server at http://{host}:{port}") socketio.run(app, host=host, port=port)