# DiffusionGemma

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

## 사용

```py
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](#static-cache-and-compilation) 참조). 더 간단한 동적 캐시 실행을 원하면 이 두 가지를 모두 제거하세요.

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

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

```py
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에서 설정합니다:

```py
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](https://huggingface.co/papers/2605.22765)의 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과 동일합니다.

```py
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](https://huggingface.co/docs/peft) 통합을 통해 로드됩니다. 이 통합은 어댑터 유형에 무관하므로, 동일한 호출로 LoRA, DoRA 또는 기타 PEFT 어댑터(예: TRL의 `SFTTrainer` 출력)를 로드할 수 있습니다. 모델 구성 요소에서 직접 어댑터를 관리하세요:

```py
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 메모리가 덮어쓰여지지 않도록 합니다). 이는 [사용법](#usage)에 표시된 설정입니다. 더 간단한 동적 캐시 실행을 원하면 `torch.compile` 호출과 `cache_implementation="static"`을 모두 제거하세요.

## 적응형 중단

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

```py
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가 교체됩니다.

```py
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]]

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

[소스](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L53)

**매개변수:**

model ([DiffusionGemmaForBlockDiffusion](https://huggingface.co/docs/transformers/main/en/model_doc/diffusion_gemma#transformers.DiffusionGemmaForBlockDiffusion)) : 블록 디퓨전 디노이저(결합된 가중치를 가진 인과적 인코더 + 양방향 디코더).

scheduler ([BlockRefinementScheduler](/docs/diffusers/main/ko/api/schedulers/block_refinement#diffusers.BlockRefinementScheduler), [DiscreteDDIMScheduler](/docs/diffusers/main/ko/api/schedulers/discrete_ddim#diffusers.DiscreteDDIMScheduler) 또는 [EntropyBoundScheduler](/docs/diffusers/main/ko/api/schedulers/entropy_bound#diffusers.EntropyBoundScheduler)) : 각 디노이징 단계에서 캔버스 토큰을 커밋하고 재노이즈하는 샘플러입니다.

processor ([ProcessorMixin](https://huggingface.co/docs/transformers/main/en/main_classes/processors#transformers.ProcessorMixin)) : 채팅 템플릿을 적용하고 생성된 토큰을 디코딩하는 데 사용되는 프로세서입니다.

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

DiffusionGemma는 블록 확산 인코더-디코더 모델입니다. 인과적 인코더가 깨끗한 프롬프트(그리고 이전에 생성된 블록)를 KV 캐시로 읽은 후, 양방향 디코더가 해당 캐시에 크로스 어텐션을 수행하여 `canvas_length` 토큰 크기의 고정된 "캔버스"를 디노이즈합니다. 생성 과정은 캔버스에 대한 외부 자기회귀 루프와 내부 디노이즈 루프를 번갈아 수행하며, 각 단계에서 후보 토큰을 샘플링하고 [BlockRefinementScheduler](/docs/diffusers/main/ko/api/schedulers/block_refinement#diffusers.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__]]

```python
__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)
```

[소스](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py#L163)

**매개변수:**

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](/docs/diffusers/main/ko/api/pipelines/diffusion_gemma#diffusers.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](/docs/diffusers/main/ko/api/pipelines/diffusion_gemma#diffusers.DiffusionGemmaPipelineOutput) 또는 `tuple`

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

block diffusion로 텍스트 생성

예:
```python
>>> 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]]

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

[소스](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/diffusion_gemma/pipeline_output.py#L25)

**매개변수:**

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

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

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