huggingface/transformers · error · ValueError

TextDiffusionStreamer only supports batch size 1

Error message

TextDiffusionStreamer only supports batch size 1

What it means

ValueError from TextDiffusionStreamer.put_draft: like TextStreamer, the diffusion streamer prints a single evolving draft to stdout, so draft token batches with batch dimension > 1 are rejected. Drafts overwrite each other via ANSI cursor save/restore, which is inherently single-stream.

Source

Thrown at src/transformers/generation/streamers.py:384

        # we recommend setting it to `False` by default.
        self._takes_logits = False
        self.sleep_time = sleep_time

    def _clear_draft(self):
        if self._has_draft:
            # Restore cursor and clear to end of screen
            print("\0338\033[J", end="", flush=True)
            self._has_draft = False

    def put_draft(self, value, **kwargs):
        """
        Receives the full sequence of draft tokens, decodes them, and prints them in yellow.
        Overwrites previous draft.
        """
        self._clear_draft()

        if len(value.shape) > 1 and value.shape[0] > 1:
            raise ValueError("TextDiffusionStreamer only supports batch size 1")
        elif len(value.shape) > 1:
            value = value[0]

        text = self.tokenizer.decode(value, **self.decode_kwargs)

        # Save cursor position
        print("\0337", end="", flush=True)
        # Print draft in yellow
        print(f"\033[33m{text}\033[0m", end="", flush=True)
        self._has_draft = True
        if self.sleep_time is not None:
            time.sleep(self.sleep_time)

    def put(self, value):
        """Receives confirmed tokens, clears draft, and prints them permanently."""
        self._clear_draft()
        super().put(value)

View on GitHub (pinned to a597f97485)

Solutions

1. Run diffusion generation with batch size 1 when using TextDiffusionStreamer.
2. Detach the streamer for batched runs.
3. If you only care about one sample, index the batch before calling put_draft.

Example fix

# before
streamer.put_draft(draft_tokens)  # (B>1, seq)

# after
streamer.put_draft(draft_tokens[0:1])

if value.dim() > 1 and value.shape[0] > 1:
    raise ValueError("TextDiffusionStreamer requires batch size 1; slice the batch first")
# or simply: value = value[:1]

When it happens

Trigger: Calling put_draft(value) with value.shape[0] > 1 in a text-diffusion generation loop (e.g. diffuLLaMA-style models) that was batched; passing the raw (batch, seq) tensor when only row 0 was intended.

Common situations: Reusing a batched diffusion pipeline with the visual streamer; forgetting to slice value[0] before put_draft in a custom loop.

Related errors

• TextStreamer only supports batch size 1
• `streamer` cannot be used with beam search (yet!). Make sure
• `crop` was called, but the current layer does not track past
• Once the sliding window size has been reached, `DynamicSlidi
• `crop` was called, but the current layer does not track past

AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/f9cdd26a92230473. Report an issue: GitHub.