diffusers/pr_14178/en/_app/immutable/nodes/148.C7HY06oa.js

import"../chunks/DsnmJJEf.js";import{i as A,h as H,C as z,H as s,a as o,D as i,E as P,s as F}from"../chunks/BtE7mKSK.js";import{p as Y,o as L,s as e,f as V,a as c,b as O,c as r,d as p,r as d,n as E}from"../chunks/jDjavuwI.js";import{E as K}from"../chunks/SrSJA0zO.js";const $='{"title":"DiffusionGemma","local":"diffusiongemma","sections":[{"title":"Usage","local":"usage","sections":[],"depth":2},{"title":"Schedulers","local":"schedulers","sections":[{"title":"Predictor-corrector sampling","local":"predictor-corrector-sampling","sections":[],"depth":3}],"depth":2},{"title":"PEFT adapters","local":"peft-adapters","sections":[],"depth":2},{"title":"Static cache and compilation","local":"static-cache-and-compilation","sections":[],"depth":2},{"title":"Adaptive stopping","local":"adaptive-stopping","sections":[],"depth":2},{"title":"Callbacks","local":"callbacks","sections":[],"depth":2},{"title":"DiffusionGemmaPipeline","local":"diffusers.DiffusionGemmaPipeline","sections":[],"depth":2},{"title":"DiffusionGemmaPipelineOutput","local":"diffusers.DiffusionGemmaPipelineOutput","sections":[],"depth":2}],"depth":1}';var ee=p('<meta name="hf:doc:metadata"/>'),se=p("<p>Examples:</p> <!>",1),oe=p(`<p></p> <!> <!> <p>DiffusionGemma is a block-diffusion encoder-decoder language model. A causal encoder reads the clean prompt (and any
previously generated blocks) into a KV cache, and a bidirectional decoder denoises a fixed-size “canvas” of <code>canvas_length</code> tokens by cross-attending to that cache. Generation alternates an outer autoregressive loop over
canvases with an inner denoising loop, where each step samples candidate tokens, commits the most confident ones via <a href="/docs/diffusers/pr_14178/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler">BlockRefinementScheduler</a> in uniform corruption mode, and renoises the rest. The model itself lives in <code>transformers</code> as <code>DiffusionGemmaForBlockDiffusion</code>; the released checkpoint is <a href="https://huggingface.co/google/diffusiongemma-26B-A4B-it" rel="nofollow"><code>google/diffusiongemma-26B-A4B-it</code></a>.</p> <!> <!> <p><code>num_inference_steps</code> is the number of denoising steps per canvas (48 matches the released checkpoint); fewer steps are
faster but lower quality. <code>cache_implementation="static"</code> lets the decoder be <code>torch.compile</code>-d with cudagraphs (see <a href="#static-cache-and-compilation">Static cache and compilation</a>); drop both for a simpler dynamic-cache run.</p> <p>For multi-turn or multimodal inputs, pass a raw <code>messages</code> conversation instead of <code>prompt</code>. It is a list of <code>&#123;"role", "content"&#125;</code> dicts in the usual chat format, which the processor runs through its chat template:</p> <!> <p>For a single user turn you can skip <code>messages</code> and pass an <code>image</code> alongside the <code>prompt</code>; the processor turns it into
the model’s image inputs automatically.</p> <!> <p>The scheduler is the sampler that denoises each canvas, and it is interchangeable: swap it to change the sampling
strategy without touching anything else. Three schedulers are available:</p> <ul><li><a href="/docs/diffusers/pr_14178/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler">BlockRefinementScheduler</a> (default): commits the most confident tokens each step (above <code>threshold</code>, plus an even
per-step quota) and renoises the rest. <code>editing_threshold</code> additionally lets it re-edit already committed tokens.</li> <li><a href="/docs/diffusers/pr_14178/en/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler">DiscreteDDIMScheduler</a>: samples each position from the exact discrete posterior of the uniform corruption process
(D3PM). It is parameter free, and the final step deterministically commits the predicted tokens.</li> <li><a href="/docs/diffusers/pr_14178/en/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler">EntropyBoundScheduler</a>: commits the lowest-entropy positions whose joint entropy stays under <code>entropy_bound</code>, so
roughly independent tokens are accepted together. It anneals its sampling temperature from <code>t_max</code> (<code>0.8</code>) on the
first step down to <code>t_min</code> (<code>0.4</code>) on the last, matching the released checkpoint’s sampler.</li></ul> <!> <p>Scheduler-specific sampling knobs (the block-refinement <code>threshold</code>/<code>top_k</code>, the entropy bound, …) are set on the
scheduler config:</p> <!> <p><code>EntropyBoundScheduler</code> anneals its sampling temperature (<code>t_max</code>/<code>t_min</code>) internally over the denoising steps; <code>DiscreteDDIMScheduler</code> and <code>BlockRefinementScheduler</code> use the flat <code>temperature</code> passed to the pipeline (<code>0.0</code> for
greedy).</p> <!> <p><code>DiscreteDDIMScheduler</code> supports the leave-one-out predictor-corrector of <a href="https://huggingface.co/papers/2605.22765" rel="nofollow">Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation</a>. It refines the canvas with <code>corrector_steps</code> Gibbs sweeps that resample the least-confident positions from the one-coordinate conditional of the noisy marginal, which leaves that marginal invariant and improves generation at no extra training cost. It works directly on the released checkpoint: for uniform diffusion the denoiser and the leave-one-out posterior are interchangeable in closed form, so the corrector recovers the leave-one-out quantities it needs without any retraining.</p> <p>The corrector sweeps are folded into the <code>num_inference_steps</code> budget rather than added on top: the pipeline runs fewer predictor steps and spends the freed forwards on correctors, so the total number of model forwards stays <code>num_inference_steps</code> and the predictor-corrector costs the same as plain ancestral sampling.</p> <!> <!> <p>The denoiser is a 🤗 Transformers model, so adapters are loaded through its native <a href="https://huggingface.co/docs/peft" rel="nofollow">PEFT</a> integration rather than the diffusers <code>load_lora_weights</code> API. Because that integration is adapter-type-agnostic, the same calls load LoRA, DoRA, or any other PEFT adapter (e.g. the output of TRL’s <code>SFTTrainer</code>). Manage adapters on the model component directly:</p> <!> <p>Adapters stay active and unmerged: DiffusionGemma ties the encoder and decoder base weights, so fusing an adapter into them would corrupt both branches.</p> <!> <p>The pipeline prefills the encoder once per block into a reusable cache (a <code>DynamicCache</code> by default). Passing <code>cache_implementation="static"</code> uses a fixed-shape <code>StaticCache</code> instead, whose shapes let you <code>torch.compile</code> the
decoder with cudagraphs for a further speedup (the pipeline marks each step and clones the logits so cudagraph memory
is not overwritten); this is the setup shown in <a href="#usage">Usage</a>. Drop both the <code>torch.compile</code> call and <code>cache_implementation="static"</code> for a simpler dynamic-cache run.</p> <!> <p>A block usually converges before all <code>num_inference_steps</code> are spent, so by default the pipeline leaves a block’s
denoising loop early once every example’s argmax prediction is stable for <code>stability_threshold</code> steps and the mean
per-token entropy falls below <code>confidence_threshold</code> (<code>0.005</code>, the value used by the released checkpoint). This roughly
halves the number of decoder forwards at matched quality and is the largest single throughput lever. Pass <code>confidence_threshold=None</code> to always run the full <code>num_inference_steps</code>:</p> <!> <!> <p>Callbacks run after each denoising step. Pass <code>callback_on_step_end_tensor_inputs</code> to select which tensors are
included in <code>callback_kwargs</code>; <code>canvas</code> (the current block tokens) and <code>logits</code> are available. Return <code>&#123;"canvas": ...&#125;</code> from the callback to replace the canvas.</p> <!> <!> <div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8"><!> <p>Pipeline for DiffusionGemma block-diffusion text generation.</p> <p>DiffusionGemma is a block-diffusion encoder-decoder model: a causal encoder reads the clean prompt (and any
previously generated blocks) into a KV cache, and a bidirectional decoder denoises a fixed-size “canvas” of <code>canvas_length</code> tokens by cross-attending to that cache. Generation alternates an outer autoregressive loop over
canvases with an inner denoising loop, where each step samples candidate tokens, commits the most confident ones
via <a href="/docs/diffusers/pr_14178/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler">BlockRefinementScheduler</a> (uniform corruption mode, <code>mask_token_id=None</code>), and renoises the rest.</p> <p>The model is expected to be a <code>DiffusionGemmaForBlockDiffusion</code> instance exposing <code>forward(input_ids, decoder_input_ids=..., self_conditioning_logits=..., ...)</code> and returning logits of shape <code>[batch, canvas_length, vocab_size]</code> over the canvas. See the model card at <a href="https://huggingface.co/google/diffusiongemma-26B-A4B-it" rel="nofollow">https://huggingface.co/google/diffusiongemma-26B-A4B-it</a>.</p> <div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8"><!> <p>Generate text with block diffusion.</p> <!></div></div> <!> <div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8"><!> <p>Output class for DiffusionGemma block-diffusion generation.</p></div> <!> <p></p>`,1);function ie(S,x){Y(x,!1),L(()=>{new URLSearchParams(window.location.search).get("fw")}),A();var u=oe();H("2vcgqd",t=>{var l=ee();F(l,"content",$),c(t,l)});var m=e(V(u),2);z(m,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"});var h=e(m,2);s(h,{title:"DiffusionGemma",local:"diffusiongemma",headingTag:"h1"});var f=e(h,4);s(f,{title:"Usage",local:"usage",headingTag:"h2"});var g=e(f,2);o(g,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwdHJhbnNmb3JtZXJzJTIwaW1wb3J0JTIwQXV0b1Byb2Nlc3NvciUyQyUyMERpZmZ1c2lvbkdlbW1hRm9yQmxvY2tEaWZmdXNpb24lMEElMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwQmxvY2tSZWZpbmVtZW50U2NoZWR1bGVyJTJDJTIwRGlmZnVzaW9uR2VtbWFQaXBlbGluZSUwQSUwQW1vZGVsX2lkJTIwJTNEJTIwJTIyZ29vZ2xlJTJGZGlmZnVzaW9uZ2VtbWEtMjZCLUE0Qi1pdCUyMiUwQW1vZGVsJTIwJTNEJTIwRGlmZnVzaW9uR2VtbWFGb3JCbG9ja0RpZmZ1c2lvbi5mcm9tX3ByZXRyYWluZWQobW9kZWxfaWQlMkMlMjBkdHlwZSUzRHRvcmNoLmJmbG9hdDE2JTJDJTIwZGV2aWNlX21hcCUzRCUyMmF1dG8lMjIpJTBBcHJvY2Vzc29yJTIwJTNEJTIwQXV0b1Byb2Nlc3Nvci5mcm9tX3ByZXRyYWluZWQobW9kZWxfaWQpJTBBc2NoZWR1bGVyJTIwJTNEJTIwQmxvY2tSZWZpbmVtZW50U2NoZWR1bGVyKCklMEElMEFwaXBlJTIwJTNEJTIwRGlmZnVzaW9uR2VtbWFQaXBlbGluZShtb2RlbCUzRG1vZGVsJTJDJTIwc2NoZWR1bGVyJTNEc2NoZWR1bGVyJTJDJTIwcHJvY2Vzc29yJTNEcHJvY2Vzc29yKSUwQXBpcGUubW9kZWwubW9kZWwuZGVjb2RlciUyMCUzRCUyMHRvcmNoLmNvbXBpbGUocGlwZS5tb2RlbC5tb2RlbC5kZWNvZGVyJTJDJTIwbW9kZSUzRCUyMnJlZHVjZS1vdmVyaGVhZCUyMiUyQyUyMGZ1bGxncmFwaCUzRFRydWUpJTBBb3V0cHV0JTIwJTNEJTIwcGlwZSglMEElMjAlMjAlMjAlMjBwcm9tcHQlM0QlMjJXaHklMjBpcyUyMHRoZSUyMHNreSUyMGJsdWUlM0YlMjIlMkMlMEElMjAlMjAlMjAlMjBnZW5fbGVuZ3RoJTNEMjU2JTJDJTBBJTIwJTIwJTIwJTIwbnVtX2luZmVyZW5jZV9zdGVwcyUzRDQ4JTJDJTBBJTIwJTIwJTIwJTIwY2FjaGVfaW1wbGVtZW50YXRpb24lM0QlMjJzdGF0aWMlMjIlMkMlMEEpJTBBcHJpbnQob3V0cHV0LnRleHRzJTVCMCU1RCk=",highlighted:`<span class="hljs-keyword">import</span> torch
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoProcessor, DiffusionGemmaForBlockDiffusion

<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = <span class="hljs-string">&quot;google/diffusiongemma-26B-A4B-it&quot;</span>
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map=<span class="hljs-string">&quot;auto&quot;</span>)
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.<span class="hljs-built_in">compile</span>(pipe.model.model.decoder, mode=<span class="hljs-string">&quot;reduce-overhead&quot;</span>, fullgraph=<span class="hljs-literal">True</span>)
output = pipe(
    prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>,
    gen_length=<span class="hljs-number">256</span>,
    num_inference_steps=<span class="hljs-number">48</span>,
    cache_implementation=<span class="hljs-string">&quot;static&quot;</span>,
)
<span class="hljs-built_in">print</span>(output.texts[<span class="hljs-number">0</span>])`,lang:"py",wrap:!1});var y=e(g,6);o(y,{code:"bWVzc2FnZXMlMjAlM0QlMjAlNUIlMEElMjAlMjAlMjAlMjAlN0IlMjJyb2xlJTIyJTNBJTIwJTIydXNlciUyMiUyQyUyMCUyMmNvbnRlbnQlMjIlM0ElMjAlMjJXaHklMjBpcyUyMHRoZSUyMHNreSUyMGJsdWUlM0YlMjIlN0QlMkMlMEElNUQlMEElMjMlMjBvciUyMHdpdGglMjBhbiUyMGltYWdlJTNBJTBBbWVzc2FnZXMlMjAlM0QlMjAlNUIlMEElMjAlMjAlMjAlMjAlN0IlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJyb2xlJTIyJTNBJTIwJTIydXNlciUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMmNvbnRlbnQlMjIlM0ElMjAlNUIlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlN0IlMjJ0eXBlJTIyJTNBJTIwJTIyaW1hZ2UlMjIlMkMlMjAlMjJpbWFnZSUyMiUzQSUyMGltYWdlJTdEJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTdCJTIydHlwZSUyMiUzQSUyMCUyMnRleHQlMjIlMkMlMjAlMjJ0ZXh0JTIyJTNBJTIwJTIyRGVzY3JpYmUlMjB0aGlzJTIwaW1hZ2UuJTIyJTdEJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTVEJTJDJTBBJTIwJTIwJTIwJTIwJTdEJTJDJTBBJTVEJTBBb3V0cHV0JTIwJTNEJTIwcGlwZShtZXNzYWdlcyUzRG1lc3NhZ2VzJTJDJTIwZ2VuX2xlbmd0aCUzRDI1Nik=",highlighted:`messages = [
    {<span class="hljs-string">&quot;role&quot;</span>: <span class="hljs-string">&quot;user&quot;</span>, <span class="hljs-string">&quot;content&quot;</span>: <span class="hljs-string">&quot;Why is the sky blue?&quot;</span>},
]
<span class="hljs-comment"># or with an image:</span>
messages = [
    {
        <span class="hljs-string">&quot;role&quot;</span>: <span class="hljs-string">&quot;user&quot;</span>,
        <span class="hljs-string">&quot;content&quot;</span>: [
            {<span class="hljs-string">&quot;type&quot;</span>: <span class="hljs-string">&quot;image&quot;</span>, <span class="hljs-string">&quot;image&quot;</span>: image},
            {<span class="hljs-string">&quot;type&quot;</span>: <span class="hljs-string">&quot;text&quot;</span>, <span class="hljs-string">&quot;text&quot;</span>: <span class="hljs-string">&quot;Describe this image.&quot;</span>},
        ],
    },
]
output = pipe(messages=messages, gen_length=<span class="hljs-number">256</span>)`,lang:"py",wrap:!1});var _=e(y,4);s(_,{title:"Schedulers",local:"schedulers",headingTag:"h2"});var b=e(_,6);o(b,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMERpc2NyZXRlRERJTVNjaGVkdWxlciUyQyUyMEVudHJvcHlCb3VuZFNjaGVkdWxlciUwQSUwQXBpcGUuc2NoZWR1bGVyJTIwJTNEJTIwRGlzY3JldGVERElNU2NoZWR1bGVyKCklMEElMjMlMjBvciUzQSUyMHBpcGUuc2NoZWR1bGVyJTIwJTNEJTIwRW50cm9weUJvdW5kU2NoZWR1bGVyKGVudHJvcHlfYm91bmQlM0QwLjEpJTBBb3V0cHV0JTIwJTNEJTIwcGlwZShwcm9tcHQlM0QlMjJXaHklMjBpcyUyMHRoZSUyMHNreSUyMGJsdWUlM0YlMjIlMkMlMjBnZW5fbGVuZ3RoJTNEMjU2JTJDJTIwbnVtX2luZmVyZW5jZV9zdGVwcyUzRDQ4KSUwQXByaW50KG91dHB1dC50ZXh0cyU1QjAlNUQp",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
<span class="hljs-comment"># or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)</span>
output = pipe(prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>, gen_length=<span class="hljs-number">256</span>, num_inference_steps=<span class="hljs-number">48</span>)
<span class="hljs-built_in">print</span>(output.texts[<span class="hljs-number">0</span>])`,lang:"py",wrap:!1});var M=e(b,4);o(M,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMEJsb2NrUmVmaW5lbWVudFNjaGVkdWxlciUwQSUwQXBpcGUuc2NoZWR1bGVyJTIwJTNEJTIwQmxvY2tSZWZpbmVtZW50U2NoZWR1bGVyLmZyb21fY29uZmlnKHBpcGUuc2NoZWR1bGVyLmNvbmZpZyUyQyUyMHRocmVzaG9sZCUzRDAuOSk=",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=<span class="hljs-number">0.9</span>)`,lang:"py",wrap:!1});var J=e(M,4);s(J,{title:"Predictor-corrector sampling",local:"predictor-corrector-sampling",headingTag:"h3"});var T=e(J,6);o(T,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMERpc2NyZXRlRERJTVNjaGVkdWxlciUwQSUwQXBpcGUuc2NoZWR1bGVyJTIwJTNEJTIwRGlzY3JldGVERElNU2NoZWR1bGVyKGNvcnJlY3Rvcl9zdGVwcyUzRDIlMkMlMjBjb3JyZWN0b3JfayUzRDEyKSUwQW91dHB1dCUyMCUzRCUyMHBpcGUocHJvbXB0JTNEJTIyV2h5JTIwaXMlMjB0aGUlMjBza3klMjBibHVlJTNGJTIyJTJDJTIwZ2VuX2xlbmd0aCUzRDI1NiUyQyUyMG51bV9pbmZlcmVuY2Vfc3RlcHMlM0Q0OCklMEFwcmludChvdXRwdXQudGV4dHMlNUIwJTVEKQ==",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=<span class="hljs-number">2</span>, corrector_k=<span class="hljs-number">12</span>)
output = pipe(prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>, gen_length=<span class="hljs-number">256</span>, num_inference_steps=<span class="hljs-number">48</span>)
<span class="hljs-built_in">print</span>(output.texts[<span class="hljs-number">0</span>])`,lang:"py",wrap:!1});var w=e(T,2);s(w,{title:"PEFT adapters",local:"peft-adapters",headingTag:"h2"});var j=e(w,4);o(j,{code:"cGlwZS5tb2RlbC5sb2FkX2FkYXB0ZXIoJTIycGF0aCUyRnRvJTJGYWRhcHRlciUyMiUyQyUyMGFkYXB0ZXJfbmFtZSUzRCUyMnNmdCUyMiklMjAlMjAlMjMlMjBMb1JBJTJDJTIwRG9SQSUyQyUyMC4uLiUwQXBpcGUubW9kZWwuc2V0X2FkYXB0ZXIoJTIyc2Z0JTIyKSUwQW91dHB1dCUyMCUzRCUyMHBpcGUocHJvbXB0JTNEJTIyV2h5JTIwaXMlMjB0aGUlMjBza3klMjBibHVlJTNGJTIyJTJDJTIwZ2VuX2xlbmd0aCUzRDI1NiklMEElMEFwaXBlLm1vZGVsLmRpc2FibGVfYWRhcHRlcnMoKSUyMCUyMCUyMyUyMHJ1biUyMHRoZSUyMGJhc2UlMjBtb2RlbCUwQXBpcGUubW9kZWwuZGVsZXRlX2FkYXB0ZXIoJTIyc2Z0JTIyKQ==",highlighted:`pipe.model.load_adapter(<span class="hljs-string">&quot;path/to/adapter&quot;</span>, adapter_name=<span class="hljs-string">&quot;sft&quot;</span>)  <span class="hljs-comment"># LoRA, DoRA, ...</span>
pipe.model.set_adapter(<span class="hljs-string">&quot;sft&quot;</span>)
output = pipe(prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>, gen_length=<span class="hljs-number">256</span>)

pipe.model.disable_adapters()  <span class="hljs-comment"># run the base model</span>
pipe.model.delete_adapter(<span class="hljs-string">&quot;sft&quot;</span>)`,lang:"py",wrap:!1});var U=e(j,4);s(U,{title:"Static cache and compilation",local:"static-cache-and-compilation",headingTag:"h2"});var G=e(U,4);s(G,{title:"Adaptive stopping",local:"adaptive-stopping",headingTag:"h2"});var k=e(G,4);o(k,{code:"b3V0cHV0JTIwJTNEJTIwcGlwZShwcm9tcHQlM0QlMjJXaHklMjBpcyUyMHRoZSUyMHNreSUyMGJsdWUlM0YlMjIlMkMlMjBnZW5fbGVuZ3RoJTNEMjU2JTJDJTIwY29uZmlkZW5jZV90aHJlc2hvbGQlM0ROb25lKSUyMCUyMCUyMyUyMGRpc2FibGUlMjBhZGFwdGl2ZSUyMHN0b3BwaW5n",highlighted:'output = pipe(prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>, gen_length=<span class="hljs-number">256</span>, confidence_threshold=<span class="hljs-literal">None</span>)  <span class="hljs-comment"># disable adaptive stopping</span>',lang:"py",wrap:!1});var v=e(k,2);s(v,{title:"Callbacks",local:"callbacks",headingTag:"h2"});var I=e(v,4);o(I,{code:"ZGVmJTIwb25fc3RlcF9lbmQocGlwZSUyQyUyMHN0ZXAlMkMlMjB0aW1lc3RlcCUyQyUyMGNhbGxiYWNrX2t3YXJncyklM0ElMEElMjAlMjAlMjAlMjBjYW52YXMlMjAlM0QlMjBjYWxsYmFja19rd2FyZ3MlNUIlMjJjYW52YXMlMjIlNUQlMEElMjAlMjAlMjAlMjAlMjMlMjBJbnNwZWN0JTIwb3IlMjBtb2RpZnklMjAlNjBjYW52YXMlNjAlMjBoZXJlLiUwQSUyMCUyMCUyMCUyMHJldHVybiUyMCU3QiUyMmNhbnZhcyUyMiUzQSUyMGNhbnZhcyU3RCUwQSUwQSUwQW91dCUyMCUzRCUyMHBpcGUoJTBBJTIwJTIwJTIwJTIwcHJvbXB0JTNEJTIyV2h5JTIwaXMlMjB0aGUlMjBza3klMjBibHVlJTNGJTIyJTJDJTBBJTIwJTIwJTIwJTIwY2FsbGJhY2tfb25fc3RlcF9lbmQlM0Rvbl9zdGVwX2VuZCUyQyUwQSUyMCUyMCUyMCUyMGNhbGxiYWNrX29uX3N0ZXBfZW5kX3RlbnNvcl9pbnB1dHMlM0QlNUIlMjJjYW52YXMlMjIlNUQlMkMlMEEp",highlighted:`<span class="hljs-keyword">def</span> <span class="hljs-title function_">on_step_end</span>(<span class="hljs-params">pipe, step, timestep, callback_kwargs</span>):
    canvas = callback_kwargs[<span class="hljs-string">&quot;canvas&quot;</span>]
    <span class="hljs-comment"># Inspect or modify \`canvas\` here.</span>
    <span class="hljs-keyword">return</span> {<span class="hljs-string">&quot;canvas&quot;</span>: canvas}


out = pipe(
    prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>,
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=[<span class="hljs-string">&quot;canvas&quot;</span>],
)`,lang:"py",wrap:!1});var Z=e(I,2);s(Z,{title:"DiffusionGemmaPipeline",local:"diffusers.DiffusionGemmaPipeline",headingTag:"h2"});var n=e(Z,2),B=r(n);i(B,{name:"class diffusers.DiffusionGemmaPipeline",anchor:"diffusers.DiffusionGemmaPipeline",source:"https://github.com/huggingface/diffusers/blob/vr_14178/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L53",parameters:[{name:"model",val:": Any"},{name:"scheduler",val:": BlockRefinementScheduler | DiscreteDDIMScheduler | EntropyBoundScheduler"},{name:"processor",val:": Any"}],parametersDescription:[{anchor:"diffusers.DiffusionGemmaPipeline.model",description:`<strong>model</strong> (<a href="https://huggingface.co/docs/transformers/main/en/model_doc/diffusion_gemma#transformers.DiffusionGemmaForBlockDiffusion" rel="nofollow">DiffusionGemmaForBlockDiffusion</a>) &#x2014;
The block-diffusion denoiser (causal encoder + bidirectional decoder with tied weights).`,name:"model"},{anchor:"diffusers.DiffusionGemmaPipeline.scheduler",description:`<strong>scheduler</strong> (<a href="/docs/diffusers/pr_14178/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler">BlockRefinementScheduler</a>, <a href="/docs/diffusers/pr_14178/en/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler">DiscreteDDIMScheduler</a> or <a href="/docs/diffusers/pr_14178/en/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler">EntropyBoundScheduler</a>) &#x2014;
The sampler that commits and renoises canvas tokens each denoising step.`,name:"scheduler"},{anchor:"diffusers.DiffusionGemmaPipeline.processor",description:`<strong>processor</strong> (<a href="https://huggingface.co/docs/transformers/main/en/main_classes/processors#transformers.ProcessorMixin" rel="nofollow">ProcessorMixin</a>) &#x2014;
The processor used to apply the chat template and decode the generated tokens.`,name:"processor"}]});var N=e(B,8),R=r(N);i(R,{name:"__call__",anchor:"diffusers.DiffusionGemmaPipeline.__call__",source:"https://github.com/huggingface/diffusers/blob/vr_14178/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L163",parameters:[{name:"prompt",val:": str | list[str] | None = None"},{name:"messages",val:": list[dict] | None = None"},{name:"image",val:": Any | list[Any] | None = None"},{name:"add_generation_prompt",val:": bool = True"},{name:"gen_length",val:": int = 256"},{name:"num_inference_steps",val:": int = 48"},{name:"temperature",val:": float = 0.0"},{name:"cache_implementation",val:": str | None = None"},{name:"eos_early_stop",val:": bool = True"},{name:"eos_token_id",val:": int | None = None"},{name:"stability_threshold",val:": int = 1"},{name:"confidence_threshold",val:": float | None = 0.005"},{name:"generator",val:": torch.Generator | None = None"},{name:"output_type",val:": str = 'text'"},{name:"return_dict",val:": bool = True"},{name:"callback_on_step_end",val:": Callable[[Any, int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None"},{name:"callback_on_step_end_tensor_inputs",val:": list[str] | None = None"}],parametersDescription:[{anchor:"diffusers.DiffusionGemmaPipeline.__call__.prompt",description:`<strong>prompt</strong> (<code>str</code> or <code>List[str]</code>, <em>optional</em>) &#x2014;
Prompt text, wrapped in a chat template and tokenized by the processor. Provide either this or
<code>messages</code>.`,name:"prompt"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.messages",description:`<strong>messages</strong> (<code>List[Dict]</code>, <em>optional</em>) &#x2014;
A raw chat conversation to encode, e.g. <code>[{&quot;role&quot;: &quot;user&quot;, &quot;content&quot;: &quot;Hello&quot;}]</code> or a multi-turn /
multimodal conversation. Use this instead of <code>prompt</code> for anything beyond a single user turn.`,name:"messages"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.image",description:`<strong>image</strong> (<code>PIL.Image.Image</code> or <code>List</code>, <em>optional</em>) &#x2014;
Image(s) to pair with <code>prompt</code> for multimodal generation; the processor turns them into the model&#x2019;s
image inputs. For richer layouts, put the image content directly in <code>messages</code>.`,name:"image"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.add_generation_prompt",description:`<strong>add_generation_prompt</strong> (<code>bool</code>, defaults to <code>True</code>) &#x2014;
Whether to add the generation prompt when applying the chat template.`,name:"add_generation_prompt"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.gen_length",description:`<strong>gen_length</strong> (<code>int</code>, defaults to <code>256</code>) &#x2014;
Number of tokens to generate, rounded up to a multiple of the model&#x2019;s <code>canvas_length</code>.`,name:"gen_length"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.num_inference_steps",description:`<strong>num_inference_steps</strong> (<code>int</code>, defaults to <code>48</code>) &#x2014;
Number of denoising steps per canvas.`,name:"num_inference_steps"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.temperature",description:`<strong>temperature</strong> (<code>float</code>, defaults to <code>0.0</code>) &#x2014;
Sampling temperature for <code>DiscreteDDIMScheduler</code>/<code>BlockRefinementScheduler</code> (<code>0.0</code> is greedy);
<code>EntropyBoundScheduler</code> ignores it and anneals its own temperature. Other sampling knobs (e.g. <code>top_k</code>,
<code>threshold</code>, <code>t_min</code>/<code>t_max</code>) are scheduler config; set them on the scheduler, e.g. <code>pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, top_k=...)</code>.`,name:"temperature"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.cache_implementation",description:`<strong>cache_implementation</strong> (<code>str</code>, <em>optional</em>) &#x2014;
Set to <code>&quot;static&quot;</code> to prefill the encoder once per block into a persistent <code>StaticCache</code> and run the
decoder against it with fixed shapes, instead of re-encoding the full sequence on every step. The fixed
shapes also let you compile the decoder, e.g. <code>pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, fullgraph=True)</code>.`,name:"cache_implementation"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.eos_early_stop",description:`<strong>eos_early_stop</strong> (<code>bool</code>, defaults to <code>True</code>) &#x2014;
Whether to stop generating further canvases once every sequence has emitted EOS.`,name:"eos_early_stop"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.eos_token_id",description:`<strong>eos_token_id</strong> (<code>int</code>, <em>optional</em>) &#x2014;
EOS token ID for early stopping. Falls back to the processor&#x2019;s tokenizer.`,name:"eos_token_id"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.stability_threshold",description:`<strong>stability_threshold</strong> (<code>int</code>, defaults to <code>1</code>) &#x2014;
Number of consecutive steps the argmax prediction must be unchanged for a block to count as stable.
Only used when <code>confidence_threshold</code> is set.`,name:"stability_threshold"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.confidence_threshold",description:`<strong>confidence_threshold</strong> (<code>float</code>, <em>optional</em>, defaults to <code>0.005</code>) &#x2014;
Leave a block&#x2019;s denoising loop early once every example is stable (see <code>stability_threshold</code>) and the
mean per-token entropy of the prediction is below this value. Speeds up generation at matched quality;
the default matches the released checkpoint. Set to <code>None</code> to always run all <code>num_inference_steps</code>.`,name:"confidence_threshold"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.generator",description:`<strong>generator</strong> (<code>torch.Generator</code>, <em>optional</em>) &#x2014;
RNG for sampling.`,name:"generator"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.output_type",description:`<strong>output_type</strong> (<code>str</code>, defaults to <code>&quot;text&quot;</code>) &#x2014;
<code>&quot;text&quot;</code> decodes sequences into strings (requires a processor); <code>&quot;seq&quot;</code> returns token IDs only.`,name:"output_type"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.return_dict",description:`<strong>return_dict</strong> (<code>bool</code>, defaults to <code>True</code>) &#x2014;
Whether to return a <a href="/docs/diffusers/pr_14178/en/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput">DiffusionGemmaPipelineOutput</a> instead of a tuple.`,name:"return_dict"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.callback_on_step_end",description:`<strong>callback_on_step_end</strong> (<code>Callable</code> or <code>PipelineCallback</code>, <em>optional</em>) &#x2014;
Callback run after each denoising step with signature <code>callback_on_step_end(self, step, timestep, callback_kwargs)</code>. Allowed tensor keys: <code>canvas</code>, <code>logits</code>.`,name:"callback_on_step_end"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.callback_on_step_end_tensor_inputs",description:`<strong>callback_on_step_end_tensor_inputs</strong> (<code>List[str]</code>, <em>optional</em>) &#x2014;
Tensor keys to pass to the callback.`,name:"callback_on_step_end_tensor_inputs"}],returnDescription:`<script context="module">export const metadata = 'undefined';<\/script>


<p>The generated token IDs (<code>sequences</code>) and, for <code>output_type="text"</code>, the decoded <code>texts</code>.</p>
`,returnType:`<script context="module">export const metadata = 'undefined';<\/script>


<p><a
  href="/docs/diffusers/pr_14178/en/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput"
>DiffusionGemmaPipelineOutput</a> or <code>tuple</code></p>
`});var q=e(R,4);K(q,{anchor:"diffusers.DiffusionGemmaPipeline.__call__.example",children:(t,l)=>{var W=se(),X=e(V(W),2);o(X,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwdHJhbnNmb3JtZXJzJTIwaW1wb3J0JTIwQXV0b1Byb2Nlc3NvciUyQyUyMERpZmZ1c2lvbkdlbW1hRm9yQmxvY2tEaWZmdXNpb24lMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwQmxvY2tSZWZpbmVtZW50U2NoZWR1bGVyJTJDJTIwRGlmZnVzaW9uR2VtbWFQaXBlbGluZSUwQSUwQW1vZGVsX2lkJTIwJTNEJTIwJTIyZ29vZ2xlJTJGZGlmZnVzaW9uZ2VtbWEtMjZCLUE0Qi1pdCUyMiUwQW1vZGVsJTIwJTNEJTIwRGlmZnVzaW9uR2VtbWFGb3JCbG9ja0RpZmZ1c2lvbi5mcm9tX3ByZXRyYWluZWQobW9kZWxfaWQlMkMlMjBkdHlwZSUzRHRvcmNoLmJmbG9hdDE2JTJDJTIwZGV2aWNlX21hcCUzRCUyMmF1dG8lMjIpJTBBcHJvY2Vzc29yJTIwJTNEJTIwQXV0b1Byb2Nlc3Nvci5mcm9tX3ByZXRyYWluZWQobW9kZWxfaWQpJTBBc2NoZWR1bGVyJTIwJTNEJTIwQmxvY2tSZWZpbmVtZW50U2NoZWR1bGVyKCklMEElMEFwaXBlJTIwJTNEJTIwRGlmZnVzaW9uR2VtbWFQaXBlbGluZShtb2RlbCUzRG1vZGVsJTJDJTIwc2NoZWR1bGVyJTNEc2NoZWR1bGVyJTJDJTIwcHJvY2Vzc29yJTNEcHJvY2Vzc29yKSUwQW91dHB1dCUyMCUzRCUyMHBpcGUocHJvbXB0JTNEJTIyV2h5JTIwaXMlMjB0aGUlMjBza3klMjBibHVlJTNGJTIyJTJDJTIwZ2VuX2xlbmd0aCUzRDI1NiklMEFwcmludChvdXRwdXQudGV4dHMlNUIwJTVEKQ==",highlighted:`<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">import</span> torch
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoProcessor, DiffusionGemmaForBlockDiffusion
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> BlockRefinementScheduler, DiffusionGemmaPipeline

<span class="hljs-meta">&gt;&gt;&gt; </span>model_id = <span class="hljs-string">&quot;google/diffusiongemma-26B-A4B-it&quot;</span>
<span class="hljs-meta">&gt;&gt;&gt; </span>model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map=<span class="hljs-string">&quot;auto&quot;</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span>processor = AutoProcessor.from_pretrained(model_id)
<span class="hljs-meta">&gt;&gt;&gt; </span>scheduler = BlockRefinementScheduler()

<span class="hljs-meta">&gt;&gt;&gt; </span>pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
<span class="hljs-meta">&gt;&gt;&gt; </span>output = pipe(prompt=<span class="hljs-string">&quot;Why is the sky blue?&quot;</span>, gen_length=<span class="hljs-number">256</span>)
<span class="hljs-meta">&gt;&gt;&gt; </span><span class="hljs-built_in">print</span>(output.texts[<span class="hljs-number">0</span>])`,lang:"python",wrap:!1}),c(t,W)},$$slots:{default:!0}}),d(N),d(n);var D=e(n,2);s(D,{title:"DiffusionGemmaPipelineOutput",local:"diffusers.DiffusionGemmaPipelineOutput",headingTag:"h2"});var a=e(D,2),Q=r(a);i(Q,{name:"class diffusers.DiffusionGemmaPipelineOutput",anchor:"diffusers.DiffusionGemmaPipelineOutput",source:"https://github.com/huggingface/diffusers/blob/vr_14178/src/diffusers/pipelines/diffusion_gemma/pipeline_output.py#L25",parameters:[{name:"sequences",val:": torch.LongTensor"},{name:"texts",val:": list[str] | None = None"}],parametersDescription:[{anchor:"diffusers.DiffusionGemmaPipelineOutput.sequences",description:`<strong>sequences</strong> (<code>torch.LongTensor</code> of shape <code>(batch_size, gen_length)</code>) &#x2014;
The generated token IDs (the prompt is stripped off).`,name:"sequences"},{anchor:"diffusers.DiffusionGemmaPipelineOutput.texts",description:`<strong>texts</strong> (<code>list[str]</code>, <em>optional</em>) &#x2014;
The decoded text, one string per sequence. Only set for <code>output_type=&quot;text&quot;</code>.`,name:"texts"}]}),E(2),d(a);var C=e(a,2);P(C,{source:"https://github.com/huggingface/diffusers/blob/main/docs/source/en/api/pipelines/diffusion_gemma.md"}),E(2),c(S,u),O()}export{ie as component};