import os import json from enum import Enum import socket import logging from contextlib import closing from invoke import task from rich import print from rich.prompt import IntPrompt, Prompt, Confirm from rich.logging import RichHandler FORMAT = "%(message)s" logging.basicConfig( level="INFO", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()] ) log = logging.getLogger("lama-cleaner") def find_free_port() -> int: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: s.bind(("", 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) return s.getsockname()[1] CONFIG_PATH = "config.json" class MODEL(str, Enum): SD15 = "sd1.5" LAMA = "lama" class DEVICE(str, Enum): CUDA = "cuda" CPU = "cpu" @task def info(c): print("Environment information".center(60, "-")) try: c.run("git --version") c.run("conda --version") c.run("which python") c.run("python --version") c.run("which pip") c.run("pip --version") c.run('pip list | grep "torch\|lama\|diffusers\|opencv\|cuda"') except: pass print("-" * 60) @task(pre=[info]) def config(c, disable_device_choice=False): # TODO: 提示选择模型,选择设备,端口,host # 如果是 sd 模型,提示接受条款和输入 huggingface token model = Prompt.ask( "Choice model", choices=[MODEL.SD15, MODEL.LAMA], default=MODEL.SD15 ) hf_access_token = "" if model == MODEL.SD15: while True: hf_access_token = Prompt.ask( "Huggingface access token (https://huggingface.co/docs/hub/security-tokens)" ) if hf_access_token == "": log.warning("Access token is required to download model") else: break if disable_device_choice: device = DEVICE.CPU else: device = Prompt.ask( "Choice device", choices=[DEVICE.CUDA, DEVICE.CPU], default=DEVICE.CUDA ) if device == DEVICE.CUDA: import torch if not torch.cuda.is_available(): log.warning( "Did not find CUDA device on your computer, fallback to cpu" ) device = DEVICE.CPU desktop = Confirm.ask("Start as desktop app?", default=True) configs = { "model": model, "device": device, "hf_access_token": hf_access_token, "desktop": desktop, } log.info(f"Save config to {CONFIG_PATH}") with open(CONFIG_PATH, "w", encoding="utf-8") as f: json.dump(configs, f, indent=2, ensure_ascii=False) log.info(f"Config finish, you can close this window.") @task(pre=[info]) def start(c): if not os.path.exists(CONFIG_PATH): log.info("Config file not exists, please run config.sh first") exit() log.info(f"Load config from {CONFIG_PATH}") with open(CONFIG_PATH, "r", encoding="utf-8") as f: configs = json.load(f) model = configs["model"] device = configs["device"] hf_access_token = configs["hf_access_token"] desktop = configs["desktop"] port = find_free_port() log.info(f"Using random port: {port}") if desktop: c.run( f"lama-cleaner --model {model} --device {device} --hf_access_token={hf_access_token} --port {port} --gui --gui-size 1400 900" ) else: c.run( f"lama-cleaner --model {model} --device {device} --hf_access_token={hf_access_token} --port {port}" )