fix(photo-ai): лестница OOM на CUDA + приёмка face-режима и GPU-оверрайда
Stage 3 закрыт. Face-режим (off/face/all, strength, warnings) был написан в Stage 1;
здесь доведена приёмка и найден баг, который невозможно было увидеть без GPU-прогона.
Главное: realesные OOM никогда не доходили до run_guarded. И realessrgan/utils.py, и
gfpgan/utils.py ловят RuntimeError вокруг вызова сети и идут дальше
(`except RuntimeError as error: print('Error', error)`), поэтому на нехватку памяти
realesrgan падал уже не RuntimeError, а UnboundLocalError на присваивании
output_tile. Наружу уходил голый 500 «Internal Server Error»: ни лестницы тайлов,
ни деградации на CPU, ни внятного текста. В gfpgan это было тихое ухудшение —
при OOM лицо молча оставалось исходным, а задание уходило в «успех».
Починка: guard_forward() оборачивает forward сетей, которые строим мы
(RealESRGANer.model, restorer.gfpgan, CodeFormer net) и превращает OOM-RuntimeError
в TileOOM. TileOOM не наследует RuntimeError, поэтому проглатывающие except его
пропускают; is_oom и обе точки run_guarded ловят его явно. Обёртка вешается на
экземпляр, идемпотентна по флагу _photo_ai_guarded и не трогает класс.
Также добавлены два предупреждения из чек-листа, которых в коде не было: вход меньше
320×320 (лица могут не найтись) и CodeFormer на не-CUDA. Предупреждение «лица не
найдены» и деградация на CPU были на месте и не менялись.
Проверено на RTX 3050 Laptop (4096 МБ, драйвер 615.71.09, CUDA 12.6):
- tile=2048, x2plus face=all, полное фото: до правки 500 + UnboundLocalError,
после 200 за 28.3 с с единственным warning «не хватило памяти при tile=2048»;
- инъекция OOM: лестница 2048 → 1024 → 512, затем переход на CPU (device: cpu,
half: false) и успешный повтор; при повторе уже на CPU — честная 500 с подсказкой
про PHOTO_AI_TILE / PHOTO_AI_MAX_PIXELS;
- 640×480: off 1.0 с / face 3.8 с / all 2.2 с, faces_found=6; фото без лиц даёт
faces_found=0 и байты, равные face=off;
- I1: 6/6 MATCH байт-в-байт против Stage 2 на CPU (jpg/.jpeg/png+70/webp+100/x4v3/anime).
docker-compose.gpu.yml: GPU выдаётся через CDI (device_ids nvidia.com/gpu=all) —
не требует правки /etc/docker/daemon.json и перезапуска демона, в отличие от
классического резервирования driver: nvidia. TORCH_VARIANT cu124 → cu126: в индексе
cu124 последний torch 2.6.0, а cu126 даёт те же 2.14.0/0.29.0, что и CPU-образ, так
что варианты сборки отличаются только CUDA-библиотеками. Образ тегируется отдельно
(whatido-photo-ai:cu126), чтобы сборка GPU-варианта не перетирала CPU-образ
whatido-photo-ai:latest — откат остаётся обычным docker compose up -d photo-ai.
README: раздел «Запуск на NVIDIA GPU» с установкой NVIDIA Container Toolkit и генерацией
CDI-спеки, оговорками про 4 ГБ VRAM (x2plus + gfpgan влезают, general-x4v3 тяжелее,
LOAD_ALL=1 лучше не включать) и описанием параметров /enhance. .env.example: команда
GPU-запуска и рекомендация по PHOTO_AI_TILE. Журналы раздела 3 и GPU-прогона — в
TODO_PHOTO_FACE_AI.md.
worker.js, server.js, схема БД и фронтенд не тронуты — они в Stage 4…6.
This commit is contained in:
dev
+35
-4
@@ -38,6 +38,7 @@ FACE_MODES = ('off', 'face', 'all')
|
||||
DEFAULT_MODEL = 'x2plus'
|
||||
DEFAULT_STRENGTH = 0.7
|
||||
MIN_TILE = 64
|
||||
MIN_FACE_SIDE = 320
|
||||
POOL_LIMIT = 2
|
||||
WARMUP_SIZE = 64
|
||||
RETRY_AFTER_SEC = 5
|
||||
@@ -233,19 +234,41 @@ class ModelNotReady(EnhanceError):
|
||||
self.label = label
|
||||
|
||||
|
||||
class TileOOM(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def error_response(status_code, message, headers=None):
|
||||
return JSONResponse(status_code=status_code, content={'ok': False, 'error': message},
|
||||
headers=headers)
|
||||
|
||||
|
||||
def is_oom(err):
|
||||
if isinstance(err, torch.cuda.OutOfMemoryError):
|
||||
if isinstance(err, (torch.cuda.OutOfMemoryError, TileOOM)):
|
||||
return True
|
||||
text = str(err).lower()
|
||||
return any(mark in text for mark in ('out of memory', 'not enough memory',
|
||||
'alloc_cpu', "can't allocate memory"))
|
||||
|
||||
|
||||
def guard_forward(model):
|
||||
if getattr(model, '_photo_ai_guarded', False):
|
||||
return model
|
||||
forward = model.forward
|
||||
|
||||
def guarded(*args, **kwargs):
|
||||
try:
|
||||
return forward(*args, **kwargs)
|
||||
except RuntimeError as err:
|
||||
if not is_oom(err):
|
||||
raise
|
||||
raise TileOOM(str(err)) from err
|
||||
|
||||
model.forward = guarded
|
||||
model._photo_ai_guarded = True
|
||||
return model
|
||||
|
||||
|
||||
def tile_ladder(base):
|
||||
if base <= 0:
|
||||
return [base]
|
||||
@@ -459,7 +482,7 @@ def build_esrgan(name, denoise):
|
||||
scale=spec['scale'],
|
||||
model_path=model_path,
|
||||
dni_weight=dni_weight,
|
||||
model=spec['arch'](),
|
||||
model=guard_forward(spec['arch']()),
|
||||
tile=state['tile'],
|
||||
tile_pad=TILE_PAD,
|
||||
pre_pad=PRE_PAD,
|
||||
@@ -511,6 +534,7 @@ def build_face_gfpgan(spec, outscale, bg):
|
||||
bg_upsampler=bg.upsampler,
|
||||
device=torch.device(state['device']),
|
||||
)
|
||||
guard_forward(restorer.gfpgan)
|
||||
|
||||
def restore(img, strength):
|
||||
cropped, _restored, output = restorer.enhance(img, has_aligned=False, only_center_face=False,
|
||||
@@ -532,6 +556,7 @@ def build_face_codeformer(spec, outscale, bg):
|
||||
connect_list=['32', '64', '128', '256'], device=state['device'], fp16=False)
|
||||
net.load_state_dict(torch.load(path, map_location=lambda storage, loc: storage))
|
||||
net.eval()
|
||||
guard_forward(net)
|
||||
helper = face_helper(outscale)
|
||||
|
||||
def restore(img, strength):
|
||||
@@ -612,7 +637,7 @@ def run_guarded(img, outscale, model_name, face, face_name, strength, warnings):
|
||||
pool.set_tile(tile)
|
||||
try:
|
||||
return process_image(img, outscale, model_name, face, face_name, strength)
|
||||
except RuntimeError as err:
|
||||
except (RuntimeError, TileOOM) as err:
|
||||
if not is_oom(err):
|
||||
raise EnhanceError('ошибка модели: %s' % err, 500) from err
|
||||
log.warning('нехватка памяти при tile=%s: %s', tile, err)
|
||||
@@ -621,7 +646,7 @@ def run_guarded(img, outscale, model_name, face, face_name, strength, warnings):
|
||||
degrade_to_cpu(warnings)
|
||||
try:
|
||||
return process_image(img, outscale, model_name, face, face_name, strength)
|
||||
except RuntimeError as err:
|
||||
except (RuntimeError, TileOOM) as err:
|
||||
if is_oom(err):
|
||||
raise EnhanceError('не хватило памяти даже на CPU: %s' % err, 500) from err
|
||||
raise EnhanceError('ошибка модели на CPU: %s' % err, 500) from err
|
||||
@@ -819,6 +844,12 @@ async def enhance(request: Request,
|
||||
outscale = min(max(int(scale), 2), 4)
|
||||
fmt = output_format(image.filename)
|
||||
warnings = []
|
||||
if face_mode != 'off':
|
||||
if min(img.shape[0], img.shape[1]) < MIN_FACE_SIDE:
|
||||
warnings.append('вход меньше %d×%d — лица могут не найтись'
|
||||
% (MIN_FACE_SIDE, MIN_FACE_SIDE))
|
||||
if face_name == 'codeformer' and not state['device'].startswith('cuda'):
|
||||
warnings.append('CodeFormer на %s медленнее GFPGAN' % state['device'])
|
||||
output, faces_found = await asyncio.to_thread(run_guarded, img, outscale, model_name,
|
||||
face_mode, face_name, float(strength), warnings)
|
||||
payload, media_type = encode_image(output, fmt, int(quality))
|
||||
|
||||
Reference in New Issue
Block a user