100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
|
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 lama")
|
|||
|
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
|
|||
|
|
|||
|
configs = {"model": model, "device": device, "hf_access_token": hf_access_token}
|
|||
|
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']
|
|||
|
port = find_free_port()
|
|||
|
log.info(f"Using random port: {port}")
|
|||
|
|
|||
|
c.run(f"lama-cleaner --model {model} --device {device} --hf_access_token={hf_access_token} --port {port} --gui --gui-size 1400 900")
|