diffusers/main/id/api/schedulers/block_refinement.md

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[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, threshold: float = 0.95, editing_threshold: float | None = None, minimal_topk: int = 1)

Sumber

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

Pada setiap langkah, penjadwal mengambil sampel token kandidat dari logit model dan mengkonfirmasi token dengan keyakinan 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 diselesaikan, 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[[diffusers.BlockRefinementScheduler.add_noise]]

add_noise(original_samples: torch.LongTensor, attention_mask: torch.LongTensor, prompt_length: int, block_length: int, mask_token_id: int, generator: torch.Generator | None = None)

Sumber

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) : Token ID 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 beserta 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[[diffusers.BlockRefinementScheduler.check_block_should_continue]]

check_block_should_continue(step_idx: int, masks_remaining: bool, editing_enabled: bool, editing_transfer_index: torch.BoolTensor, post_steps: int, max_post_steps: int, finished: torch.BoolTensor)

Sumber

Parameter:

step_idx (int) : Indeks langkah penyempurnaan saat ini di dalam blok ini.

masks_remaining (bool) : Apakah masih ada token mask yang tersisa di dalam 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[[diffusers.BlockRefinementScheduler.check_eos_finished]]

check_eos_finished(cur_x: torch.LongTensor, sampled_tokens: torch.LongTensor, final_transfer: torch.BoolTensor, finished: torch.BoolTensor, eos_token_id: int, mask_token_id: int, prompt_length: int)

Sumber

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 penjadwal 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[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

get_num_transfer_tokens(block_length: int, num_inference_steps: int)

Sumber

Distribusikan block_length token commits secara merata di antara num_inference_steps langkah.

langkah[[diffusers.BlockRefinementScheduler.step]]

step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.LongTensor, mask_token_id: int | None = None, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'auto', threshold: float | None = None, editing_threshold: float | None = None, minimal_topk: int | None = None, prompt_mask: torch.BoolTensor | None = None, generator: torch.Generator | None = None, return_dict: bool = True)

Sumber

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) : Token ID yang digunakan untuk posisi yang dimasker. Ketika None, scheduler berjalan dalam mode korupsi seragam: 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) : Batasan pemotongan sampling nukleus.

top_k (int, opsional) : Batasan pemotongan sampling top-k.

sampling_method (str) : Metode sampling (auto, greedy, multinomial).

threshold (float, opsional) : Ambang kepercayaan untuk mengonfirmasi token. Bawaan ke nilai konfigurasi.

editing_threshold (float, opsional) : Ambang kepercayaan untuk mengedit token non-mask; harus bernilai positif untuk mengaktifkan pengeditan. Bawaan adalah nilai konfigurasi.

minimal_topk (int, opsional) : Jumlah token minimum yang akan 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 sunting token yang sudah ada.

BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput(prev_sample: torch.LongTensor, transfer_index: torch.BoolTensor, editing_transfer_index: torch.BoolTensor, sampled_tokens: torch.LongTensor, sampled_probs: torch.Tensor, pred_logits: torch.Tensor)

Sumber

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.