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

DiffusionGemma

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

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

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

スケジューラー

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

• 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 は、Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation で提案される 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 と同じになります。

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統合を通じて読み込まれます。この統合はアダプターの種類に依存しないため、同じ呼び出しでLoRA、DoRA、またはその他のPEFTアダプター（例: TRLのSFTTrainerの出力）を読み込むことができます。モデルコンポーネント上で直接アダプターを管理してください:

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に示されている設定です。よりシンプルな動的キャッシュでの実行を行う場合は、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

コールバック

Callback は各デノイジングステップの後に実行されます。callback_on_step_end_tensor_inputs を渡すことで、callback_kwargs に含めるテンソルを選択できます。canvas（現在のブロックトークン）と logits が利用可能です。Callback から {"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のブロック拡散テキスト生成用のPipeline。

DiffusionGemma はブロック拡散型エンコーダー・デコーダーモデルです。因果エンコーダーがクリーンなプロンプト（および以前に生成されたブロック）を読み取り KV キャッシュに格納し、双方向デコーダーがそのキャッシュへのクロスアテンションを用いて、canvas_length トークンからなる固定サイズの「キャンバス」をデノイズします。生成は、キャンバスをまたぐ外側の自己回帰ループと、内側のデノイズループを交互に実行します。デノイズループの各ステップでは、候補トークンをサンプリングし、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]]

__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, 任意) : 早期停止用の 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を返すかどうか。

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 または tuple

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

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

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