diffusers/v0.43.0/ko/api/pipelines/diffusion_gemma.md

DiffusionGemma

DiffusionGemma는 블록 확산(block-diffusion) 인코더-디코더 언어 모델입니다. 인과(causal) 인코더가 클린 프롬프트(그리고 이전에 생성된 블록)를 KV 캐시로 읽어들이고, 양방향(bidirectional) 디코더가 해당 캐시에 크로스 어텐션을 수행하여 canvas_length 토큰 크기의 고정 크기 "캔버스"를 디노이즈(denoise)합니다. 생성 과정은 캔버스에 대한 외부 자기회귀(autoregressive) 루프와 내부 디노이즈 루프를 번갈아 수행하며, 각 단계에서 후보 토큰을 샘플링하고 균일한 손상(uniform corruption) 모드의 BlockRefinementScheduler를 통해 가장 신뢰도가 높은 토큰을 확정(commit)한 후 나머지 토큰에 다시 노이즈를 적용합니다. 모델 자체는 transformers에 DiffusionGemmaForBlockDiffusion으로 포함되어 있으며, 공개된 체크포인트는 google/diffusiongemma-26B-A4B-it입니다.

사용

import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = "google/diffusiongemma-26B-A4B-it"
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead", fullgraph=True)
output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
)
print(output.texts[0])

num_inference_steps는 캔버스당 디노이징 단계 수입니다(48은 공개된 체크포인트와 일치합니다). 단계 수가 적을수록 빠르지만 품질이 낮아집니다. cache_implementation="static"을 사용하면 디코더를 cudagraphs와 함께 torch.compile할 수 있습니다(Static cache and compilation 참조). 더 간단한 동적 캐시 실행을 원하면 이 두 가지를 모두 제거하세요.

멀티턴 또는 멀티모달 입력의 경우, prompt 대신 원시 messages 대화를 전달합니다. 이는 일반적인 채팅 형식의 {"role", "content"} 딕셔너리 목록이며, 프로세서가 이를 채팅 템플릿을 통해 처리합니다:

messages = [
    {"role": "user", "content": "Why is the sky blue?"},
]
# or with an image:
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": "Describe this image."},
        ],
    },
]
output = pipe(messages=messages, gen_length=256)

단일 사용자 턴의 경우 messages를 생략하고 prompt와 함께 image를 전달할 수 있습니다. 프로세서가 이를 자동으로 모델의 이미지 입력으로 변환합니다.

스케줄러

스케줄러는 각 캔버스를 디노이즈하는 샘플러이며, 교체 가능합니다. 다른 부분을 건드리지 않고 스케줄러만 교체하여 샘플링 전략을 변경할 수 있습니다. 사용할 수 있는 스케줄러는 세 가지입니다:

• BlockRefinementScheduler (기본값): 각 단계에서 가장 확신도가 높은 토큰(threshold 이상, 그리고 단계별 균등 할당분 포함)을 확정하고 나머지는 다시 노이즈를 적용합니다. editing_threshold를 통해 이미 확정된 토큰을 다시 수정할 수도 있습니다.
• DiscreteDDIMScheduler: 균일한 손상 과정(D3PM)의 정확한 이산 사후 분포에서 각 위치를 샘플링합니다. 매개변수가 없으며, 마지막 단계에서 예측된 토큰을 결정적으로 확정합니다.
• EntropyBoundScheduler: entropy_bound 이하의 결합 엔트로피를 유지하는 가장 낮은 엔트로피 위치들을 커밋하여, 대략적으로 독립적인 토큰들이 함께 수락되도록 합니다. 첫 번째 단계에서 t_max(0.8)로 시작하여 마지막 단계에서 t_min(0.4)까지 샘플링 온도를 어닐링하여, 공개된 체크포인트의 샘플러와 일치하도록 합니다.

from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
# or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

Scheduler별 샘플링 파라미터(블록 정밀화 threshold/top_k, 엔트로피 상한 등)는 scheduler config에서 설정합니다:

from diffusers import BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)

EntropyBoundScheduler는 디노이징 단계 동안 내부적으로 샘플링 온도(t_max/t_min)를 감쇠합니다. DiscreteDDIMScheduler와 BlockRefinementScheduler는 파이프라인에 전달되는 고정된 temperature를 사용합니다(그리디의 경우 0.0).

Predictor-corrector 샘플링

DiscreteDDIMScheduler는 Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation의 leave-one-out predictor-corrector를 지원합니다. corrector_steps개의 Gibbs sweeps를 수행하여 노이즈된 marginal의 one-coordinate conditional에서 가장 확신이 낮은 위치들을 다시 샘플링함으로써 캔버스를 정제합니다. 이 과정은 해당 marginal을 불변으로 유지하면서 추가적인 학습 비용 없이 생성 품질을 향상시킵니다. 공개된 체크포인트에서 바로 동작합니다. uniform diffusion의 경우 denoiser와 leave-one-out posterior가 closed form으로 상호 교환 가능하므로, corrector는 재학습 없이 필요한 leave-one-out 양을 복원할 수 있습니다.

corrector sweeps는 num_inference_steps 예산에 추가되는 것이 아니라 그 안에 포함됩니다. 파이프라인은 predictor 단계 수를 줄이고, 확보된 forward를 corrector에 사용하므로 모델 forward의 총 수는 num_inference_steps로 유지되며, predictor-corrector의 비용은 일반적인 ancestral sampling과 동일합니다.

from diffusers import DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=12)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

PEPT 어댑터

Denoiser는 🤗 Transformers 모델이므로, 어댑터는 diffusers의 load_lora_weights API가 아니라 해당 모델의 네이티브 PEFT 통합을 통해 로드됩니다. 이 통합은 어댑터 유형에 무관하므로, 동일한 호출로 LoRA, DoRA 또는 기타 PEFT 어댑터(예: TRL의 SFTTrainer 출력)를 로드할 수 있습니다. 모델 구성 요소에서 직접 어댑터를 관리하세요:

pipe.model.load_adapter("path/to/adapter", adapter_name="sft")  # LoRA, DoRA, ...
pipe.model.set_adapter("sft")
output = pipe(prompt="Why is the sky blue?", gen_length=256)

pipe.model.disable_adapters()  # run the base model
pipe.model.delete_adapter("sft")

Adapters는 활성화된 상태로 유지되며 병합되지 않습니다: DiffusionGemma는 인코더와 디코더의 기본 가중치를 연결(tie)하므로, 어댑터를 이 가중치에 병합하면 두 가지 분기(branch) 모두 손상됩니다.

정적 캐시 및 컴파일

pipeline은 블록마다 인코더를 한 번만 사전 입력하여 재사용 가능한 캐시(기본값은 DynamicCache)에 저장합니다. cache_implementation="static"을 전달하면 고정된 모양의 StaticCache를 대신 사용하며, 이 모양 덕분에 cudagraphs를 사용하여 디코더를 torch.compile하여 추가적인 속도 향상을 얻을 수 있습니다(pipeline은 각 스텝을 마킹하고 logits를 복제하여 cudagraph 메모리가 덮어쓰여지지 않도록 합니다). 이는 사용법에 표시된 설정입니다. 더 간단한 동적 캐시 실행을 원하면 torch.compile 호출과 cache_implementation="static"을 모두 제거하세요.

적응형 중단

블록은 일반적으로 num_inference_steps가 모두 소진되기 전에 수렴합니다. 따라서 기본적으로 파이프라인은 각 배치 예시의 argmax 예측이 stability_threshold 단계 동안 안정적이고 토큰당 평균 엔트로피가 confidence_threshold(발행된 체크포인트에서 사용된 값인 0.005) 미만으로 떨어지면 해당 예시를 동결합니다. 모든 예시가 동결되면 디노이징 루프가 종료됩니다. 이는 동일한 품질에서 디코더 포워드를 약 절반으로 줄여 주며, 단일 스루풋 레버 중 가장 큰 효과를 냅니다. confidence_threshold=None을 전달하면 항상 전체 num_inference_steps를 실행합니다:

output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None)  # disable adaptive stopping

콜백

Callbacks는 각 denoising 단계 이후에 실행됩니다. callback_on_step_end_tensor_inputs를 전달하여 callback_kwargs에 포함할 텐서를 선택할 수 있으며, canvas(현재 블록 토큰)와 logits가 사용 가능합니다. Callback에서 {"canvas": ...}를 반환하면 canvas가 교체됩니다.

def on_step_end(pipe, step, timestep, callback_kwargs):
    canvas = callback_kwargs["canvas"]
    # Inspect or modify `canvas` here.
    return {"canvas": canvas}

out = pipe(
    prompt="Why is the sky blue?",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["canvas"],
)

DiffusionGemmaPipeline[[diffusers.DiffusionGemmaPipeline]]

diffusers.DiffusionGemmaPipeline[[diffusers.DiffusionGemmaPipeline]]

diffusers.DiffusionGemmaPipeline(model: Any, scheduler: BlockRefinementScheduler | DiscreteDDIMScheduler | EntropyBoundScheduler, processor: Any)

소스

매개변수:

model (DiffusionGemmaForBlockDiffusion) : 블록 디퓨전 디노이저(결합된 가중치를 가진 인과적 인코더 + 양방향 디코더).

scheduler (BlockRefinementScheduler, DiscreteDDIMScheduler 또는 EntropyBoundScheduler) : 각 디노이징 단계에서 캔버스 토큰을 커밋하고 재노이즈하는 샘플러입니다.

processor (ProcessorMixin) : 채팅 템플릿을 적용하고 생성된 토큰을 디코딩하는 데 사용되는 프로세서입니다.

DiffusionGemma 블록 확산 텍스트 생성을 위한 파이프라인.

DiffusionGemma는 블록 확산 인코더-디코더 모델입니다. 인과적 인코더가 깨끗한 프롬프트(그리고 이전에 생성된 블록)를 KV 캐시로 읽은 후, 양방향 디코더가 해당 캐시에 크로스 어텐션을 수행하여 canvas_length 토큰 크기의 고정된 "캔버스"를 디노이즈합니다. 생성 과정은 캔버스에 대한 외부 자기회귀 루프와 내부 디노이즈 루프를 번갈아 수행하며, 각 단계에서 후보 토큰을 샘플링하고 BlockRefinementScheduler(일괄 손상 모드, mask_token_id=None)를 통해 가장 확신이 높은 토큰을 확정하고 나머지는 다시 노이즈를 적용합니다.

모델은 DiffusionGemmaForBlockDiffusion 인스턴스로, forward(input_ids, decoder_input_ids=..., self_conditioning_logits=..., ...)를 노출하고 캔버스에 대해 [batch, canvas_length, vocab_size] 모양의 logits를 반환해야 합니다. 모델 카드는 https://huggingface.co/google/diffusiongemma-26B-A4B-it에서 확인하세요.

호출[[diffusers.DiffusionGemmaPipeline.call]]

__call__(prompt: str | list[str] | None = None, messages: list[dict] | None = None, image: Any | list[Any] | None = None, add_generation_prompt: bool = True, gen_length: int = 256, num_inference_steps: int = 48, temperature: float = 0.0, cache_implementation: str | None = None, eos_early_stop: bool = True, eos_token_id: int | None = None, stability_threshold: int = 1, confidence_threshold: float | None = 0.005, generator: torch.Generator | None = None, output_type: str = 'text', return_dict: bool = True, callback_on_step_end: Callable[[Any, int, int, dict], dict] | PipelineCallback | MultiPipelineCallbacks | None = None, callback_on_step_end_tensor_inputs: list[str] | None = None)

소스

매개변수:

prompt (str 또는 List[str], 선택 사항) : 채팅 템플릿으로 감싸고 프로세서로 토큰화된 프롬프트 텍스트입니다. 이 항목 또는 messages 중 하나를 제공해야 합니다.

messages (List[Dict], 선택) : 인코딩할 원시 채팅 대화입니다. 예: [{"role": "user", "content": "Hello"}] 또는 다중 턴/멀티모달 대화. 단일 사용자 턴을 넘어가는 모든 경우에 prompt 대신 이 항목을 사용하세요.

image (PIL.Image.Image 또는 List, 선택 사항) : 멀티모달 생성을 위해 prompt와 짝지어질 이미지입니다. 프로세서가 이를 모델의 이미지 입력으로 변환합니다. 더 풍부한 레이아웃을 원하면 이미지 내용을 messages에 직접 넣으세요.

add_generation_prompt (bool, 기본값 True) : 채팅 템플릿을 적용할 때 생성 프롬프트를 추가할지 여부.

gen_length (int, 기본값 256) : 생성할 토큰 수입니다. 모델의 canvas_length 배수로 올림됩니다.

num_inference_steps (int, 기본값 48) : 캔버스당 디노이징 단계 수.

temperature (float, 기본값 0.0) : DiscreteDDIMScheduler/BlockRefinementScheduler의 샘플링 온도 (0.0는 greedy); EntropyBoundScheduler는 이를 무시하고 자체 온도를 anneal합니다. 기타 샘플링 파라미터(예: top_k, threshold, t_min/t_max)는 스케줄러 설정이며, 스케줄러에 설정해야 합니다. 예: pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, top_k=...).

cache_implementation (str, 선택 사항) : "static"으로 설정하면, 매 스텝마다 전체 시퀀스를 다시 인코딩하는 대신 블록당 한 번씩 인코더를 StaticCache에 미리 채우고 고정된 shape로 데코더를 실행합니다. 고정된 shape 덕분에 데코더를 컴파일할 수도 있습니다. 예: pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, fullgraph=True).

eos_early_stop (bool, 기본값 True) : 모든 시퀀스가 EOS를 출력한 후 추가 캔버스 생성을 중단할지 여부.

eos_token_id (int, 선택 사항) : 조기 정지를 위한 EOS 토큰 ID입니다. 프로세서의 토크나이저로 대체됩니다.

stability_threshold (int, 기본값 1) : 해당 예제가 안정적으로 간주되려면 argmax 예측이 연속으로 몇 단계 동안 변하지 않아야 하는지를 나타냅니다. confidence_threshold가 설정된 경우에만 사용됩니다.

confidence_threshold (float, 선택 사항, 기본값 0.005) : 각 예제가 안정화되면(stability_threshold 참조) 해당 스케줄러 모양의 예측 로짓의 토큰당 평균 엔트로피가 이 값보다 낮을 때 예제를 고정합니다. 모든 예제가 고정되면 블록의 디노이징 루프가 종료됩니다. 동일한 품질에서 생성 속도를 높이며, 기본값은 공개된 체크포인트와 일치합니다. num_inference_steps를 항상 모두 실행하려면 None으로 설정하세요.

generator (torch.Generator, 선택 사항) : 샘플링을 위한 RNG.

output_type (str, 기본값 "text") : "text"는 시퀀스를 문자열로 디코딩합니다(프로세서 필요). "seq"는 토큰 ID만 반환합니다.

return_dict (bool, 기본값 True) : 튜플 대신 DiffusionGemmaPipelineOutput를 반환할지 여부.

callback_on_step_end (Callable 또는 PipelineCallback, 선택 사항) : 각 노이즈 제거 단계 이후에 실행되는 콜백으로, 시그니처는 callback_on_step_end(self, step, timestep, callback_kwargs)입니다. 허용되는 텐서 키: canvas, logits.

callback_on_step_end_tensor_inputs (List[str], 선택 사항) : 콜백에 전달할 Tensor 키.

반환값: DiffusionGemmaPipelineOutput 또는 tuple

생성된 토큰 ID(sequences) 및 output_type="text"의 경우 디코딩된 texts.

block diffusion로 텍스트 생성

예:

>>> import torch
>>> from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
>>> from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

>>> model_id = "google/diffusiongemma-26B-A4B-it"
>>> model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
>>> output = pipe(prompt="Why is the sky blue?", gen_length=256)
>>> print(output.texts[0])

DiffusionGemmaPipelineOutput[[diffusers.DiffusionGemmaPipelineOutput]]

diffusers.DiffusionGemmaPipelineOutput[[diffusers.DiffusionGemmaPipelineOutput]]

diffusers.DiffusionGemmaPipelineOutput(sequences: torch.LongTensor, texts: list[str] | None = None)

소스

매개변수:

sequences (torch.LongTensor, shape (batch_size, gen_length)) : 생성된 토큰 ID(프롬프트는 제거됨).

texts (list[str], 선택 사항) : 디코딩된 텍스트로, 시퀀스당 하나의 문자열입니다. output_type="text"인 경우에만 설정됩니다.

DiffusionGemma 블록 확산 생성을 위한 출력 클래스입니다.