Реалтайм диффузия в 100fps.

Именно столько выдает на данный момент связка 4090+ultra 7 265k при запуске диффузии на 512х512 и hidden size 64.
Я взял taesd, а из unet выкинул все модули с attention и заменил все на convolution. Также диффузия является одношаговой, то есть мы просим сгенерировать картинку сразу. Получается такой пайплайн. vae.decode(scheduler.step(unet(vae.encode(image)))). Scheduler я использую пока что тот, что был в работе pix2pix turbo DDPMScheduler. Вероятно его смена тоже может улучшить изображение. Для компиляции vae, unet использовал stable-fast.

Замеряю этим кодом. Картинки уже преобразованы и лежат в оперативке на CPU (это момент когда я буду их получать из буфера в reshade), потом перевожу ее в CUDA, однако если представить что все картинки сразу в CUDA тогда будет 191 FPS! С учетом того что играбельность начинается с 20 fps, вполне норм.

import time

amount = 400
start = time.time()
for image in images[:amount]:
    output_image = model(image.cuda())
total = time.time() - start
print(f"total {total}s, {total/amount}s per img {1/(total/amount)}fps")
# total 4.003062725067139s, 0.010007656812667846s per img 99.92349045524668fps

Сначала запустил обучение на dim/nfs_pix2pix_1920_1080_v6, wandb, 15_000 шагов, в целом результат удовлетворительный, модель почти идеально научилась воспроизводить обучающую выборку, несмотря на то что я существенно сократил размер модели и выкинул attention слои. Примеры можно посмотреть в wandb. Скачать тут dim/nfs_pix2pix_1739912522

Однако потери есть. Так как изначальные картинки мутные и детали там не явно выражены, при обучении, микродетали теряются и картинка получается мыльной. На итоговом видео это выглядит уже как просто перекраска, а не ремастер. Тогда я решил улучшить качество данных, сделав апскейл двойным прогоном. workflow. Долгое время пытался заставить апскелить FLUX, но сдался и все сделал на RealVisXL_V4.0_Lightning. Примеры выше\ниже хз где телега их разместит.

Теперь точно можно сказать что картинка улучшилась, причем такого качества точно никогда не достичь обычными рукописными алгоритмами или же основанные на простых функциях. Данные решают. Сделал 2 варианта датасета, то что получается после 1 аскейла и 2. Данные тут dim/nfs_pix2pix_1920_1080_v5_upscale_1x_raw dim/nfs_pix2pix_1920_1080_v5_upscale_2x_raw искать на https://huggingface.co/datasets. К сожалению я пока не понял почему теряется качество при сохранении и можно этого как-то избежать, поэтому также выложил просто пожатые архивы render_nfs_4screens_5_sdxl_1_upscale_1x_raw.tar.gz и render_nfs_4screens_5_sdxl_1_upscale_2x_raw.tar.gz. У данных апскейлов есть галлюцинации на объектах в далеке, которые я не стал фиксить. Апскейл 852 картинок занял примерно 36-48 часов на одной 4090. Делать батчами, ускорять и тд не пробовал.

После обучения на dim/nfs_pix2pix_1920_1080_v5_upscale_2x_raw wandb картинки стали четче. Теперь думаю можно сказать что модель не справляется. Наверное нужна другая архитектура, например Sana. Скачать тут dim/nfs_pix2pix_1740323104

Что видео? На мой взгдят картинка стала хуже. Из-за нестабильной картинки в апскейле, в далеке мы получаем кашу как и в оригинале. Однако если закрыть глаза ровно по середине можно понять что уже не так все плохо. Проблема в данных.

Как бы мне не хотелось разбираться с нейронками, так как не совсем понятны ограничения конечного пайплайна, думаю внедрить нейронки что есть end2nd, чтобы была возможность поиграть. А потом уже запустить ресерч по поиску архитектур, улучшению данных и тд.

Кстати базовый код на С++ уже готов, https://github.com/dmitrymailk/reshade/blob/main/examples/01-framerate_limit%20copy/negative_filter.cpp. Он меняет цвета картинки на противоположные. Тестил только на 3050ti, когда все выключено 90fps, когда код просто выполняет загрузку картики в массив, а потом обратно возвращает ее в пайплайн(короч ничего не делает над картинкой), тогда 72fps. Если же я включаю фильтр 50fps (правильно фильтр это просто 2 вложенных цикла). Теперь нужно как-то вместо данного цикла воткнуть нейронку.

У меня наконец получилось подружить диффузию с reshade.

По итогу получился совершенно проклятый код. Пайплайн получился примерно следующий. Сначала создается семафор на питоне который ждет пока его не тригернут, тем самым не закрывая программу на python. Далее на основе с++ API reshade, получаю доступ к backbuffer, сохраняю данную картинку в промежуточный массив. Создаю shared memory в с++ для inter process communication с python. Далее копирую изображение в shared memory и тригерю питоновский семафор, одновременно вешая другой семафор в с++ который ожидает пока его не тригернут из python. В самом питоне уже просто запихиваю данное изображение в pytorch через протокол buffer. Далее как обычно передают изображение в нейронку, результат сохраняю в shared memory, тригерю семафор для с++.

Зачем так проклято? Насколько я понял, моя игра работает в х32, и если я соберу addon для reshade в х64, он просто не увидит его и ничего не заработает. Зачем нужен х64? Только в нем работают нейронки на питоне да и большинство программ для их внедрения больше всего любят х64. К сожалению я узнал это только когда написал обертку для pybind, которая запускает интерпретатор питона из с++ и умеет импортировать любые модули для их дальнейшего использования, но это работает только с x32 питоном.

Из грустных новостей у меня так и не получилось поставить stable fast на windows, и вообще с нейронками на windows больно работать. Поэтому прилось запускать ее на голом торче, что дает на моем железе 70fps. Однако самый главный фактор который роняет фпс в игре оказалась не нейронка, а копирование картиночного массива туда-сюда. Из мыслей как это можно исправить это как-то уйти от перегонки с CPU на GPU, или хотябы данные массивы сделать shared memory.

Насчет конечного результата для пользователя пока не все так хорошо, картинка слишком убогая, потому что изначальаня задумка была для того чтобы картинка стала как в апскейлах, которые я показывал в прошлый раз. Результат можно посмотреть на ютубе или тут.

Код для аддона пока что лежит тут. Из интересного появилась новая restyle фича от runway, которая в теории должна помочь с генерацией синтетики для обучения.