IOPaint/lama_cleaner/server.py

623 lines
18 KiB
Python
Raw Normal View History

2022-04-18 09:01:10 +02:00
#!/usr/bin/env python3
2023-03-29 16:05:34 +02:00
import os
2023-05-13 07:45:27 +02:00
import hashlib
2023-05-09 13:07:12 +02:00
2023-03-29 16:05:34 +02:00
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
2022-04-18 09:01:10 +02:00
2023-03-26 06:37:58 +02:00
import imghdr
2022-04-18 09:01:10 +02:00
import io
import logging
import multiprocessing
2022-09-15 16:21:27 +02:00
import random
2022-04-18 09:01:10 +02:00
import time
from pathlib import Path
import cv2
import numpy as np
2023-03-26 06:37:58 +02:00
import torch
from PIL import Image
2022-04-18 09:01:10 +02:00
from loguru import logger
2023-03-19 15:40:23 +01:00
from lama_cleaner.const import SD15_MODELS
2023-03-26 06:37:58 +02:00
from lama_cleaner.file_manager import FileManager
2023-03-22 05:57:18 +01:00
from lama_cleaner.model.utils import torch_gc
2022-04-18 09:01:10 +02:00
from lama_cleaner.model_manager import ModelManager
2023-03-26 07:39:09 +02:00
from lama_cleaner.plugins import (
InteractiveSeg,
RemoveBG,
RealESRGANUpscaler,
MakeGIF,
GFPGANPlugin,
2023-03-30 15:02:18 +02:00
RestoreFormerPlugin,
2023-05-13 07:45:27 +02:00
AnimeSeg,
2023-03-26 07:39:09 +02:00
)
2022-04-18 09:01:10 +02:00
from lama_cleaner.schema import Config
try:
torch._C._jit_override_can_fuse_on_cpu(False)
torch._C._jit_override_can_fuse_on_gpu(False)
torch._C._jit_set_texpr_fuser_enabled(False)
torch._C._jit_set_nvfuser_enabled(False)
except:
pass
2023-02-06 15:00:47 +01:00
from flask import (
Flask,
request,
send_file,
cli,
make_response,
send_from_directory,
jsonify,
)
2023-05-07 07:25:46 +02:00
from flask_socketio import SocketIO
2022-04-18 16:54:34 +02:00
2022-04-18 15:30:49 +02:00
# Disable ability for Flask to display warning about using a development server in a production environment.
# https://gist.github.com/jerblack/735b9953ba1ab6234abb43174210d356
cli.show_server_banner = lambda *_: None
2022-04-18 09:01:10 +02:00
from flask_cors import CORS
from lama_cleaner.helper import (
load_img,
numpy_to_bytes,
resize_max_size,
2023-02-06 15:00:47 +01:00
pil_to_bytes,
2022-04-18 09:01:10 +02:00
)
NUM_THREADS = str(multiprocessing.cpu_count())
# fix libomp problem on windows https://github.com/Sanster/lama-cleaner/issues/56
2022-09-15 16:21:27 +02:00
os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
2022-04-18 09:01:10 +02:00
os.environ["OMP_NUM_THREADS"] = NUM_THREADS
os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS
os.environ["MKL_NUM_THREADS"] = NUM_THREADS
os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS
os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS
if os.environ.get("CACHE_DIR"):
os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"]
BUILD_DIR = os.environ.get("LAMA_CLEANER_BUILD_DIR", "app/build")
2022-04-18 09:29:29 +02:00
class NoFlaskwebgui(logging.Filter):
def filter(self, record):
2023-05-13 07:45:27 +02:00
msg = record.getMessage()
if "Running on http:" in msg:
print(msg[msg.index("Running on http:") :])
return (
"flaskwebgui-keep-server-alive" not in msg
and "socket.io" not in msg
and "This is a development server." not in msg
)
2022-04-18 09:01:10 +02:00
2022-04-18 09:29:29 +02:00
logging.getLogger("werkzeug").addFilter(NoFlaskwebgui())
2022-04-18 09:01:10 +02:00
app = Flask(__name__, static_folder=os.path.join(BUILD_DIR, "static"))
app.config["JSON_AS_ASCII"] = False
CORS(app, expose_headers=["Content-Disposition"])
2023-05-13 07:45:27 +02:00
sio_logger = logging.getLogger("sio-logger")
sio_logger.setLevel(logging.ERROR)
2023-05-07 10:58:55 +02:00
socketio = SocketIO(app, cors_allowed_origins="*", async_mode="threading")
2022-04-18 09:01:10 +02:00
model: ModelManager = None
2023-01-08 13:53:55 +01:00
thumb: FileManager = None
output_dir: str = None
2022-04-18 09:01:10 +02:00
device = None
input_image_path: str = None
2022-11-13 06:14:37 +01:00
is_disable_model_switch: bool = False
2023-03-19 15:40:23 +01:00
is_controlnet: bool = False
2023-05-18 07:07:12 +02:00
controlnet_method: str = "control_v11p_sd15_canny"
2023-01-01 15:36:11 +01:00
is_enable_file_manager: bool = False
is_enable_auto_saving: bool = False
2022-11-16 10:59:39 +01:00
is_desktop: bool = False
image_quality: int = 95
2023-03-22 05:57:18 +01:00
plugins = {}
2022-04-18 09:01:10 +02:00
2022-11-18 14:40:12 +01:00
2022-04-18 09:01:10 +02:00
def get_image_ext(img_bytes):
w = imghdr.what("", img_bytes)
if w is None:
w = "jpeg"
return w
2022-10-15 16:32:25 +02:00
def diffuser_callback(i, t, latents):
2023-05-07 10:58:55 +02:00
socketio.emit("diffusion_progress", {"step": i})
2022-09-15 16:21:27 +02:00
2023-01-05 15:07:39 +01:00
@app.route("/save_image", methods=["POST"])
def save_image():
if output_dir is None:
return "--output-dir is None", 500
2023-01-05 15:07:39 +01:00
input = request.files
filename = request.form["filename"]
2023-01-05 15:07:39 +01:00
origin_image_bytes = input["image"].read() # RGB
ext = get_image_ext(origin_image_bytes)
2023-06-27 04:08:36 +02:00
image, alpha_channel, exif_infos = load_img(origin_image_bytes, return_exif=True)
save_path = os.path.join(output_dir, filename)
2023-06-27 04:08:36 +02:00
if alpha_channel is not None:
if alpha_channel.shape[:2] != image.shape[:2]:
alpha_channel = cv2.resize(
alpha_channel, dsize=(image.shape[1], image.shape[0])
)
image = np.concatenate((image, alpha_channel[:, :, np.newaxis]), axis=-1)
pil_image = Image.fromarray(image)
img_bytes = pil_to_bytes(
2023-06-27 04:08:36 +02:00
pil_image,
ext,
quality=image_quality,
exif_infos=exif_infos,
)
with open(save_path, "wb") as fw:
fw.write(img_bytes)
2023-02-06 15:00:47 +01:00
return "ok", 200
2023-01-05 15:07:39 +01:00
2023-01-07 13:51:05 +01:00
@app.route("/medias/<tab>")
def medias(tab):
2023-02-06 15:00:47 +01:00
if tab == "image":
2023-01-08 13:53:55 +01:00
response = make_response(jsonify(thumb.media_names), 200)
else:
response = make_response(jsonify(thumb.output_media_names), 200)
2023-01-08 14:59:26 +01:00
# response.last_modified = thumb.modified_time[tab]
2023-01-08 13:53:55 +01:00
# response.cache_control.no_cache = True
# response.cache_control.max_age = 0
2023-01-08 14:59:26 +01:00
# response.make_conditional(request)
2023-01-08 13:53:55 +01:00
return response
2022-12-31 14:07:08 +01:00
2023-02-06 15:00:47 +01:00
@app.route("/media/<tab>/<filename>")
2023-01-07 13:51:05 +01:00
def media_file(tab, filename):
2023-02-06 15:00:47 +01:00
if tab == "image":
2023-01-07 13:51:05 +01:00
return send_from_directory(thumb.root_directory, filename)
return send_from_directory(thumb.output_dir, filename)
2022-12-31 14:07:08 +01:00
2023-02-06 15:00:47 +01:00
@app.route("/media_thumbnail/<tab>/<filename>")
2023-01-07 13:51:05 +01:00
def media_thumbnail_file(tab, filename):
2022-12-31 14:07:08 +01:00
args = request.args
2023-02-06 15:00:47 +01:00
width = args.get("width")
height = args.get("height")
2022-12-31 14:07:08 +01:00
if width is None and height is None:
width = 256
if width:
width = int(float(width))
if height:
height = int(float(height))
2023-01-07 13:51:05 +01:00
directory = thumb.root_directory
2023-02-06 15:00:47 +01:00
if tab == "output":
2023-01-07 13:51:05 +01:00
directory = thumb.output_dir
2023-02-06 15:00:47 +01:00
thumb_filename, (width, height) = thumb.get_thumbnail(
directory, filename, width, height
)
2022-12-31 14:07:08 +01:00
thumb_filepath = f"{app.config['THUMBNAIL_MEDIA_THUMBNAIL_ROOT']}{thumb_filename}"
response = make_response(send_file(thumb_filepath))
response.headers["X-Width"] = str(width)
response.headers["X-Height"] = str(height)
return response
2022-04-18 09:01:10 +02:00
@app.route("/inpaint", methods=["POST"])
def process():
input = request.files
# RGB
origin_image_bytes = input["image"].read()
2023-05-07 10:58:55 +02:00
image, alpha_channel, exif_infos = load_img(origin_image_bytes, return_exif=True)
2022-12-10 15:06:15 +01:00
mask, _ = load_img(input["mask"].read(), gray=True)
2022-11-27 14:25:27 +01:00
mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
if image.shape[:2] != mask.shape[:2]:
2023-02-06 15:00:47 +01:00
return (
f"Mask shape{mask.shape[:2]} not queal to Image shape{image.shape[:2]}",
400,
)
2022-04-18 09:01:10 +02:00
original_shape = image.shape
interpolation = cv2.INTER_CUBIC
form = request.form
2023-03-25 07:11:00 +01:00
size_limit = max(image.shape)
2022-04-18 09:01:10 +02:00
2022-12-10 15:06:15 +01:00
if "paintByExampleImage" in input:
2023-02-06 15:00:47 +01:00
paint_by_example_example_image, _ = load_img(
input["paintByExampleImage"].read()
)
2022-12-10 15:06:15 +01:00
paint_by_example_example_image = Image.fromarray(paint_by_example_example_image)
else:
paint_by_example_example_image = None
2022-04-18 09:01:10 +02:00
config = Config(
2022-04-18 16:54:34 +02:00
ldm_steps=form["ldmSteps"],
2022-06-12 07:14:17 +02:00
ldm_sampler=form["ldmSampler"],
2022-04-18 16:54:34 +02:00
hd_strategy=form["hdStrategy"],
2022-07-13 03:04:28 +02:00
zits_wireframe=form["zitsWireframe"],
2022-04-18 16:54:34 +02:00
hd_strategy_crop_margin=form["hdStrategyCropMargin"],
hd_strategy_crop_trigger_size=form["hdStrategyCropTrigerSize"],
hd_strategy_resize_limit=form["hdStrategyResizeLimit"],
2022-09-20 16:43:20 +02:00
prompt=form["prompt"],
2022-11-08 14:58:48 +01:00
negative_prompt=form["negativePrompt"],
2022-09-20 16:43:20 +02:00
use_croper=form["useCroper"],
croper_x=form["croperX"],
croper_y=form["croperY"],
croper_height=form["croperHeight"],
croper_width=form["croperWidth"],
2023-01-05 15:07:39 +01:00
sd_scale=form["sdScale"],
2022-09-22 15:50:41 +02:00
sd_mask_blur=form["sdMaskBlur"],
2022-09-15 16:21:27 +02:00
sd_strength=form["sdStrength"],
sd_steps=form["sdSteps"],
sd_guidance_scale=form["sdGuidanceScale"],
sd_sampler=form["sdSampler"],
sd_seed=form["sdSeed"],
sd_match_histograms=form["sdMatchHistograms"],
2022-10-09 15:32:13 +02:00
cv2_flag=form["cv2Flag"],
2023-02-06 15:00:47 +01:00
cv2_radius=form["cv2Radius"],
2022-12-10 15:06:15 +01:00
paint_by_example_steps=form["paintByExampleSteps"],
paint_by_example_guidance_scale=form["paintByExampleGuidanceScale"],
paint_by_example_mask_blur=form["paintByExampleMaskBlur"],
paint_by_example_seed=form["paintByExampleSeed"],
paint_by_example_match_histograms=form["paintByExampleMatchHistograms"],
paint_by_example_example_image=paint_by_example_example_image,
2023-01-28 14:13:21 +01:00
p2p_steps=form["p2pSteps"],
p2p_image_guidance_scale=form["p2pImageGuidanceScale"],
p2p_guidance_scale=form["p2pGuidanceScale"],
2023-03-19 15:40:23 +01:00
controlnet_conditioning_scale=form["controlnet_conditioning_scale"],
2023-05-13 07:45:27 +02:00
controlnet_method=form["controlnet_method"],
2022-04-18 09:01:10 +02:00
)
2022-09-15 16:21:27 +02:00
if config.sd_seed == -1:
2022-10-15 16:32:25 +02:00
config.sd_seed = random.randint(1, 999999999)
2022-12-10 15:06:15 +01:00
if config.paint_by_example_seed == -1:
config.paint_by_example_seed = random.randint(1, 999999999)
2022-09-15 16:21:27 +02:00
2022-04-18 09:01:10 +02:00
logger.info(f"Origin image shape: {original_shape}")
image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
start = time.time()
try:
res_np_img = model(image, mask, config)
except RuntimeError as e:
if "CUDA out of memory. " in str(e):
2022-11-18 14:40:12 +01:00
# NOTE: the string may change?
return "CUDA out of memory", 500
2022-11-18 14:40:12 +01:00
else:
logger.exception(e)
2023-05-18 07:07:12 +02:00
return f"{str(e)}", 500
finally:
logger.info(f"process time: {(time.time() - start) * 1000}ms")
2023-05-20 06:35:36 +02:00
torch_gc()
2022-04-18 09:01:10 +02:00
2023-02-06 15:00:47 +01:00
res_np_img = cv2.cvtColor(res_np_img.astype(np.uint8), cv2.COLOR_BGR2RGB)
2022-04-18 09:01:10 +02:00
if alpha_channel is not None:
if alpha_channel.shape[:2] != res_np_img.shape[:2]:
alpha_channel = cv2.resize(
alpha_channel, dsize=(res_np_img.shape[1], res_np_img.shape[0])
)
res_np_img = np.concatenate(
(res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
)
ext = get_image_ext(origin_image_bytes)
2022-09-20 16:43:20 +02:00
2023-05-07 10:58:55 +02:00
bytes_io = io.BytesIO(
pil_to_bytes(
Image.fromarray(res_np_img),
ext,
quality=image_quality,
exif_infos=exif_infos,
)
)
2023-02-06 15:00:47 +01:00
2022-09-20 16:43:20 +02:00
response = make_response(
send_file(
2023-02-06 15:00:47 +01:00
# io.BytesIO(numpy_to_bytes(res_np_img, ext)),
bytes_io,
2022-09-20 16:43:20 +02:00
mimetype=f"image/{ext}",
)
2022-04-18 09:01:10 +02:00
)
2022-09-20 16:43:20 +02:00
response.headers["X-Seed"] = str(config.sd_seed)
2023-05-07 07:25:46 +02:00
2023-05-07 10:58:55 +02:00
socketio.emit("diffusion_finish")
2022-09-20 16:43:20 +02:00
return response
2022-04-18 09:01:10 +02:00
2023-03-25 02:53:22 +01:00
@app.route("/run_plugin", methods=["POST"])
2023-03-22 05:57:18 +01:00
def run_plugin():
form = request.form
files = request.files
name = form["name"]
if name not in plugins:
return "Plugin not found", 500
origin_image_bytes = files["image"].read() # RGB
2023-05-07 10:58:55 +02:00
rgb_np_img, alpha_channel, exif_infos = load_img(
origin_image_bytes, return_exif=True
)
2022-11-27 14:25:27 +01:00
start = time.time()
2023-03-30 15:02:18 +02:00
try:
2023-04-06 15:55:20 +02:00
form = dict(form)
if name == InteractiveSeg.name:
img_md5 = hashlib.md5(origin_image_bytes).hexdigest()
form["img_md5"] = img_md5
2023-03-30 15:02:18 +02:00
bgr_res = plugins[name](rgb_np_img, files, form)
except RuntimeError as e:
torch.cuda.empty_cache()
if "CUDA out of memory. " in str(e):
# NOTE: the string may change?
return "CUDA out of memory", 500
else:
logger.exception(e)
return "Internal Server Error", 500
2023-03-22 05:57:18 +01:00
logger.info(f"{name} process time: {(time.time() - start) * 1000}ms")
torch_gc()
2023-03-25 02:53:22 +01:00
if name == MakeGIF.name:
return send_file(
2023-03-26 14:42:31 +02:00
io.BytesIO(bgr_res),
2023-03-25 02:53:22 +01:00
mimetype="image/gif",
as_attachment=True,
2023-04-16 07:45:34 +02:00
download_name=form["filename"],
2023-03-25 02:53:22 +01:00
)
2023-03-26 15:02:51 +02:00
if name == InteractiveSeg.name:
return make_response(
send_file(
io.BytesIO(numpy_to_bytes(bgr_res, "png")),
mimetype="image/png",
)
)
2023-03-26 14:42:31 +02:00
2023-05-09 13:07:12 +02:00
if name in [RemoveBG.name, AnimeSeg.name]:
rgb_res = bgr_res
2023-03-26 14:42:31 +02:00
ext = "png"
2023-03-25 02:53:22 +01:00
else:
2023-03-26 14:42:31 +02:00
rgb_res = cv2.cvtColor(bgr_res, cv2.COLOR_BGR2RGB)
ext = get_image_ext(origin_image_bytes)
if alpha_channel is not None:
if alpha_channel.shape[:2] != rgb_res.shape[:2]:
alpha_channel = cv2.resize(
alpha_channel, dsize=(rgb_res.shape[1], rgb_res.shape[0])
)
rgb_res = np.concatenate(
(rgb_res, alpha_channel[:, :, np.newaxis]), axis=-1
2023-03-25 02:53:22 +01:00
)
2023-03-26 14:42:31 +02:00
response = make_response(
send_file(
io.BytesIO(
pil_to_bytes(
2023-05-07 10:58:55 +02:00
Image.fromarray(rgb_res),
ext,
quality=image_quality,
exif_infos=exif_infos,
2023-03-26 14:42:31 +02:00
)
),
mimetype=f"image/{ext}",
2022-11-27 14:25:27 +01:00
)
2023-03-26 14:42:31 +02:00
)
2022-11-27 14:25:27 +01:00
return response
2023-03-25 02:53:22 +01:00
@app.route("/server_config", methods=["GET"])
def get_server_config():
return {
"isControlNet": is_controlnet,
2023-05-18 07:07:12 +02:00
"controlNetMethod": controlnet_method,
2023-03-25 02:53:22 +01:00
"isDisableModelSwitchState": is_disable_model_switch,
2023-03-31 16:06:17 +02:00
"isEnableAutoSaving": is_enable_auto_saving,
"enableFileManager": is_enable_file_manager,
2023-03-25 02:53:22 +01:00
"plugins": list(plugins.keys()),
}, 200
2023-03-22 05:57:18 +01:00
2022-04-18 09:01:10 +02:00
@app.route("/model")
def current_model():
return model.name, 200
2022-11-18 14:40:12 +01:00
2022-04-18 09:01:10 +02:00
@app.route("/model_downloaded/<name>")
def model_downloaded(name):
return str(model.is_downloaded(name)), 200
2022-11-18 14:40:12 +01:00
2022-11-16 10:59:39 +01:00
@app.route("/is_desktop")
def get_is_desktop():
return str(is_desktop), 200
2022-04-18 09:01:10 +02:00
2022-11-18 14:40:12 +01:00
2022-04-18 09:01:10 +02:00
@app.route("/model", methods=["POST"])
def switch_model():
2022-11-25 01:55:10 +01:00
if is_disable_model_switch:
return "Switch model is disabled", 400
2022-04-18 09:01:10 +02:00
new_name = request.form.get("name")
if new_name == model.name:
return "Same model", 200
try:
model.switch(new_name)
except NotImplementedError:
return f"{new_name} not implemented", 403
return f"ok, switch to {new_name}", 200
@app.route("/")
def index():
2023-04-16 07:45:34 +02:00
return send_file(os.path.join(BUILD_DIR, "index.html"))
2022-04-18 09:01:10 +02:00
@app.route("/inputimage")
def set_input_photo():
if input_image_path:
with open(input_image_path, "rb") as f:
image_in_bytes = f.read()
return send_file(
input_image_path,
as_attachment=True,
2023-05-07 11:00:28 +02:00
download_name=Path(input_image_path).name,
2022-04-18 09:01:10 +02:00
mimetype=f"image/{get_image_ext(image_in_bytes)}",
)
else:
return "No Input Image"
2023-03-22 05:57:18 +01:00
def build_plugins(args):
global plugins
if args.enable_interactive_seg:
logger.info(f"Initialize {InteractiveSeg.name} plugin")
2023-04-06 15:55:20 +02:00
plugins[InteractiveSeg.name] = InteractiveSeg(
args.interactive_seg_model, args.interactive_seg_device
)
2023-05-09 13:07:12 +02:00
2023-03-22 05:57:18 +01:00
if args.enable_remove_bg:
logger.info(f"Initialize {RemoveBG.name} plugin")
plugins[RemoveBG.name] = RemoveBG()
2023-05-09 13:07:12 +02:00
if args.enable_anime_seg:
logger.info(f"Initialize {AnimeSeg.name} plugin")
plugins[AnimeSeg.name] = AnimeSeg()
2023-03-22 05:57:18 +01:00
if args.enable_realesrgan:
2023-03-25 03:15:44 +01:00
logger.info(
f"Initialize {RealESRGANUpscaler.name} plugin: {args.realesrgan_model}, {args.realesrgan_device}"
)
plugins[RealESRGANUpscaler.name] = RealESRGANUpscaler(
2023-04-03 07:32:04 +02:00
args.realesrgan_model,
args.realesrgan_device,
no_half=args.realesrgan_no_half,
2023-03-25 03:15:44 +01:00
)
2023-05-09 13:07:12 +02:00
2023-03-26 07:39:09 +02:00
if args.enable_gfpgan:
logger.info(f"Initialize {GFPGANPlugin.name} plugin")
2023-04-06 15:55:20 +02:00
if args.enable_realesrgan:
logger.info("Use realesrgan as GFPGAN background upscaler")
else:
logger.info(
f"GFPGAN no background upscaler, use --enable-realesrgan to enable it"
)
2023-03-26 14:52:06 +02:00
plugins[GFPGANPlugin.name] = GFPGANPlugin(
args.gfpgan_device, upscaler=plugins.get(RealESRGANUpscaler.name, None)
)
2023-05-09 13:07:12 +02:00
2023-03-30 10:07:38 +02:00
if args.enable_restoreformer:
logger.info(f"Initialize {RestoreFormerPlugin.name} plugin")
plugins[RestoreFormerPlugin.name] = RestoreFormerPlugin(
2023-03-30 15:02:18 +02:00
args.restoreformer_device,
upscaler=plugins.get(RealESRGANUpscaler.name, None),
2023-03-30 10:07:38 +02:00
)
2023-05-09 13:07:12 +02:00
2023-03-25 02:53:22 +01:00
if args.enable_gif:
logger.info(f"Initialize GIF plugin")
plugins[MakeGIF.name] = MakeGIF()
2023-01-08 13:53:55 +01:00
2022-04-18 09:01:10 +02:00
def main(args):
global model
global device
global input_image_path
2022-11-13 06:14:37 +01:00
global is_disable_model_switch
2023-01-01 15:36:11 +01:00
global is_enable_file_manager
2022-11-16 10:59:39 +01:00
global is_desktop
2023-01-08 13:53:55 +01:00
global thumb
global output_dir
global is_enable_auto_saving
2023-03-19 15:40:23 +01:00
global is_controlnet
2023-05-18 07:07:12 +02:00
global controlnet_method
global image_quality
2023-03-22 05:57:18 +01:00
build_plugins(args)
image_quality = args.quality
2023-03-19 15:40:23 +01:00
if args.sd_controlnet and args.model in SD15_MODELS:
is_controlnet = True
2023-05-18 07:07:12 +02:00
controlnet_method = args.sd_controlnet_method
output_dir = args.output_dir
2023-03-31 16:06:17 +02:00
if output_dir:
is_enable_auto_saving = True
2022-04-18 09:01:10 +02:00
device = torch.device(args.device)
2022-11-13 06:14:37 +01:00
is_disable_model_switch = args.disable_model_switch
2022-11-16 10:59:39 +01:00
is_desktop = args.gui
2022-11-13 06:14:37 +01:00
if is_disable_model_switch:
2023-02-06 15:00:47 +01:00
logger.info(
f"Start with --disable-model-switch, model switch on frontend is disable"
)
2022-04-18 09:01:10 +02:00
2023-01-08 14:26:57 +01:00
if args.input and os.path.isdir(args.input):
2023-01-08 13:53:55 +01:00
logger.info(f"Initialize file manager")
thumb = FileManager(app)
is_enable_file_manager = True
2022-12-31 14:07:08 +01:00
app.config["THUMBNAIL_MEDIA_ROOT"] = args.input
2023-02-06 15:00:47 +01:00
app.config["THUMBNAIL_MEDIA_THUMBNAIL_ROOT"] = os.path.join(
args.output_dir, "lama_cleaner_thumbnails"
)
2023-01-05 15:07:39 +01:00
thumb.output_dir = Path(args.output_dir)
2023-01-08 14:59:26 +01:00
# thumb.start()
2023-01-08 13:53:55 +01:00
# try:
# while True:
# time.sleep(1)
# finally:
# thumb.image_dir_observer.stop()
# thumb.image_dir_observer.join()
# thumb.output_dir_observer.stop()
# thumb.output_dir_observer.join()
2022-12-31 14:07:08 +01:00
else:
input_image_path = args.input
2022-09-20 16:43:20 +02:00
model = ModelManager(
name=args.model,
2023-03-19 15:40:23 +01:00
sd_controlnet=args.sd_controlnet,
2023-05-02 08:09:40 +02:00
sd_controlnet_method=args.sd_controlnet_method,
2022-09-20 16:43:20 +02:00
device=device,
2023-01-03 14:30:33 +01:00
no_half=args.no_half,
2022-09-20 16:43:20 +02:00
hf_access_token=args.hf_access_token,
disable_nsfw=args.sd_disable_nsfw or args.disable_nsfw,
2022-09-29 06:20:55 +02:00
sd_cpu_textencoder=args.sd_cpu_textencoder,
2022-09-29 07:13:09 +02:00
sd_run_local=args.sd_run_local,
2023-03-29 16:05:34 +02:00
sd_local_model_path=args.sd_local_model_path,
2023-01-05 15:07:39 +01:00
local_files_only=args.local_files_only,
cpu_offload=args.cpu_offload,
enable_xformers=args.sd_enable_xformers or args.enable_xformers,
2022-10-15 16:32:25 +02:00
callback=diffuser_callback,
2022-09-20 16:43:20 +02:00
)
2022-04-18 09:01:10 +02:00
if args.gui:
app_width, app_height = args.gui_size
from flaskwebgui import FlaskUI
2022-04-18 16:54:34 +02:00
ui = FlaskUI(
2023-02-06 15:00:47 +01:00
app,
2023-05-07 07:25:46 +02:00
socketio=socketio,
2023-02-06 15:00:47 +01:00
width=app_width,
height=app_height,
host=args.host,
port=args.port,
close_server_on_exit=not args.no_gui_auto_close,
2022-04-18 16:54:34 +02:00
)
2022-04-18 16:28:47 +02:00
ui.run()
2022-04-18 09:01:10 +02:00
else:
2023-05-13 07:45:27 +02:00
socketio.run(
app,
host=args.host,
port=args.port,
debug=args.debug,
allow_unsafe_werkzeug=True,
)