Изучение простых оптимизаций для SDXL

Освоение простых оптимизаций для SDXL

Открыть в Colab

Stable Diffusion XL (SDXL) – последняя модель латентной диффузии от Stability AI для создания высококачественных и фотореалистичных изображений. Она преодолевает проблемы предыдущих моделей Stable Diffusion, такие как правильное воспроизведение рук и текста, а также корректные композиции в пространстве. Кроме того, SDXL также более контекстно осознает и требует меньше слов в своём запросе для создания более красивых изображений.

Однако все эти улучшения происходят за счет значительно большей модели. Насколько она больше? Базовая модель SDXL имеет 3,5 миллиарда параметров (в частности, UNet), что примерно в 3 раза больше, чем предыдущая модель Stable Diffusion.

Для изучения того, как мы можем оптимизировать SDXL для скорости вывода и использования памяти, мы провели некоторые тесты на графическом процессоре A100 (40 ГБ). Для каждого запуска вывода мы генерируем 4 изображения и повторяем это 3 раза. При вычислении задержки вывода мы рассматриваем только конечную итерацию из 3 итераций.

Так что если вы запустите SDXL “из коробки”, не изменяя его и используя значения с полной точностью, а также используя механизм внимания по умолчанию, это займет 28 ГБ памяти и 72,2 секунды!

from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")pipeline.unet.set_default_attn_processor()

Это не очень практично и может замедлить вас, поскольку вы часто создаете более 4 изображений. А если у вас нет более мощного графического процессора, вы столкнетесь с раздражающим сообщением об ошибке “нехватка памяти”. Так как же мы можем оптимизировать SDXL для увеличения скорости вывода и уменьшения его использования памяти?

В 🤗 Diffusers у нас есть множество трюков и техник оптимизации, чтобы помочь вам запускать модели, требующие большого объема памяти, такие как SDXL, и мы покажем вам, как это сделать! Мы сосредоточимся на скорости вывода и использовании памяти.

Скорость вывода

Диффузия – это случайный процесс, поэтому нет гарантии, что вы получите изображение, которое вам понравится. Часто вам придется запускать вывод несколько раз и итерировать, и поэтому оптимизация скорости крайне важна. В этом разделе рассматривается использование весов с нижней точностью, а также включение памяти-эффективного механизма внимания и torch.compile из PyTorch 2.0 для увеличения скорости и сокращения времени вывода.

Нижняя точность

Веса модели сохраняются с определенной точностью, которая выражается в виде типа данных с плавающей запятой. Стандартным типом данных с плавающей запятой является float32 (fp32), который может точно представлять широкий диапазон плавающих чисел. Для вывода вам часто не нужна такая точность, поэтому вы должны использовать float16 (fp16), который записывает более узкий диапазон плавающих чисел. Это означает, что fp16 занимает в два раза меньше памяти по сравнению с fp32, и работает вдвое быстрее за счет более простых вычислений. Кроме того, современные графические карты имеют оптимизированное оборудование для выполнения вычислений с использованием fp16, что делает его еще быстрее.

С 🤗 Diffusers вы можете использовать fp16 для вывода, указав параметр torch.dtype, чтобы преобразовать веса при загрузке модели:

from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0",    torch_dtype=torch.float16,).to("cuda")pipeline.unet.set_default_attn_processor()

По сравнению с полностью неоптимизированной конвейерной системой SDXL, использование fp16 занимает 21,7 ГБ памяти и занимает всего 14,8 секунд. Вы почти ускоряете вывод на целую минуту!

Памяти-эффективное внимание

Блоки внимания, используемые в модулях трансформеров, могут быть огромным узким местом, поскольку память увеличивается квадратично с увеличением длины входных последовательностей. Это может быстро занять много памяти и привести к ошибке “нехватка памяти”. 😬

Алгоритмы памяти-эффективного внимания стремятся уменьшить объем памяти, необходимый для вычисления внимания, путем использования разреженности или плитки. Эти оптимизированные алгоритмы раньше в основном были доступны как сторонние библиотеки, которые нужно было устанавливать отдельно. Но начиная с PyTorch 2.0 это больше не так. PyTorch 2 ввел масштабированное внимание скалярное произведение (SDPA), которое предлагает объединенные реализации алгоритмов Flash Attention, памяти-эффективного внимания (xFormers) и реализацию PyTorch на C++. SDPA, вероятно, является самым простым способом ускорить вывод: если вы используете PyTorch ≥ 2.0 с 🤗 Diffusers, он автоматически включен по умолчанию!

from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16).to("cuda") 

По сравнению с полностью неоптимизированной SDXL-трубкой, использование fp16 и SDPA требует такого же объема памяти, а время вывода улучшается до 11,4 секунды. Давайте используем это в качестве новой отправной точки для сравнения с другими оптимизациями.

torch.compile

В PyTorch 2.0 также был представлен API torch.compile для компиляции вашего PyTorch-кода в более оптимизированные ядра для вывода на лету (JIT). В отличие от других решений компилятора, для torch.compile требуются минимальные изменения в вашем существующем коде, и все это так просто, как упаковка вашей модели с помощью функции.

С параметром mode вы можете оптимизировать использование памяти или скорость вывода во время компиляции, что дает вам гораздо больше гибкости.

from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16).to("cuda")pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)

По сравнению с предыдущей отправной точкой (fp16 + SDPA), обертывание UNet с помощью torch.compile улучшает время вывода до 10,2 секунды.

Память модели

Сегодня модели становятся все больше, что создает проблему вмещения их в память. В этом разделе рассматриваются способы сокращения памяти этих огромных моделей, чтобы можно было запускать их на обычных графических процессорах. Эти методы включают выполняемые на ЦП, декодирование латентных переменных в изображения на несколько шагов, а не сразу все, и использование укороченной версии автоэнкодера.

Выгрузка модели на ЦП

Выгрузка модели позволяет сэкономить память, загружая UNet в память графического процессора, в то время как другие компоненты модели диффузии (текстовые энкодеры, VAE) загружаются на ЦП. Таким образом, UNet может выполняться в течение нескольких итераций на графическом процессоре, пока он больше не понадобится.

from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipeline.enable_model_cpu_offload()

По сравнению с отправной точкой, сейчас требуется 20,2 ГБ памяти, что экономит 1,5 ГБ памяти.

Последовательная выгрузка на ЦП

Еще один тип выгрузки, который может сэкономить больше памяти за счет медленного вывода – последовательная выгрузка на ЦП. Вместо выгрузки всей модели – как UNet – веса модели, хранящиеся в разных подмодулях UNet, выгружаются на ЦП и загружаются на графический процессор только перед вызовом прямого прохода. Фактически, каждый раз загружаются только части модели, что позволяет сэкономить еще больше памяти. Единственный недостаток состоит в том, что он существенно медленнее, потому что вы загружаете и выгружаете подмодули много раз.

from diffusers import StableDiffusionXLPipelinepipeline = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipeline.enable_sequential_cpu_offload()

По сравнению с отправной точкой, здесь требуется 19,9 ГБ памяти, но время вывода увеличивается до 67 секунд.

Раскрой

В SDXL вариационный энкодер (VAE) декодирует уточненные латентности (предсказанные UNet) в реалистичные изображения. Требования к памяти этого шага масштабируются с числом изображений, которые требуется предсказать (размер пакета). В зависимости от разрешения изображения и доступной графической памяти, это может быть довольно затратно по памяти.

Здесь полезно использование “раскроя”. Тензор входных данных для декодирования разбивается на сегменты, и вычисление производится за несколько шагов. Это позволяет сэкономить память и работать с большими размерами пакетов.

pipe = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipe = pipe.to("cuda")pipe.enable_vae_slicing()

С использованием срезанных вычислений мы сокращаем объем памяти до 15,4 ГБ. Если добавить последовательное отложение ЦП, то память дополнительно сокращается до 11,45 ГБ, что позволяет генерировать 4 изображения (1024×1024) на запрос. Однако, при последовательном отложении также увеличивается задержка вывода.

Кеширование вычислений

Любая модель генерации изображений с учетом текста обычно использует текстовый энкодер для вычисления эмбеддингов из входного предложения. SDXL использует два текстовых энкодера! Это существенно влияет на задержку вывода. Однако, поскольку эти эмбеддинги остаются неизменными на протяжении обратного процесса диффузии, мы можем предварительно вычислить их и повторно использовать при необходимости. Таким образом, после вычисления текстовых эмбеддингов мы можем удалить их из памяти.

Сначала загрузите текстовые энкодеры и соответствующие токенизаторы и вычислите эмбеддинги из входного предложения:

tokenizers = [tokenizer, tokenizer_2]text_encoders = [text_encoder, text_encoder_2](    prompt_embeds,    negative_prompt_embeds,    pooled_prompt_embeds,    negative_pooled_prompt_embeds) = encode_prompt(tokenizers, text_encoders, prompt)

Затем очистите память GPU, чтобы удалить текстовые энкодеры:

del text_encoder, text_encoder_2, tokenizer, tokenizer_2flush()

Теперь эмбеддинги готовы пройти SDXL-пайплайн:

from diffusers import StableDiffusionXLPipelinepipe = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0",    text_encoder=None,    text_encoder_2=None,    tokenizer=None,    tokenizer_2=None,    torch_dtype=torch.float16,)pipe = pipe.to("cuda")call_args = dict(        prompt_embeds=prompt_embeds,        negative_prompt_embeds=negative_prompt_embeds,        pooled_prompt_embeds=pooled_prompt_embeds,        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,        num_images_per_prompt=num_images_per_prompt,        num_inference_steps=num_inference_steps,)image = pipe(**call_args).images[0]

В сочетании с SDPA и fp16, мы можем сократить объем памяти до 21,9 ГБ. Другие вышеупомянутые техники оптимизации памяти также могут использоваться с кешированными вычислениями.

Маленький автоэнкодер

Как уже упоминалось ранее, VAE декодирует латенты в изображения. Естественно, этот шаг прямо зависит от размера VAE. Давайте просто используем меньший автоэнкодер! Tiny Autoencoder от madebyollin, доступный в the Hub, занимает всего 10 МБ и является урезанной версией исходного VAE, используемого SDXL.

from diffusers import AutoencoderTinypipe = StableDiffusionXLPipeline.from_pretrained(    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)pipe = pipe.to("cuda")

С такой настройкой мы сокращаем требования к памяти до 15,6 ГБ, одновременно сокращая задержку вывода.

Заключение

В заключение и для подведения итогов по нашим оптимизациям:

Мы надеемся, что эти оптимизации сделают запуск ваших любимых пайплайнов легким и комфортным. Попробуйте эти техники и поделитесь своими изображениями с нами! 🤗