Локальная генерация изображений и видео с помощью квантованных моделей

Вы когда-нибудь хотели генерировать изображения в своём стиле локально, но думали, что для этого нужна дорогая видеокарта? Я тоже так думал, пока не нашёл способ запустить FLUX на 12 ГБ VRAM. В этой статье я покажу, как с помощью квантования, LoRA и градиентных чекпоинтов обойти ограничения памяти и получить качественные результаты. Вы узнаете точный workflow, который я использовал, и сможете повторить его самостоятельно.

Проблема: как запустить большую модель на слабой видеокарте

Современные модели генерации изображений, такие как FLUX, требуют 24+ ГБ VRAM. Но что делать, если у вас только 12 ГБ? Ответ — комбинация трюков: квантование, заморозка весов, предвычисление эмбеддингов и управление памятью. Ниже разберём каждый этап.

Шаг 1: Квантование модели

Квантование уменьшает размер весов модели за счёт снижения точности чисел. Вместо 16-битных чисел с плавающей точкой (FP16) мы используем 4-битные (NF4). Это сжимает модель примерно в 4 раза. Библиотека bitsandbytes делает это легко.

from transformers import BitsAndBytesConfig
import torch

# Конфигурация 4-битного квантования
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Загрузка модели с квантованием
model = AutoModelForCausalLM.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=quant_config,
    device_map="auto"
)

Шаг 2: Тонкая настройка с LoRA

LoRA (Low-Rank Adaptation) позволяет дообучить модель на ваших изображениях, не трогая основные веса. Вы тренируете маленькие матрицы — «линзы», которые накладываются на базовую модель. Это экономит память и время.

from peft import LoraConfig, get_peft_model

# Конфигурация LoRA
lora_config = LoraConfig(
    r=16,  # ранг адаптера
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],  # слои для адаптации
    lora_dropout=0.1,
    bias="none"
)

# Применяем LoRA к квантованной модели
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # <1% параметров обучаемы

Шаг 3: Предвычисление эмбеддингов

Чтобы не пересчитывать эмбеддинги каждый раз, кэшируем их. Это особенно полезно для текстовых описаний (промптов).

from diffusers import FluxPipeline
import torch

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

# Предвычисляем и кэшируем эмбеддинги промпта
prompt = "a cat in a hat"
with torch.no_grad():
    prompt_embeds = pipe.encode_prompt(prompt)
    # сохраняем в файл или память
    torch.save(prompt_embeds, "prompt_embeds.pt")

Шаг 4: Управление памятью

Используем маленькое разрешение, batch size = 1, градиентное накопление и чекпоинты. Вот пример цикла обучения:

from torch.utils.data import DataLoader
from diffusers import FluxPipeline
import torch

# Гиперпараметры для экономии памяти
batch_size = 1
resolution = 256  # низкое разрешение
accumulation_steps = 4  # градиентное накопление
gradient_checkpointing = True  # чекпоинты

# Загрузка данных
dataloader = DataLoader(your_dataset, batch_size=batch_size)

# Включение градиентных чекпоинтов
if gradient_checkpointing:
    pipe.unet.enable_gradient_checkpointing()
    pipe.text_encoder.enable_gradient_checkpointing()

# Оптимизатор 8-битный (экономит память)
import bitsandbytes as bnb
optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=1e-4)

# Цикл обучения
for step, batch in enumerate(dataloader):
    images = batch["images"].to("cuda")
    prompts = batch["prompts"]
    
    # Forward pass
    loss = pipe(images, prompts).loss
    loss = loss / accumulation_steps
    loss.backward()
    
    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
        print(f"Step {step}, Loss: {loss.item():.4f}")

Практический результат: генерация в вашем стиле

После обучения LoRA на 14 своих изображениях (с метаданными в JSONL) я получил стилизованные картинки. Вот сравнение:

• Без обучения: случайные изображения, не похожие на мой стиль.
• С обучением (100 шагов): отдалённое сходство.
• С обучением (700 шагов): почти точное попадание в стиль.

Вы можете регулировать количество шагов: чем больше, тем точнее, но дольше.

Генерация видео

На основе той же модели можно генерировать короткие видео. Я запустил 5-секундный ролик на 4070 — заняло около часа. Код аналогичен, только добавляется временная размерность.

# Псевдокод для видео (упрощённо)
from diffusers import FluxVideoPipeline

video_pipe = FluxVideoPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-video",
    torch_dtype=torch.bfloat16
)
video_pipe.to("cuda")

video_frames = video_pipe(
    prompt="a painting coming to life",
    num_frames=16,
    height=256,
    width=256
).frames

Что делать прямо сейчас

1. Установите зависимости: pip install torch transformers diffusers peft bitsandbytes
2. Скачайте квантованную модель FLUX с Hugging Face (например, black-forest-labs/FLUX.1-dev)
3. Создайте папку с 10-20 своими изображениями и файлом metadata.jsonl с описаниями.
4. Запустите обучение LoRA по примеру выше.
5. Генерируйте изображения с помощью обученной модели.

Не бойтесь экспериментировать — даже на 12 ГБ можно получить отличные результаты. Если у вас 4060, тоже должно работать (чуть медленнее). Делитесь результатами в комментариях!