Buckets:
| 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_14409/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>{"role", "content"}</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_14409/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_14409/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_14409/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>{"canvas": ...}</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_14409/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">"google/diffusiongemma-26B-A4B-it"</span> | |
| model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map=<span class="hljs-string">"auto"</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">"reduce-overhead"</span>, fullgraph=<span class="hljs-literal">True</span>) | |
| output = pipe( | |
| prompt=<span class="hljs-string">"Why is the sky blue?"</span>, | |
| gen_length=<span class="hljs-number">256</span>, | |
| num_inference_steps=<span class="hljs-number">48</span>, | |
| cache_implementation=<span class="hljs-string">"static"</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">"role"</span>: <span class="hljs-string">"user"</span>, <span class="hljs-string">"content"</span>: <span class="hljs-string">"Why is the sky blue?"</span>}, | |
| ] | |
| <span class="hljs-comment"># or with an image:</span> | |
| messages = [ | |
| { | |
| <span class="hljs-string">"role"</span>: <span class="hljs-string">"user"</span>, | |
| <span class="hljs-string">"content"</span>: [ | |
| {<span class="hljs-string">"type"</span>: <span class="hljs-string">"image"</span>, <span class="hljs-string">"image"</span>: image}, | |
| {<span class="hljs-string">"type"</span>: <span class="hljs-string">"text"</span>, <span class="hljs-string">"text"</span>: <span class="hljs-string">"Describe this image."</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">"Why is the sky blue?"</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">"Why is the sky blue?"</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">"path/to/adapter"</span>, adapter_name=<span class="hljs-string">"sft"</span>) <span class="hljs-comment"># LoRA, DoRA, ...</span> | |
| pipe.model.set_adapter(<span class="hljs-string">"sft"</span>) | |
| output = pipe(prompt=<span class="hljs-string">"Why is the sky blue?"</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">"sft"</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">"Why is the sky blue?"</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">"canvas"</span>] | |
| <span class="hljs-comment"># Inspect or modify \`canvas\` here.</span> | |
| <span class="hljs-keyword">return</span> {<span class="hljs-string">"canvas"</span>: canvas} | |
| out = pipe( | |
| prompt=<span class="hljs-string">"Why is the sky blue?"</span>, | |
| callback_on_step_end=on_step_end, | |
| callback_on_step_end_tensor_inputs=[<span class="hljs-string">"canvas"</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_14409/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>) — | |
| 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_14409/en/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler">BlockRefinementScheduler</a>, <a href="/docs/diffusers/pr_14409/en/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler">DiscreteDDIMScheduler</a> or <a href="/docs/diffusers/pr_14409/en/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler">EntropyBoundScheduler</a>) — | |
| 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>) — | |
| 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_14409/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>) — | |
| 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>) — | |
| A raw chat conversation to encode, e.g. <code>[{"role": "user", "content": "Hello"}]</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>) — | |
| Image(s) to pair with <code>prompt</code> for multimodal generation; the processor turns them into the model’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>) — | |
| 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>) — | |
| Number of tokens to generate, rounded up to a multiple of the model’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>) — | |
| 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>) — | |
| 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>) — | |
| Set to <code>"static"</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>) — | |
| 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>) — | |
| EOS token ID for early stopping. Falls back to the processor’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>) — | |
| 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>) — | |
| Leave a block’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>) — | |
| RNG for sampling.`,name:"generator"},{anchor:"diffusers.DiffusionGemmaPipeline.__call__.output_type",description:`<strong>output_type</strong> (<code>str</code>, defaults to <code>"text"</code>) — | |
| <code>"text"</code> decodes sequences into strings (requires a processor); <code>"seq"</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>) — | |
| Whether to return a <a href="/docs/diffusers/pr_14409/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>) — | |
| 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>) — | |
| 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_14409/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">>>> </span><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-meta">>>> </span><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoProcessor, DiffusionGemmaForBlockDiffusion | |
| <span class="hljs-meta">>>> </span><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> BlockRefinementScheduler, DiffusionGemmaPipeline | |
| <span class="hljs-meta">>>> </span>model_id = <span class="hljs-string">"google/diffusiongemma-26B-A4B-it"</span> | |
| <span class="hljs-meta">>>> </span>model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map=<span class="hljs-string">"auto"</span>) | |
| <span class="hljs-meta">>>> </span>processor = AutoProcessor.from_pretrained(model_id) | |
| <span class="hljs-meta">>>> </span>scheduler = BlockRefinementScheduler() | |
| <span class="hljs-meta">>>> </span>pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor) | |
| <span class="hljs-meta">>>> </span>output = pipe(prompt=<span class="hljs-string">"Why is the sky blue?"</span>, gen_length=<span class="hljs-number">256</span>) | |
| <span class="hljs-meta">>>> </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_14409/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>) — | |
| 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>) — | |
| The decoded text, one string per sequence. Only set for <code>output_type="text"</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}; | |
Xet Storage Details
- Size:
- 34.7 kB
- Xet hash:
- 6e76627e611c003db99b81f8260d8951cd4effbd84589f52813dcf8ae912b964
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.