# DiffusionGemma

DiffusionGemma は、ブロック拡散型エンコーダー・デコーダー言語モデルです。因果エンコーダーがクリーンなプロンプト（およびそれまでに生成されたブロック）を KV キャッシュに読み込み、双方向デコーダーがそのキャッシュへのクロスアテンションを通じて、`canvas_length` トークンからなる固定サイズの「キャンバス」をノイズ除去します。生成は、キャンバスに対する外側の自己回帰ループと、内側のノイズ除去ループを交互に実行します。各ステップでは候補トークンをサンプリングし、[BlockRefinementScheduler](/docs/diffusers/main/ja/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler) を使用して uniform corruption モードで最も確信度の高いトークンを確定し、残りのトークンには再びノイズを付与します。モデル本体は `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"` を指定すると、デコーダーを cudagraphs を使用して `torch.compile` できます（[Static cache and compilation](#static-cache-and-compilation) を参照）。よりシンプルな動的キャッシュでの実行を行う場合は、両方を省略してください。

マルチターンまたはマルチモーダルの入力では、`prompt` の代わりに生の `messages` 会話を渡します。これは通常のチャット形式の `{"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)
```

1回のユーザーターンでは、`messages`を省略し、`prompt`と並んで`image`を渡すことができます。プロセッサはこれを自動的にモデルの画像入力に変換します。

## スケジューラー

スケジューラーは、各キャンバスをデノイズするサンプラーであり、交換可能です。スケジューラーを切り替えるだけで、その他の部分を変更せずにサンプリング戦略を変更できます。利用可能なスケジューラーは3種類あります。

- [BlockRefinementScheduler](/docs/diffusers/main/ja/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler)（デフォルト）: 各ステップで最も確信度の高いトークン（`threshold`を超えるもの、およびステップごとの均等な割り当て分）を確定し、残りのトークンには再ノイズを適用します。`editing_threshold` を設定すると、すでに確定済みのトークンも再編集できるようになります。
- [DiscreteDDIMScheduler](/docs/diffusers/main/ja/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler): 一様破損プロセス（D3PM）の正確な離散事後分布から各位置をサンプリングします。パラメータを必要とせず、最終ステップでは予測されたトークンを決定論的に確定します。
- [EntropyBoundScheduler](/docs/diffusers/main/ja/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` は、[Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation](https://huggingface.co/papers/2605.22765) で提案される leave-one-out predictor-corrector をサポートしています。この手法は、`corrector_steps` 回の Gibbs sweep を実行して、ノイズ付き周辺分布の 1 座標条件分布から最も確信度の低い位置を再サンプリングすることでキャンバスを洗練させます。これにより、その周辺分布は不変のまま、追加の学習コストなしで生成品質が向上します。公開されているチェックポイントに直接適用できます。Uniform diffusion の場合、denoiser と leave-one-out 事後分布は閉形式で互換性があるため、corrector は再学習なしに必要な leave-one-out 量を復元できます。

corrector のスウィープは `num_inference_steps` の予算に組み込まれ、上乗せされるわけではありません。パイプラインは predictor のステップ数を減らし、解放されたフォワードを corrector に割り当てるため、モデルのフォワード総数は `num_inference_steps` のままであり、predictor-corrector のコストは通常の ancestral sampling と同じになります。

```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モデルであるため、アダプターはdiffusersの`load_lora_weights` APIではなく、そのネイティブな[PEFT](https://huggingface.co/docs/peft)統合を通じて読み込まれます。この統合はアダプターの種類に依存しないため、同じ呼び出しでLoRA、DoRA、またはその他のPEFTアダプター（例: TRLの`SFTTrainer`の出力）を読み込むことができます。モデルコンポーネント上で直接アダプターを管理してください:

```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")
```

Adapters はマージされず、アクティブな状態のまま保持されます。DiffusionGemma はエンコーダーとデコーダーのベース重みを結合しているため、アダプターをこれらに融合すると、両方のブランチが破損します。

## 静的キャッシュとコンパイル

パイプラインは、ブロックごとにエンコーダーを1回だけ事前計算し、再利用可能なキャッシュ（デフォルトでは`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
```

## コールバック

Callback は各デノイジングステップの後に実行されます。`callback_on_step_end_tensor_inputs` を渡すことで、`callback_kwargs` に含めるテンソルを選択できます。`canvas`（現在のブロックトークン）と `logits` が利用可能です。Callback から `{"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/ja/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler)、[DiscreteDDIMScheduler](/docs/diffusers/main/ja/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler) または [EntropyBoundScheduler](/docs/diffusers/main/ja/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler)) : 各デノイジングステップでキャンバスのトークンをコミットし、再ノイズ化するサンプラー。

processor ([ProcessorMixin](https://huggingface.co/docs/transformers/main/en/main_classes/processors#transformers.ProcessorMixin)) : チャットテンプレートを適用し、生成されたトークンをデコードするために使用されるプロセッサ。

DiffusionGemmaのブロック拡散テキスト生成用のPipeline。

DiffusionGemma はブロック拡散型エンコーダー・デコーダーモデルです。因果エンコーダーがクリーンなプロンプト（および以前に生成されたブロック）を読み取り KV キャッシュに格納し、双方向デコーダーがそのキャッシュへのクロスアテンションを用いて、`canvas_length` トークンからなる固定サイズの「キャンバス」をデノイズします。生成は、キャンバスをまたぐ外側の自己回帰ループと、内側のデノイズループを交互に実行します。デノイズループの各ステップでは、候補トークンをサンプリングし、[BlockRefinementScheduler](/docs/diffusers/main/ja/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]` の logits を返す必要があります。モデルカードは 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`, *任意*) : 早期停止用の EOS トークン ID。プロセッサのトークナイザーにフォールバックします。

stability_threshold (`int`、デフォルトは`1`) : あるサンプルのargmax予測が、そのサンプルを安定とみなすために、何ステップ連続で変化しない必要があるかを指定します。`confidence_threshold`が設定されている場合のみ使用されます。

confidence_threshold (`float`、*任意*、デフォルトは `0.005`) : 各サンプルが安定した時点で（`stability_threshold` を参照）、そのスケジューラー形状の予測 logits のトークンごとの平均エントロピーがこの値を下回った場合に、そのサンプルをフリーズします。すべてのサンプルがフリーズされると、ブロックのデノイジングループが終了します。同等の品質を保ちつつ生成を高速化します。デフォルト値は公開されたチェックポイントと一致します。常にすべての `num_inference_steps` を実行するには `None` を設定してください。

generator (`torch.Generator`, *任意*) : サンプリング用の乱数生成器。

output_type (`str`、デフォルトは`"text"`) : `"text"`はシーケンスを文字列にデコードします（プロセッサが必要です）。`"seq"`はトークンIDのみを返します。

return_dict (`bool`、デフォルトは`True`) : タプルではなく[DiffusionGemmaPipelineOutput](/docs/diffusers/main/ja/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput)を返すかどうか。

callback_on_step_end (`Callable` または `PipelineCallback`、*任意*) : 各デノイジングステップの後に実行されるコールバックで、シグネチャは `callback_on_step_end(self, step, timestep, callback_kwargs)` です。許可されるテンソルキー: `canvas`、`logits`。

callback_on_step_end_tensor_inputs (`List[str]`、*任意*) : コールバックに渡すテンソルのキー。

**戻り値:** [DiffusionGemmaPipelineOutput](/docs/diffusers/main/ja/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput) または `tuple`

生成されたトークン ID(`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`、shape `(batch_size, gen_length)`) : 生成されたトークン ID（プロンプトは除去されています）。

texts (`list[str]`、*任意*) : 復号化されたテキスト。シーケンスごとに1文字列です。`output_type="text"`の場合のみ設定されます。

DiffusionGemmaのブロック拡散生成用の出力クラス。