Обзор и запуск Stable Diffusion 3.5
Марк Блуменау
Stable Diffusion 3.5 — это самая свежая (на момент написания) модель из семейства Stable Diffusion от Stability AI, выпущенная 22 октября 2024 года. Модель представлена в трёх вариантах:
- Large (8 млрд параметров)
- Large Turbo (ускоренная версия Large)
- Medium (2.5 млрд параметров)
Основные улучшения включают более точное следование промпту, повышенное качество изображений и оптимизации для работы на различных ускорителях.
Для раскрытия потенциала модели Medium достаточно всего лишь 10 Гб видеопамяти, что позволяет её запускать на игровых домашних компьютерах без особых сложностей. Если же у вас менее 10 Гб, то модель также есть в квантизованном варианте, который позволяет уменьшить количество требуемой памяти. И, конечно, существует возможность использования offloading, когда веса сменяют друг друга в памяти видеокарты. Однако, стоит учитывать, что последний подход будет гораздо более долгим.
Отметим также, что запускать модель возможно не только на видеокартах NVIDIA, но и картах AMD с поддержкой ROCm (6000-7000 серия игровых карт) и даже на видеокартах Intel. При должном старании модель запустится и на Ascend NPU от Huawei или Apple Silicon, но сегодня мы экзотикой заниматься не будем, и воспользуемся картой NVIDIA (секрет прост, запускать будем на компьютере автора с RTX3060 :) ). Запускать будем версию с большим количеством оптимизаций, чтобы Вы могли попробовать и на более слабых картах, код из примера ниже потреблял порядка 1,5 Гб видеопамяти. Генерация при этом занимала 41 секунду.
Если у Вас другие ускорители, то стоит обратить внимание на странички StabilityAI HuggingFace и StabilityAI, где указаны даже конкретные игровые видеокарты.
Для запуска нам понадобится установленный драйвер с CUDA,а также conda для создания виртуального окружения. Инструкция написана для операционной системы Linux, но если у вас стоит Windows, то потребуется добавить также WSL и драйвер для CUDA под WSL (см. инструкцию на сайте Microsoft).
Установка зависимостей
# Создаём виртуальное окружение conda create -n "sd3.5" python=3.12 conda activate sd3.5 # Обновляем pip pip install --upgrade pip # Устанавливаем библиотеку Diffusers и сопутствующие пакеты pip install diffusers transformers accelerate sentencepiece protobuf pip install xformers pip install -U "huggingface_hub[cli]"
Для скачивания модели требуется создать аккаунт на HuggingFace, а также согласиться с условиями доступа к модели. Не переживайте, можно указать страну Россия, а использование скаченной модели бесплатно, если не планируете зарабатывать суммы, превышающие 1 млн $. Далее потребуется залогиниться через свой терминал (командную строку) при помощи токена (получить можно здесь, достаточно выбрать Read вариант):
# Логинимся в Hugging Face huggingface-cli login
Пример запуска
import torch
from diffusers import StableDiffusion3Pipeline
# Загружаем pipeline "StableDiffusion3Pipeline" с заранее обученной моделью
# ignore_mismatched_sizes=True — игнорируем несовпадение размеров весов при загрузке
# low_cpu_mem_usage=False — не применяем оптимизации по снижению использования CPU-памяти при загрузке (даже так потребление CPU RAM с системой и всем-всем открытым было меньше 16 GB)
# torch_dtype=torch.float16 — используем половинную (half) точность для тензоров (FP16), чтобы снизить потребление VRAM
# variant="fp16" — загружаем оптимизированные для FP16 ядра, если они доступны
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium",
ignore_mismatched_sizes=True,
low_cpu_mem_usage=False,
torch_dtype=torch.float16,
variant="fp16"
)
# Включаем разделение (slicing) внимания (attention) для снижения пикового потребления видеопамяти
# При slicing внимание обрабатывается маленькими порциями, что может замедлить работу
pipe.enable_attention_slicing()
# Включаем последовательный offload: поочередно выгружаем части моделей на CPU до вызова forward()
# Это позволяет держать большую часть весов в оперативной памяти (обычной RAM), освобождая VRAM
pipe.enable_sequential_cpu_offload()
# Включаем flash-attention из xformers для более эффективного и быстрого вычисления внимания
# Позволяет дополнительно снизить потребление VRAM и ускорить инференс
pipe.enable_xformers_memory_efficient_attention()
# Генерируем изображение по текстовому запросу
# text_prompt — строка с описанием желаемой картинки
# num_inference_steps=40 — число шагов диффузии (чем больше, тем выше качество, но медленнее)
# guidance_scale=7.5 — масштаб условного управления (CFG): баланс соответствия тексту и креативности
# height, width — размер финального изображения в пикселях
image = pipe(
"Awesome high detail capybara in space with a sign 'AiEdu' in front of it, clearly visible.",
num_inference_steps=40,
guidance_scale=7.5,
height=512, width=512
).images[0]
# Сохраняем картинку
image.save("output.png")