Discrete diffusion in diffusers by kashif · Pull Request #12911 · huggingface/diffusers

kashif · 2026-01-04T23:33:53Z

What does this PR do?

Add experimental support for discrete token diffusion methods and pipeline

moved llada2 to its own PR: #13226

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

yiyixuxu · 2026-01-15T21:29:23Z

Thanks for this PR!
cc @dg845, can you take a look here? it's related to Dream 7B #12091 you are working on

dg845 · 2026-01-23T01:14:47Z

Thanks for the PR! Some preliminary design questions and comments:

I think it could be useful to have a natural place to implement logic which is common to discrete diffusion models. Would something like a DiscreteDiffusionPipelineMixin make sense? For example, I think _resolve_start_token_id, _normalize_prefix_ids, _top_p_filtering, etc. could be candidates as mixin methods. (A possible alternative could be to put the methods in DiffusionPipeline, but it feels a little weird to put the methods there because they aren't applicable to continuous diffusion models.) But maybe this is premature, since we might not know what logic will end up being useful for all (or most) discrete diffusion models.
1. One motivation for this is that we often want to do semi-autoregressive (SAR) sampling for discrete diffusion models, so it would be useful to have autoregressive sampling techniques such as top-$p$ sampling, top-$k$ sampling, etc. So I think it would be nice to have a place where these methods can be implemented and tested once, and then new discrete diffusion models that support SAR sampling can have easy access to them without having to copy them every time.

Similarly, would it make sense to have a TokenizerTextProcessor class which handles text pre-processing and and post-processing, analogous to how VaeImageProcessor handles image pre- and post-processing? It's probably less necessary as we don't need to do as much normalization as for images, but I could see this being useful for handling e.g. chat templates like in the SDAR and LLaDA 2 pipelines.

As an aside, this could also be useful for existing (continuous) diffusion models, some of which have pretty involved text processing, such as pipelines like SanaPipeline that use a _text_preprocessing method:

diffusers/src/diffusers/pipelines/sana/pipeline_sana.py

Lines 548 to 549 in d4f97d1

    
           # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing 
        
           def _text_preprocessing(self, text, clean_caption=False):

Currently it looks like the pipelines only support denoising models with a transformers-like interface. But we would probably want to implement some discrete diffusion transformers in diffusers, which currently doesn't enforce that interface. So I think we should think about how we can handle both cases gracefully in discrete diffusion pipelines. (One solution could be to simply adopt the transformers interface for all discrete denoising models in diffusers, but that could be unnecessarily restrictive.)

dg845 · 2026-01-23T01:20:25Z

+        self.register_to_config(
+            seq_len=seq_len,
+            num_inference_steps=num_inference_steps,
+            inject_start_token=inject_start_token,
+        )


Generally we don't register default __call__ arguments to the config, but rather set them as default arguments to the __call__ method:

diffusers/src/diffusers/pipelines/ltx2/pipeline_ltx2.py

Lines 744 to 752 in d4f97d1

def __call__(

self,

prompt: Union[str, List[str]] = None,

negative_prompt: Optional[Union[str, List[str]]] = None,

height: int = 512,

width: int = 768,

num_frames: int = 121,

frame_rate: float = 24.0,

num_inference_steps: int = 40,

dg845 · 2026-01-23T01:26:33Z

+        *,
+        batch_size: int = 1,


diffusers pipelines usually don't set __call__ arguments to be keyword-only. (That's not to say that there are no arguments for it, but because other pipelines allow positional arguments I think the expectation is that discrete diffusion pipelines will allow them as well.)

dg845 · 2026-01-23T01:28:37Z

+        if seq_len is None:
+            seq_len = int(self.config.seq_len)
+        if num_inference_steps is None:
+            num_inference_steps = int(self.config.num_inference_steps)
+        if inject_start_token is None:
+            inject_start_token = bool(self.config.inject_start_token)


Following up on #12911 (comment), this logic could be removed if we don't register default arguments to the config.

dg845 · 2026-01-23T01:32:21Z

+        if infill_mask is not None:
+            if infill_mask.shape != (batch_size, seq_len):
+                raise ValueError(
+                    f"`infill_mask` must have shape {(batch_size, seq_len)}, got {tuple(infill_mask.shape)}."
+                )


I think input checking and exceptions should be moved to a check_inputs method, which is the usual practice for diffusers pipelines:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2.py

Lines 686 to 693 in d4f97d1

def check_inputs(

self,

prompt,

height,

width,

prompt_embeds=None,

callback_on_step_end_tensor_inputs=None,

):

dg845 · 2026-01-23T01:36:13Z

+                return int(token_id)
+        return None
+
+    def _init_latents(


We usually name methods which sample latents from the prior distribution prepare_latents:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2.py

Line 619 in d4f97d1

def prepare_latents(

dg845 · 2026-01-23T01:53:41Z

+        if hasattr(self.scheduler, "forward_process") and getattr(self.scheduler, "forward_process") == "uniform":
+            # Uniform prior over token IDs. Mirror scheduler's exclude-mask behavior.
+            if getattr(self.scheduler, "exclude_mask_from_uniform", False) and hasattr(
+                self.scheduler, "_sample_uniform_tokens"
+            ):
+                return self.scheduler._sample_uniform_tokens(
+                    torch.Size((batch_size, seq_len)),
+                    device=device,
+                    dtype=torch.long,
+                    generator=generator,
+                )
+            vocab_size = int(getattr(self.scheduler, "vocab_size", 0))
+            if vocab_size <= 0:
+                raise ValueError("Scheduler must define `vocab_size` for uniform prior sampling.")
+            return torch.randint(
+                0, vocab_size, (batch_size, seq_len), device=device, dtype=torch.long, generator=generator
+            )


Suggestion: maybe it would be cleaner to define a scheduler method called (say) sample_prior which samples from the prior distribution based on the configured forward_process? So if self.forward_process == "uniform", we would call _sample_uniform_tokens under the hood in sample_prior to sample from a uniform prior distribution.

I think this would allow for more graceful support of other possible forward processes, and make the pipeline code cleaner (as most of the logic would be handled inside the scheduler).

dg845 · 2026-04-02T06:34:43Z

+        self.vocab_size = int(vocab_size)
+        self.mask_token_id = int(mask_token_id)
+        self.num_train_timesteps = int(num_train_timesteps)
+        self.t_eps = float(t_eps)


Suggested change

self.vocab_size = int(vocab_size)

self.mask_token_id = int(mask_token_id)

self.num_train_timesteps = int(num_train_timesteps)

self.t_eps = float(t_eps)

I think this is unnecessary as register_to_config should make these available as self.config.vocab_size, self.config.mask_token_id, etc.

dg845 · 2026-04-02T06:39:58Z

+        p_uniform = max(math.exp(-float(clip_noise)), float(p_uniform))
+        log_B = float(gamma) * math.log(2.0) + math.log(p_uniform) - math.log(1.0 - p_uniform)
+        log_B = float(np.clip(log_B, -float(clip_noise), float(clip_noise)))
+        self.log_B = float(log_B)
+        self.log_gamma = float(math.log(float(gamma)))


Suggested change

p_uniform = max(math.exp(-float(clip_noise)), float(p_uniform))

log_B = float(gamma) * math.log(2.0) + math.log(p_uniform) - math.log(1.0 - p_uniform)

log_B = float(np.clip(log_B, -float(clip_noise), float(clip_noise)))

self.log_B = float(log_B)

self.log_gamma = float(math.log(float(gamma)))

p_uniform = max(math.exp(clip_noise), p_uniform)

log_B = gamma * math.log(2.0) + math.log(p_uniform) - math.log(1.0 - p_uniform)

log_B = np.clip(log_B, -clip_noise, clip_noise)

self.log_B = float(log_B)

self.log_gamma = math.log(gamma)

Can we remove all the int(...) and float(...) casts here and elsewhere?

dg845 · 2026-04-02T06:43:37Z

+
+class HybridTokenDiffusionScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Hybrid-transition discrete token diffusion scheduler.


Can the __init__ arguments (such as p_uniform, clip_noise, gamma, etc.) be documented in the docstring here, including what they mean and what values might be reasonable for them?

dg845 · 2026-04-02T06:55:12Z

+        p_uniform: float = 0.0,
+        clip_noise: float = 20.0,


The default parameters here set the effective p_uniform to $\exp(-20)$. If I understand correctly, p_uniform corresponds to the maximum probability that tokens will transition to another token uniformly at random instead of to the mask token. Since $\exp(-20)$ is very small, I believe the default schedule is thus essentially identical to a normal absorbing/masked diffusion schedule. Is there a reasonable default setting such that this scheduler behaves noticeably different from standard absorbing diffusion (e.g. as implemented by TokenDiffusionScheduler)?

dg845 · 2026-04-02T07:00:25Z

+        elif noise_type == "cosine":
+            return 1.0 - (1.0 - eps) * torch.cos(t * math.pi / 2.0)


The alpha schedules here are different from the alpha schedules defined in TokenDiffusionScheduler, is this intended? See also #12911 (comment).

dg845 · 2026-04-02T07:04:34Z

+
+    def _compute_move_chance(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the probability that a token has been masked (move chance) at continuous time *t*.


If I understand correctly, the move chance is given by $1 - \alpha(t)$. I think documenting the relationship between the move_chance and alpha_t here would be useful since it would be easier to compare this scheduler with similar schedulers such as TokenDiffusionScheduler.

dg845 · 2026-04-02T07:05:42Z

+            `torch.Tensor`: Move chance at each timestep value, same shape as *t*.
+        """
+        noise_type = self.config.noise_type
+        eps = 1e-3


I think eps should be configurable via __init__, following TokenDiffusionScheduler.

dg845 · 2026-04-02T07:07:15Z

+        # Compute move chances at t and s = t - dt
+        # ------------------------------------------------------------------
+        move_chance_t = self._compute_move_chance(t).to(dtype=torch.float64)
+        move_chance_s = self._compute_move_chance(t - dt).to(dtype=torch.float64)


I think getting s from the timestep schedule would be better in case we want to support non-linspace timestep schedules.

dg845 · 2026-04-02T07:11:46Z

+        # Subs parameterization: mask token gets -inf, then log_softmax normalizes.
+        # For unmasked positions, the distribution is forced to be the identity.
+        logits[..., mask_token_id] = -1e9
+        logits = logits - torch.logsumexp(logits, dim=-1, keepdim=True)


My understanding is that this computes a log softmax, would using something like torch.special.log_softmax be better here? Also I think it might be more clear to rename this to something like log_probs.

dg845 · 2026-04-02T07:13:17Z

+        gumbel_noise = -(torch.rand_like(q_xs, generator=generator) + 1e-10).log()
+        gumbel_noise = (1e-10 + gumbel_noise).clamp(min=1e-30)
+        x_block = (q_xs / gumbel_noise).argmax(dim=-1)


Could we reuse the _gumbel_argmax function defined in scheduling_token_diffusion.py here?

dg845 · 2026-04-02T07:26:44Z

+    def step(
+        self,
+        draft_tokens: torch.LongTensor,
+        target_logits: torch.Tensor,


Would it be possible to refactor the DFlash step method such that it follows the standard step interface?

diffusers/src/diffusers/schedulers/scheduling_block_refinement.py

Lines 166 to 170 in e365d74

def step(

self,

model_output: torch.Tensor,

timestep: int | torch.Tensor,

sample: torch.LongTensor,

Pipelines: - TokenDiffusion: add usage example, remove unused attention_mask, add sample_prior to scheduler, inline enforce_fixed_masks - HybridTokenDiffusion: consolidate into thin wrapper over TokenDiffusion - SDAR: rename cur_x to block_x, remove int() casts, remove unused vars - DFlash: inline _get_target_layer_ids, remove eval() calls, remove from_pretrained override, simplify model support - BD3LM: refactor block logic parallel with LLaDA2 Schedulers: - TokenDiffusion: pre-compute alpha schedule, cleaner if/elif in step, add sample_prior method - HybridTokenDiffusion: remove redundant self.xxx assignments, document params, remove int/float casts - BD3LM: make eps configurable, document move_chance vs alpha_t, use log_softmax, get s from timestep schedule, reuse _gumbel_argmax - DFlash: refactor step to standard model_output interface

…behavior

…indow context Three features from the official BD3LM repo now supported: - first_hitting: faster equivalent sampler from Zheng et al. 2025 that uses t *= u^(1/num_masked) instead of uniform timestep spacing - variable_length: entropy-based and EOS-based stopping for variable-length generation, truncating the output accordingly - context_size: sliding window to cap the context passed to the model All three features added to the scheduler as methods (compute_first_hitting_timestep, check_variable_length_stop) and exposed as pipeline __call__ parameters.

kashif added 11 commits December 20, 2025 14:45

initial

bb70cab

use the scheduler

f4990a9

added Block-wise sampling

c74d66d

add hybrid

9a9470c

hybrid sample training

886a76d

initial trainer

6f94f4f

added llada2

c1bf227

add sample_llada2.py

f2b9223

added api

cf18acd

fix llada2 sampling

9209924

formatting

5835d2c

kashif marked this pull request as draft January 4, 2026 23:34

kashif changed the title ~~Discrete diffusion in diffuers~~ Discrete diffusion in diffusers Jan 4, 2026

kashif added 7 commits January 6, 2026 11:58

Merge branch 'main' into diff-d2

264886a

Merge branch 'main' into diff-d2

541c03b

make fix-copies

1b3ffd7

fix docs

c263f16

add dflash pipeline

b97e11c

added SDAR JET pipeline and scheduler

471bfd3

Merge branch 'main' into diff-d2

ec4b2a4

kashif added 2 commits January 18, 2026 18:34

Merge branch 'main' into diff-d2

2f8a48b

Merge branch 'main' into diff-d2

3772ca1

dg845 reviewed Jan 23, 2026

View reviewed changes

dg845 reviewed Apr 2, 2026

View reviewed changes

kashif added 6 commits April 2, 2026 18:40

Merge upstream/main into diff-d2

640ac86

Add paper references for alpha schedules, document p_uniform default …

9772051

…behavior

fix: sample_hybrid_token_diffusion.py return_text -> output_type

9d37561

Add BD3LM sampling script matching current pipeline API

b3831f0

Merge branch 'main' into diff-d2

7558c79

github-actions bot added documentation Improvements or additions to documentation tests utils pipelines examples schedulers size/L PR with diff > 200 LOC labels Apr 14, 2026

github-actions bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 14, 2026

fix formatting

76e12ce

github-actions bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 14, 2026

	def __call__(
	self,
	prompt: Union[str, List[str]] = None,
	negative_prompt: Optional[Union[str, List[str]]] = None,
	height: int = 512,
	width: int = 768,
	num_frames: int = 121,
	frame_rate: float = 24.0,
	num_inference_steps: int = 40,

	def check_inputs(
	self,
	prompt,
	height,
	width,
	prompt_embeds=None,
	callback_on_step_end_tensor_inputs=None,
	):

		elif noise_type == "cosine":
		return 1.0 - (1.0 - eps) * torch.cos(t * math.pi / 2.0)

	def step(
	self,
	model_output: torch.Tensor,
	timestep: int \| torch.Tensor,
	sample: torch.LongTensor,

Conversation

kashif commented Jan 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

yiyixuxu commented Jan 15, 2026

Uh oh!

dg845 commented Jan 23, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kashif commented Jan 4, 2026 •

edited

Loading