Реалтайм диффузия в 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, которая в теории должна помочь с генерацией синтетики для обучения.

Личные итоги года.
Ничего не делал по данному проекту последние 8 месяцев, занимался всем подряд. Например пытался понять как можно создать универсальный компрессор для текста чтобы по итогу модель кушала меньше токенов и генерировала больше текста. Спойлер, ничего практически значимого не вышло. Потом пытался разобраться как на малом скейле (4 A100 GPU) прям с нуля максимально эффективно делать претрейн модельки в 1B, при этом оставаясь в экосистеме transformers и библиотеке lm eval, делая эвалюацию прямо в моменте обучения. Если кому интересно код вот тут https://github.com/dmitrymailk/vision_genai/blob/master/train_gym/massive_train/pretrain_edu/pretrain_hf_trainer.py Наверное позже напишу более подробный пост, но пока я этим занимался случилась некоторая инфляция этих знаний, Карпатый(https://github.com/karpathy/nanochat) и команда SmolLM(The Smol Training Playbook: The Secrets to Building World-Class LLMs https://huggingface.co/spaces/HuggingFaceTB/smol-training-playbook) выпустила свои статьи про подобные запуски. Хотя на мой взгляд это не отменяет ценность моей репы. Еще попутно значительно ускорил recurrent memory transformer, но чет особо не пошло. Оставшуюся часть года занимался исследованием области Computer Use, точно потом сделаю обзор локальных моделей на конец 2025. Спойлер, довольно неплохая модель оказалась fara https://github.com/microsoft/fara в целом нормально себя чувствует на картах до 24 гб, и стабильно может выполнять очень простые действия в вебе. В целом планирую двигаться в этой области в следующем году, она как раз такая техничная и комплексная, как я люблю.

Желаю всем больше нейрослопа в ленте в этом году.
Последние месяц я решил возобновить работу над проектом. За время простоя вышло много видео моделей, лор и техник. Так как я сильно привязался к идее покадрового редактирования кадра, а затем его сглаживания, я продолжал копать в эту сторону. По итогу получился следующий пайплайн, общую идею которую я прочитал в ditto https://github.com/EzioBy/Ditto . Если очень кратко, их главная заслуга в том что они смогли очень сильно автоматизировать процесс создания датасета с инструкциями для редактирования видео. Припомощи WAN VACE и QWEN-EDIT, qwen редактирует ключевые кадры, а wan-vace уже на основе карты глубины выдает итоговое видео. Я поэкспериментировал с этим подходом, по итогу получился следующий алгоритм.

Этап 1: Первичная покадровая стилизация (SDXL)
На вход поступает последовательность кадров, независимо обработанных через SDXL.

Этап 2: Восстановление временной структуры (Wan VACE 14B)
Применяется модель VACE в режиме Masked Video-to-Video (MV2V).
Берется 250 кадров. Фиксирую по маске каждый 16-й кадр. Все промежуточные кадры стартуют свою генерацию со стилизованных изображений. Ключевое отличие тут от стандартного использования VACE, я попробовал заменить карту глубины между этими кадрами на также стилизованные кадры и это позволило сохранить стиль напротяжении всего видео и пофиксить все серьезные ошибки мерцания и деформации объектов.

Этап 3: Устранение шума (SeedVR2)
Сглаженное видео передается в модель SeedVR2, специализирующуюся на реставрации видео. Она отлично убирает оставшиеся серьезные артефакты, которые остались от покадровой стилизации. Но, на моих данных она замыливает текстуры. Для этого их нужно добавить обратно.

Этап 4: Восстановление деталей (Wan Video 14B Video-to-Video)
Для возвращения текстур и исправления финальных темпоральных ошибок выполняется финальная обработка моделью Wan Video 14B в режиме Video-to-Video. Процесс генерации запускается с экстремально низким уровнем шумоподавления (denoise = 0.09) и ограничением в 2 шага (steps = 2). Значение CFG устанавливается на 1.0, сэмплер – uni_pc.

Этап 5: Восстановление текстуры дорожного полотна
Поскольку дорожное полотно занимает значительную часть кадра, потеря его детализации негативно сказывается на общем восприятии четкости видео. Для этого я при помощи florence-2 и SAM2 нахожу дорогу, смешиваю плавно по гауссу оригинальное изображение со сглаженным кадром. Тоже самое что я делал c VEnhancer.

На выходе мы имеем наконец, фактически идентичное общему стилю, с сохранением деталей из оригинального SDXL, плавное видео. На данном этапе я решил остановиться, я считаю этого пайплайна в целом достаточно для генерации неограниченного количества синтетики для обучения MVP.

Сейчас я сфокусировался на создании действительно качественной одношаговой диффузионки. В целом для этого в этом году вышли все работы: repa-e(end-to-end обучение диффузионки совместно с VAE), latent bridge matching(рабочая и многообещающая теория на замену flow matching), seedvr2(одношаговая диффузионка для реконструкции, которая много юзает из работы The GAN is dead; long live the GAN!), DMD2(и различные техники дистиляции), rectified flow(выпрямление решающих потоков), [TwinFlow](https://github.com/inclusionAI/TwinFlow) и тд. Просто на мой взгляд теперь надо создать end2end алгоритм который бы бесшовно адаптировал вот этот генерируемый набор кадров в качественный реалтаймовый фильтр для 4090, а затем и 4060 Laptop. Теперь это просто програмерская задачка по перебору гиперпараметров и переписыванию кода. Думаю что через пару месяцев буду ездить в NFS с фотореалистичной графикой с видеорегистраторов.

Отказался от использования костыльного решения с reshade. Теперь я напрямую использую захват экрана от windows API, передаю данную текстуру в cuda, затем в pipeline на tensorRT. Это в свою очередь позволяет достичь 56-60 кадров с одношаговой (VAE->UNET->VAE) на моей ноутбучной RTX 4060 и ryzen 9 7945HX на слегка заниженных частотах. Но это все хорошо с NFS, я пробовал запускать в dirt 3, dirt 2 colin mc rae, там 32-35 всего. С точки зрения пользователя это работает примерно так, запускаем игру, запускаем бинарник через консоль и выбираем нужное окно, и прога начинает захват. Затем можно нажать f9 и поверх целевого окна рисуется новое окно которое пропускает все инпуты в целевую прогу. Там я еще влепил RTX video super resolution, потому что 512х512 на моем 2к экране тяжело смотреть, просадка примерно 4-5 FPS, честно говоря вообще того не стоит лучше потом на какой нибудь nvidia image scaling или обычный бикубик. DLSS не внедрить, там нужно кучу внутряков от движка передавать, что противоречит моей философии основываться только на входной картинке.

Конечно я сейчас использую довольно жирный unet, но большее время занимает VAE. Сначала я использовал fal/FLUX.2-Tiny-AutoEncoder, но он слишком мылил и убивал детали в декодированной картинке, так что пришлось натренировать свою на целевых данных, что дало четкость и мелкие текстуры. Для этого юзаю слегка модифицированный https://github.com/cloneofsimo/vqgan-training.

К своему сожалению или облегчению я обнаружил что обучение на сглаженных картинках при помощи видеогенераторов не особо нужно в текущем сетапе обучения.
Благодаря тому что вышел flux klein 4B, можно обучить любую лору на редактирование. У меня уже был датасет гиперреалистичных изображений полученных при помощи двойного апскейла SDXL. Их главная проблема что в далеке они создают чудовищые искажения объектов что не подходит под задачу ремастера. Но для объектов вблизи и средней давности получается неплохо. Сначала я попробовал обучить лору на всем датасете без фильтрации артефактов, получилось идеально воспроизвести все ошибки и на тесте. отлично подумал я и завайбкодил за час инструмент для разметки данных, из 4к картинок меня устроило всего 241 dim/nfs_pix2pix_1920_1080_v6_upscale_2x_raw_filtered. Соответственно после обучения на нем результат стал значительно лучше на тестовых данных, которые модель раньше не видела. Но все равно осталась странная проблема со скачком цветов, логики генерации растительности и тд. Я списал это на проблему шума. Поэтому попробовал применить structured noise https://yuzeng-at-tri.github.io/ppd-page/ что значительно улучшило результаты и я уже точно могу сказать что сейчас проблема только в данных на которых тренирую лору(да в них еще есть артефакты на дальние автомобили и объекты). Интересный факт, лора обучаласть на размерность 1 на q,k,v и итоговая лора весит всего 503 килобайта. Килобайта ало. Сама модель конечно весит в BF16 7.75 GB, лора тут https://huggingface.co/dim/nfs_pix2pix_1920_1080_v6_upscale_2x_raw_filtered_noise_lora_1_klein_4B_512x512

В своих попытках снизить галлюцинации, я так все зарегулировал что модель превратилась в просто покрасочный фильтр, который не изменяет геометрию. Из плюсов максимально стабильная картинка без кипения, из минусов пропадает как таковой смысл модели, детальный асфальт или листву она не создает. Вообще я думал над дистиляцией или встройкой GAN, я пробовал r3gan и patch gan, но они в какой-то момент у меня взрываются, я не умею это готовить. Также начал разбирать reflow meanflow, piflow, twinflow, senseflow, но пока на уровне самих репозиториев, наверное это должно работать лучше. Еще думаю над Latent Adversarial Diffusion Distillation, ее вроде как stable diffusion использовали.

Теперь имея на руках прогу для реального инференса я могу сказать что такой размер модели это предел без каких-то хитрых кернелов или алгоритмов, зато понятно пространство перебора параметров и методов.