Linux安装

1
2
3
4
5
6
7
安装python3.10
conda create -n ai310 python=3.10
source activate ai310
pip install -r requirements.txt

安装对应cuda版本的pytorch
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

修改启动端口

1
python -m uvicorn --host 0.0.0.0 --port=6969 --workers 1 main:app

多显卡优化

1
export CUDA_VISIBLE_DEVICES=0,1

内存分配优化(alloc memory error)

加上这个可以分配更大的图,2080Ti 11G可以输出 512*768的

1
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

并发模型,可以多个任务次序执行而不会崩溃

修改 main.py

1
2
3
4
5
6
7
8
9
10
11
12
# 1.nn.DataParallel支持
- model, config, model_hash = init_config_model()
- model, config, model_hash = init_config_model()
+ model = nn.DataParallel(model)

# 2.
- images = model.module.sample_two_stages(request)
+ images = model.sample_two_stages(request)

# 3. 这个两处相同修改
- images = model.module.sample(request)
+ images = model.sample(request)

最后大家一个试用地址


https://novelai.shawncx.top/