# DiffusionGemma

DiffusionGemma — это блочно-диффузионная модель языка с архитектурой «кодировщик-декодировщик». Причинный кодировщик считывает чистый промпт (а также ранее сгенерированные блоки) в KV-кэш, а двунаправленный декодер выполняет денуазинг фиксированного «холста» из `canvas_length` токенов, обращаясь к этому кэшу через кросс-внимание. Генерация чередует внешний авто-регрессивный цикл по холстам с внутренним циклом денуазинга, на каждом шаге которого отбираются кандидаты токенов, наиболее уверенные из них фиксируются через [BlockRefinementScheduler](/docs/diffusers/main/ru/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler) в режиме равномерного повреждения, а остальные подвергаются повторному шумовому воздействию. Сама модель находится в `transformers` в виде `DiffusionGemmaForBlockDiffusion`; опубликованная контрольная точка — [`google/diffusiongemma-26B-A4B-it`](https://huggingface.co/google/diffusiongemma-26B-A4B-it).

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

```py
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 (см. [Статический кэш и компиляция](#static-cache-and-compilation)); для более простого запуска с динамическим кэшем уберите оба параметра.

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

```py
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](/docs/diffusers/main/ru/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler) (по умолчанию): на каждом шаге фиксирует наиболее уверенные токены (те, что выше `threshold`, плюс равномерная квота на каждый шаг) и заново добавляет шум к остальным. Параметр `editing_threshold` дополнительно позволяет переопределять уже зафиксированные токены.
- [DiscreteDDIMScheduler](/docs/diffusers/main/ru/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler): выборка каждой позиции выполняется из точного дискретного апостериорного распределения равномерного процесса повреждения (D3PM). Этот планировщик не имеет параметров, а на последнем шаге предсказанные токены фиксируются детерминированно.
- [EntropyBoundScheduler](/docs/diffusers/main/ru/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler): фиксирует позиции с наименьшей энтропией, совместная энтропия которых не превышает `entropy_bound`, благодаря чему приблизительно независимые токены принимаются вместе. Температура выборки аннелируется от `t_max` (`0.8`) на первом шаге до `t_min` (`0.4`) на последнем, что соответствует сэмплеру опубликованной контрольной точки.

```py
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` для уточнения блока, граница энтропии, ...) задаются в конфигурации планировщика:

```py
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](https://huggingface.co/papers/2605.22765). Он уточняет холст с помощью `corrector_steps` шагов Гиббса, которые пересэмплируют позиции с наименьшей уверенностью из однокоординатного условного распределения шумного маргинального распределения, что оставляет это маргинальное распределение инвариантным и улучшает генерацию без дополнительных затрат на обучение. Он работает непосредственно с опубликованной контрольной точкой: для uniform diffusion денойзер и leave-one-out апостериорное распределение взаимозаменяемы в замкнутой форме, поэтому корректор восстанавливает необходимые ему leave-one-out величины без какого-либо переобучения.

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

```py
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](https://huggingface.co/docs/peft), а не через API `load_lora_weights` diffusers. Поскольку эта интеграция не зависит от типа адаптера, те же вызовы загружают LoRA, DoRA или любой другой адаптер PEFT (например, результат работы `SFTTrainer` из TRL). Управляйте адаптерами непосредственно на компоненте модели:

```py
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](#usage). Чтобы выполнить более простой прогон с динамическим кэшем, удалите и вызов `torch.compile`, и параметр `cache_implementation="static"`.

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

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

```py
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.

```py
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]]

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

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L53)

**Параметры:**

model ([DiffusionGemmaForBlockDiffusion](https://huggingface.co/docs/transformers/main/en/model_doc/diffusion_gemma#transformers.DiffusionGemmaForBlockDiffusion)) : блок-диффузионный денуизер (каузальный энкодер + двунаправленный декодер с общими весами).

scheduler ([BlockRefinementScheduler](/docs/diffusers/main/ru/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler), [DiscreteDDIMScheduler](/docs/diffusers/main/ru/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler) или [EntropyBoundScheduler](/docs/diffusers/main/ru/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler)) : сэмплер, который на каждом шаге денойзинга фиксирует и заново шумит токены холста.

processor ([ProcessorMixin](https://huggingface.co/docs/transformers/main/en/main_classes/processors#transformers.ProcessorMixin)) : процессор, используемый для применения шаблона чата и декодирования сгенерированных токенов.

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

DiffusionGemma — это модель encoder-decoder с блочной диффузией: каузальный энкодер считывает чистый промпт (а также ранее сгенерированные блоки) в KV-кэш, а двунаправленный декодер выполняет денуазинг фиксированного «холста» из `canvas_length` токенов, обращаясь к этому кэшу через cross-attention. Генерация чередует внешний авто-регрессивный цикл по холстам с внутренним циклом денуазинга, на каждом шаге которого отбираются кандидаты токенов, наиболее уверенные из них фиксируются через [BlockRefinementScheduler](/docs/diffusers/main/ru/api/schedulers/block_refinement#diffusers.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__]]

```python
__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)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L163)

**Параметры:**

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](/docs/diffusers/main/ru/api/pipelines/diffusion_gemma#diffusers.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](/docs/diffusers/main/ru/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput) или `tuple`

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

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

Примеры:
```python
>>> 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]]

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

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/diffusion_gemma/pipeline_output.py#L25)

**Параметры:**

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

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

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