xsl
443cfa298f
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
28 lines
694 B
Python
28 lines
694 B
Python
import pynvml
|
|
threshold = 0.9
|
|
def get_gpu(need_gpu_id):
|
|
used = get_gpu_threshold(need_gpu_id)
|
|
#小于一定的
|
|
if used > threshold:
|
|
need_use = get_use_gpu()
|
|
return need_use[0]
|
|
else:
|
|
return need_gpu_id
|
|
|
|
def get_use_gpu():
|
|
use=[]
|
|
for index in range(pynvml.nvmlDeviceGetCount()):
|
|
used = get_gpu_threshold(index)
|
|
if used > threshold:
|
|
use.append(index)
|
|
return use
|
|
|
|
def get_gpu_count():
|
|
return pynvml.nvmlDeviceGetCount()
|
|
|
|
def get_gpu_threshold(index):
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
|
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
used = meminfo.used / meminfo.total
|
|
return used
|