Diffusers documentation

BlockRefinementScheduler

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

BlockRefinementScheduler

BlockRefinementScheduler mengelola penyempurnaan iteratif per blok untuk difusi token diskrit. Pada setiap langkah, scheduler ini mengunci token yang paling yakin dan secara opsional mengedit token yang sudah dikunci ketika model memprediksi token lain dengan keyakinan tinggi.

Scheduler ini digunakan oleh LLaDA2Pipeline.

BlockRefinementScheduler

kelas diffusers.BlockRefinementScheduler

< >

( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )

Penjadwal untuk penyempurnaan iteratif per blok (commit-by-confidence).

Pada setiap langkah, penjadwal mengambil sampel token kandidat dari logit model dan mengkonfirmasi token dengan kepercayaan tertinggi. Jumlah token yang dikonfirmasi per langkah ditentukan dengan mendistribusikan panjang blok secara merata di antara jumlah langkah penyempurnaan.

Secara opsional mendukung pengeditan: setelah semua token mask teresolusi, token dapat diganti jika model memprediksi token yang berbeda dengan kepercayaan di atas editing_threshold positif (None, 0.0, atau nilai negatif menonaktifkan pengeditan).

add_noise

< >

( original_samples: torch.LongTensorattention_mask: torch.LongTensorprompt_length: intblock_length: intmask_token_id: intgenerator: torch.Generator | None = None ) → tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

Parameter

  • original_samples (torch.LongTensor dengan bentuk (batch_size, seq_len)) — ID token bersih.
  • attention_mask (torch.LongTensor dengan bentuk (batch_size, seq_len)) — Mask padding (1 untuk valid, 0 untuk padding).
  • prompt_length (int) — Jumlah token prompt awal yang dipertahankan tanpa masker.
  • block_length (int) — Ukuran blok untuk masking.
  • mask_token_id (int) — ID token yang digunakan untuk posisi yang dimasker.
  • generator (torch.Generator, opsional) — RNG untuk reproduktibilitas.

Mengembalikan

tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

(noisy, noisy_rev, masked, masked_rev) — dua urutan berisik yang saling melengkapi dan mask boolean yang sesuai.

Terapkan proses maju (noising) untuk masking blok semi-otoregresif.

Untuk setiap blok setelah prompt, sebagian acak dari token valid (non-padding) diganti dengan mask_token_id. Dua tampilan komplementer dikembalikan: noisy dan noisy_rev, di mana posisi yang dimasker pada salah satu adalah posisi yang tidak dimasker pada yang lain.

check_block_should_continue

< >

( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool

Parameter

  • step_idx (int) — Indeks langkah penyempurnaan saat ini di dalam blok ini.
  • masks_remaining (bool) — Apakah masih ada token mask yang tersisa di blok.
  • editing_enabled (bool) — Apakah mode pengeditan sedang aktif.
  • editing_transfer_index (torch.BoolTensor) — Token mana yang diedit pada langkah ini.
  • post_steps (int) — Jumlah langkah pengeditan pasca-masker yang telah dilakukan sejauh ini.
  • max_post_steps (int) — Jumlah langkah pengeditan pasca-masker maksimum yang diizinkan.
  • finished (torch.BoolTensor) — Flag selesai per batch (dari deteksi EOS).

Mengembalikan

bool

True jika penyempurnaan harus dilanjutkan, False untuk berhenti.

Tentukan apakah loop penyempurnaan internal harus dilanjutkan untuk blok saat ini.

check_eos_finished

< >

( cur_x: torch.LongTensorsampled_tokens: torch.LongTensorfinal_transfer: torch.BoolTensorfinished: torch.BoolTensoreos_token_id: intmask_token_id: intprompt_length: int ) → torch.BoolTensor

Parameter

  • cur_x (torch.LongTensor dengan bentuk (batch_size, seq_len)) — Urutan penuh saat ini yang mencakup semua blok hingga jendela saat ini.
  • sampled_tokens (torch.LongTensor dengan bentuk (batch_size, block_length)) — Token yang disampelkan oleh scheduler pada langkah ini.
  • final_transfer (torch.BoolTensor dengan bentuk (batch_size, block_length)) — Mask gabungan dari posisi yang dikomit dan diedit.
  • finished (torch.BoolTensor dengan bentuk (batch_size,)) — Flag selesai per-batch saat ini.
  • eos_token_id (int) — ID token EOS.
  • mask_token_id (int) — ID token mask.
  • prompt_length (int) — Jumlah token prompt di awal urutan.

Mengembalikan

torch.BoolTensor

Perbarui bendera selesai.

Perbarui flag selesai per-batch ketika token EOS dikomit.

get_num_transfer_tokens

< >

( block_length: intnum_inference_steps: int )

Distribusikan block_length token commits secara merata di antara num_inference_steps langkah.

langkah

< >

( model_output: torch.Tensortimestep: int | torch.Tensorsample: torch.LongTensormask_token_id: int | None = Nonetemperature: float = 0.0top_p: float | None = Nonetop_k: int | None = Nonesampling_method: str = 'auto'threshold: float | None = Noneediting_threshold: float | None = Noneminimal_topk: int | None = Noneprompt_mask: torch.BoolTensor | None = Nonegenerator: torch.Generator | None = Nonereturn_dict: bool = True )

Parameter

  • model_output (torch.Tensor dengan bentuk (batch_size, block_length, vocab_size)) — Logit mentah dari model untuk blok saat ini.
  • timestep (int atau torch.Tensor) — Indeks langkah saat ini dalam jadwal penyempurnaan blok.
  • sample (torch.LongTensor dengan bentuk (batch_size, block_length)) — ID token blok saat ini (mengandung token mask untuk posisi yang belum dikomit dalam mode berbasis mask).
  • mask_token_id (int, opsional) — ID token yang digunakan untuk posisi yang dimasker. Ketika None, penjadwal berjalan dalam mode korupsi seragam: ia melacak posisi yang sudah dikomit secara internal (di-reset pada timestep == 0) dan menambahkan noise kembali pada posisi yang belum dikomit dengan token acak seragam, sesuai dengan sampler penyempurnaan blok DiffusionGemma.
  • temperature (float) — Suhu sampling.
  • top_p (float, opsional) — Batas pemotongan nucleus sampling.
  • top_k (int, opsional) — Batas pemotongan sampling top-k.
  • sampling_method (str) — Metode sampling (auto, greedy, multinomial).
  • threshold (float, opsional) — Ambang batas kepercayaan untuk mengunci token. Bawaan ke nilai konfigurasi.
  • editing_threshold (float, opsional) — Ambang batas kepercayaan untuk mengedit token non-mask; harus bernilai positif untuk mengaktifkan pengeditan. Secara default menggunakan nilai config.
  • minimal_topk (int, opsional) — Token minimum yang harus dikomit per langkah. Bawaan ke nilai konfigurasi.
  • prompt_mask (torch.BoolTensor, opsional) — Mask boolean dengan bentuk (block_length,) di mana True menandai posisi prompt (tidak dapat diedit).
  • generator (torch.Generator, opsional) — RNG untuk sampling.
  • return_dict (bool) — Apakah akan mengembalikan BlockRefinementSchedulerOutput atau tuple.

Lakukan satu langkah penyempurnaan: ambil sampel dari logits, tetapkan token yang yakin, dan secara opsional edit yang sudah ada.

BlockRefinementSchedulerOutput

kelas diffusers.BlockRefinementSchedulerOutput

< >

( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )

Parameter

  • prev_sample (torch.LongTensor dengan bentuk (batch_size, block_length)) — Token blok yang diperbarui setelah langkah penyempurnaan saat ini.
  • transfer_index (torch.BoolTensor dengan bentuk (batch_size, block_length)) — Mask boolean yang menunjukkan token mana yang telah dikomit (mask-filling).
  • editing_transfer_index (torch.BoolTensor dengan bentuk (batch_size, block_length)) — Mask boolean yang menunjukkan token mana yang diedit (penggantian non-mask).
  • sampled_tokens (torch.LongTensor dengan bentuk (batch_size, block_length)) — ID token yang disampelkan dari logits model.
  • sampled_probs (torch.Tensor dengan bentuk (batch_size, block_length)) — Probabilitas dari token yang disampelkan.
  • pred_logits (torch.Tensor dengan bentuk (batch_size, block_length, vocab_size)) — Logits denoiser, diteruskan untuk self-conditioning langkah berikutnya.

Kelas output untuk penjadwalan penyempurnaan blok.

Update on GitHub