diffusers/v0.43.0/ru/api/pipelines/diffusion_gemma.md

DiffusionGemma

DiffusionGemma — это блочно-диффузионная модель языка с архитектурой «кодировщик-декодировщик». Причинный кодировщик считывает чистый промпт (а также ранее сгенерированные блоки) в KV-кэш, а двунаправленный декодер выполняет денуазинг фиксированного «холста» из canvas_length токенов, обращаясь к этому кэшу через кросс-внимание. Генерация чередует внешний авто-регрессивный цикл по холстам с внутренним циклом денуазинга, на каждом шаге которого отбираются кандидаты токенов, наиболее уверенные из них фиксируются через BlockRefinementScheduler в режиме равномерного повреждения, а остальные подвергаются повторному шумовому воздействию. Сама модель находится в transformers в виде DiffusionGemmaForBlockDiffusion; опубликованная контрольная точка — google/diffusiongemma-26B-A4B-it.

Использование

import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = "google/diffusiongemma-26B-A4B-it"
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead", fullgraph=True)
output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
)
print(output.texts[0])

num_inference_steps — количество шагов денойзинга на холст (48 соответствует опубликованной контрольной точке); меньшее число шагов ускоряет процесс, но снижает качество. cache_implementation="static" позволяет декодеру компилироваться с помощью torch.compile с использованием cudagraphs (см. Статический кэш и компиляция); для более простого запуска с динамическим кэшем уберите оба параметра.

Для многораундовых или мультимодальных входов передайте исходный разговор messages вместо prompt. Это список словарей {"role", "content"} в стандартном формате чата, который процессор обрабатывает через свой шаблон чата:

messages = [
    {"role": "user", "content": "Why is the sky blue?"},
]
# or with an image:
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": "Describe this image."},
        ],
    },
]
output = pipe(messages=messages, gen_length=256)

Для одного хода пользователя можно пропустить messages и передать image вместе с prompt; процессор автоматически преобразует их во входные данные изображения для модели.

Планировщики

Планировщик — это сэмплер, который выполняет денуазинг каждого холста, и он заменяем: вы можете заменить его, чтобы изменить стратегию сэмплирования, не изменяя ничего другого. Доступны три планировщика:

• BlockRefinementScheduler (по умолчанию): на каждом шаге фиксирует наиболее уверенные токены (те, что выше threshold, плюс равномерная квота на каждый шаг) и заново добавляет шум к остальным. Параметр editing_threshold дополнительно позволяет переопределять уже зафиксированные токены.
• DiscreteDDIMScheduler: выборка каждой позиции выполняется из точного дискретного апостериорного распределения равномерного процесса повреждения (D3PM). Этот планировщик не имеет параметров, а на последнем шаге предсказанные токены фиксируются детерминированно.
• EntropyBoundScheduler: фиксирует позиции с наименьшей энтропией, совместная энтропия которых не превышает entropy_bound, благодаря чему приблизительно независимые токены принимаются вместе. Температура выборки аннелируется от t_max (0.8) на первом шаге до t_min (0.4) на последнем, что соответствует сэмплеру опубликованной контрольной точки.

from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
# or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

Параметры выборки, специфичные для планировщика (threshold/top_k для уточнения блока, граница энтропии, ...) задаются в конфигурации планировщика:

from diffusers import BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)

EntropyBoundScheduler внутренним образом аннелирует температуру выборки (t_max/t_min) на шагах денуазинга; DiscreteDDIMScheduler и BlockRefinementScheduler используют плоскую temperature, переданную в пайплайн (0.0 для жадного алгоритма).

Сэмплирование с предсказателем-корректором

DiscreteDDIMScheduler поддерживает предиктор-корректор leave-one-out из работы Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. Он уточняет холст с помощью corrector_steps шагов Гиббса, которые пересэмплируют позиции с наименьшей уверенностью из однокоординатного условного распределения шумного маргинального распределения, что оставляет это маргинальное распределение инвариантным и улучшает генерацию без дополнительных затрат на обучение. Он работает непосредственно с опубликованной контрольной точкой: для uniform diffusion денойзер и leave-one-out апостериорное распределение взаимозаменяемы в замкнутой форме, поэтому корректор восстанавливает необходимые ему leave-one-out величины без какого-либо переобучения.

Шаги корректора включаются в бюджет num_inference_steps, а не добавляются сверх него: пайплайн выполняет меньше шагов предиктора и тратит высвобожденные проходы на корректоры, поэтому общее число проходов модели остаётся равным num_inference_steps, а стоимость предиктора-корректора совпадает со стоимостью обычного анцестрального сэмплирования.

from diffusers import DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=12)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

PEFT-адаптеры

Денойзер представляет собой модель 🤗 Transformers, поэтому адаптеры загружаются через его нативную интеграцию PEFT, а не через API load_lora_weights diffusers. Поскольку эта интеграция не зависит от типа адаптера, те же вызовы загружают LoRA, DoRA или любой другой адаптер PEFT (например, результат работы SFTTrainer из TRL). Управляйте адаптерами непосредственно на компоненте модели:

pipe.model.load_adapter("path/to/adapter", adapter_name="sft")  # LoRA, DoRA, ...
pipe.model.set_adapter("sft")
output = pipe(prompt="Why is the sky blue?", gen_length=256)

pipe.model.disable_adapters()  # run the base model
pipe.model.delete_adapter("sft")

Адаптеры остаются активными и не объединяются: DiffusionGemma связывает базовые веса энкодера и декодера, поэтому слияние адаптера с ними повредит обе ветви.

Статический кэш и компиляция

Пайплайн выполняет предварительное заполнение энкодера один раз на блок в переиспользуемый кэш (по умолчанию это DynamicCache). Передача cache_implementation="static" использует вместо него StaticCache с фиксированной формой, что позволяет torch.compile декодер с cudagraphs для дополнительного ускорения (пайплайн помечает каждый шаг и клонирует logits, чтобы память cudagraph не была перезаписана); именно такая конфигурация показана в разделе Usage. Чтобы выполнить более простой прогон с динамическим кэшем, удалите и вызов torch.compile, и параметр cache_implementation="static".

Адаптивная остановка

Блок, как правило, сходится до того, как будут израсходованы все num_inference_steps, поэтому по умолчанию конвейер замораживает каждый пример в батче, когда его предсказание argmax остаётся стабильным в течение stability_threshold шагов, а средняя энтропия на токен опускается ниже confidence_threshold (0.005, значение, используемое в опубликованном чекпоинте). Цикл денойзинга завершается, как только все примеры заморожены. Это примерно вдвое сокращает количество проходов декодера при сопоставимом качестве и является главным рычагом повышения пропускной способности. Передайте confidence_threshold=None, чтобы всегда выполнять полный цикл из num_inference_steps:

output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None)  # disable adaptive stopping

Обработчики событий

Обработчики (callbacks) выполняются после каждого шага удаления шума. Передайте callback_on_step_end_tensor_inputs, чтобы выбрать, какие тензоры будут включены в callback_kwargs; доступны canvas (текущие токены блока) и logits. Верните из обработчика {"canvas": ...}, чтобы заменить canvas.

def on_step_end(pipe, step, timestep, callback_kwargs):
    canvas = callback_kwargs["canvas"]
    # Inspect or modify `canvas` here.
    return {"canvas": canvas}

out = pipe(
    prompt="Why is the sky blue?",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["canvas"],
)

DiffusionGemmaPipeline[[diffusers.DiffusionGemmaPipeline]]

diffusers.DiffusionGemmaPipeline[[diffusers.DiffusionGemmaPipeline]]

diffusers.DiffusionGemmaPipeline(model: Any, scheduler: BlockRefinementScheduler | DiscreteDDIMScheduler | EntropyBoundScheduler, processor: Any)

Источник

Параметры:

model (DiffusionGemmaForBlockDiffusion) : блок-диффузионный денуизер (каузальный энкодер + двунаправленный декодер с общими весами).

scheduler (BlockRefinementScheduler, DiscreteDDIMScheduler или EntropyBoundScheduler) : сэмплер, который на каждом шаге денойзинга фиксирует и заново шумит токены холста.

processor (ProcessorMixin) : процессор, используемый для применения шаблона чата и декодирования сгенерированных токенов.

Пайплайн для блочно-диффузионной генерации текста в DiffusionGemma.

DiffusionGemma — это модель encoder-decoder с блочной диффузией: каузальный энкодер считывает чистый промпт (а также ранее сгенерированные блоки) в KV-кэш, а двунаправленный декодер выполняет денуазинг фиксированного «холста» из canvas_length токенов, обращаясь к этому кэшу через cross-attention. Генерация чередует внешний авто-регрессивный цикл по холстам с внутренним циклом денуазинга, на каждом шаге которого отбираются кандидаты токенов, наиболее уверенные из них фиксируются через BlockRefinementScheduler (режим равномерного повреждения, mask_token_id=None), а остальные подвергаются повторному шумовому воздействию.

Ожидается, что модель будет экземпляром DiffusionGemmaForBlockDiffusion, который предоставляет метод forward(input_ids, decoder_input_ids=..., self_conditioning_logits=..., ...) и возвращает логиты формы [batch, canvas_length, vocab_size] для холста. См. карточку модели по адресу https://huggingface.co/google/diffusiongemma-26B-A4B-it.

call[[diffusers.DiffusionGemmaPipeline.call]]

__call__(prompt: str | list[str] | None = None, messages: list[dict] | None = None, image: Any | list[Any] | None = None, add_generation_prompt: bool = True, gen_length: int = 256, num_inference_steps: int = 48, temperature: float = 0.0, cache_implementation: str | None = None, eos_early_stop: bool = True, eos_token_id: int | None = None, stability_threshold: int = 1, confidence_threshold: float | None = 0.005, generator: torch.Generator | None = None, output_type: str = 'text', return_dict: bool = True, callback_on_step_end: Callable[[Any, int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None, callback_on_step_end_tensor_inputs: list[str] | None = None)

Источник

Параметры:

prompt (str или List[str], необязательный) : Текст подсказки, обернутый в шаблон чата и токенизированный процессором. Укажите либо это, либо messages.

messages (List[Dict], необязательный) : Сырой чат-диалог для кодирования, например [{"role": "user", "content": "Hello"}] или многораундовый/мультимодальный диалог. Используйте это вместо prompt для всего, что выходит за рамки одного пользовательского хода.

image (PIL.Image.Image или List, необязательно) : Изображение(я) для связывания с prompt при мультимодальной генерации; процессор преобразует их во входные данные изображения для модели. Для более сложных компоновок помещайте содержимое изображения непосредственно в messages.

add_generation_prompt (bool, по умолчанию True) : добавлять ли промпт генерации при применении шаблона чата.

gen_length (int, по умолчанию 256) : Количество токенов для генерации, округленное вверх до кратного canvas_length модели.

num_inference_steps (int, по умолчанию 48) : Количество шагов денойзинга на холст.

temperature (float, по умолчанию 0.0) : Температура выборки для DiscreteDDIMScheduler/BlockRefinementScheduler (0.0 — жадный алгоритм); EntropyBoundScheduler игнорирует этот параметр и самостоятельно аннелирует свою температуру. Другие параметры выборки (например, top_k, threshold, t_min/t_max) задаются в конфигурации планировщика; установите их на планировщике, например, pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, top_k=...).

cache_implementation (str, необязательный) : Установите значение "static", чтобы один раз на блок заполнить энкодер в постоянный StaticCache и запускать декодер с фиксированными формами, вместо повторного кодирования всей последовательности на каждом шаге. Фиксированные формы также позволяют компилировать декодер, например, pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, fullgraph=True).

eos_early_stop (bool, по умолчанию True) : следует ли прекратить генерацию новых холстов, как только каждая последовательность выдаст EOS.

eos_token_id (int, optional) : ID токена EOS для раннего завершения. Если не задан, используется токенизатор процессора.

stability_threshold (int, по умолчанию 1) : Количество последовательных шагов, в течение которых предсказание argmax для примера должно оставаться неизменным, чтобы этот пример считался стабильным. Используется только при заданном confidence_threshold.

confidence_threshold (float, optional, по умолчанию 0.005) : Замораживает каждый пример, как только он становится стабильным (см. stability_threshold), а средняя энтропия на токен предсказанных логитов, сформированных планировщиком, оказывается ниже этого значения. Цикл денойзинга блока завершается, когда все примеры заморожены. Ускоряет генерацию при сохранении качества; значение по умолчанию соответствует опубликованной контрольной точке. Установите None, чтобы всегда выполнять все num_inference_steps.

generator (torch.Generator, необязательный) : генератор случайных чисел для выборки.

output_type (str, по умолчанию "text") : "text" декодирует последовательности в строки (требуется процессор); "seq" возвращает только идентификаторы токенов.

return_dict (bool, по умолчанию True) : Возвращать ли DiffusionGemmaPipelineOutput вместо кортежа.

callback_on_step_end (Callable или PipelineCallback, необязательный) : Callback, выполняемый после каждого шага денойзинга с сигнатурой callback_on_step_end(self, step, timestep, callback_kwargs). Допустимые ключи тензоров: canvas, logits.

callback_on_step_end_tensor_inputs (List[str], необязательный) : Ключи тензоров, которые нужно передать в обратный вызов.

Возвращает: DiffusionGemmaPipelineOutput или tuple

Идентификаторы сгенерированных токенов (sequences) и, при output_type="text", декодированные texts.

Генерация текста с помощью блочной диффузии.

Примеры:

>>> import torch
>>> from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
>>> from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

>>> model_id = "google/diffusiongemma-26B-A4B-it"
>>> model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
>>> output = pipe(prompt="Why is the sky blue?", gen_length=256)
>>> print(output.texts[0])

DiffusionGemmaPipelineOutput[[diffusers.DiffusionGemmaPipelineOutput]]

diffusers.DiffusionGemmaPipelineOutput[[diffusers.DiffusionGemmaPipelineOutput]]

diffusers.DiffusionGemmaPipelineOutput(sequences: torch.LongTensor, texts: list[str] | None = None)

Источник

Параметры:

sequences (torch.LongTensor формы (batch_size, gen_length)) : идентификаторы сгенерированных токенов (промпт удалён).

texts (list[str], необязательный) : Декодированный текст, одна строка на последовательность. Задается только для output_type="text".

Класс вывода для генерации с блочной диффузией в DiffusionGemma.