DiffusionGemma

DiffusionGemma é um modelo de linguagem encoder-decoder de difusão em blocos. Um encoder causal lê o prompt limpo (e quaisquer blocos gerados anteriormente) para um cache KV, e um decoder bidirecional remove o ruído de uma “tela” de tamanho fixo com canvas_length tokens, fazendo cross-attention para esse cache. A geração alterna um loop autoregressivo externo sobre as telas com um loop interno de remoção de ruído, em que cada etapa amostra tokens candidatos, confirma os mais confiantes via BlockRefinementScheduler em modo de corrupção uniforme e re-aplica ruído aos demais. O próprio modelo está em transformers como DiffusionGemmaForBlockDiffusion; o checkpoint publicado é google/diffusiongemma-26B-A4B-it.

Uso

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 é o número de etapas de remoção de ruído por canvas (48 corresponde ao checkpoint publicado); menos etapas são mais rápidas, mas de menor qualidade. cache_implementation="static" permite que o decodificador seja compilado com torch.compile usando cudagraphs (veja Cache estático e compilação); remova ambos para uma execução mais simples com cache dinâmico.

Para entradas de múltiplas rodadas ou multimodais, passe uma conversa bruta messages em vez de prompt. Trata-se de uma lista de dicts {"role", "content"} no formato de chat usual, que o processador executa por meio do seu modelo de chat:

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)

Para um turno de usuário único, você pode omitir messages e passar uma image junto com o prompt; o processador a converte automaticamente nos inputs de imagem do modelo.

Agendadores

O scheduler é o sampler que remove o ruído de cada canvas e é intercambiável: troque-o para alterar a estratégia de amostragem sem tocar em mais nada. Três schedulers estão disponíveis:

• BlockRefinementScheduler (padrão): confirma os tokens mais confiantes a cada etapa (acima de threshold, além de uma cota uniforme por etapa) e re-aplica ruído aos demais. editing_threshold permite, adicionalmente, reeditar tokens já confirmados.
• DiscreteDDIMScheduler: amostra cada posição da posterior discreta exata do processo de corrupção uniforme (D3PM). Ele é livre de parâmetros e o passo final compromete deterministicamente os tokens previstos.
• EntropyBoundScheduler: confirma as posições de menor entropia cuja entropia conjunta permanece abaixo de entropy_bound, de modo que tokens aproximadamente independentes sejam aceitos juntos. Ele anela sua temperatura de amostragem de t_max (0.8) no primeiro passo até t_min (0.4) no último, correspondendo ao amostrador do checkpoint liberado.

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])

Os parâmetros de amostragem específicos do agendador (o threshold/top_k de refinamento de bloco, o limite de entropia, …) são definidos na configuração do agendador:

from diffusers import BlockRefinementScheduler

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

EntropyBoundScheduler reduz gradualmente a temperatura de amostragem (t_max/t_min) internamente ao longo das etapas de desnívelamento; DiscreteDDIMScheduler e BlockRefinementScheduler usam a temperature plana passada para o pipeline (0.0 para greedy).

Amostragem preditor-corretor

DiscreteDDIMScheduler suporta o preditor-corretor leave-one-out de Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. Ele refina o canvas com corrector_steps varreduras de Gibbs que reamostram as posições de menor confiança a partir da condicional de uma coordenada da marginal ruidosa, o que mantém essa marginal invariante e melhora a geração sem custo adicional de treinamento. Ele funciona diretamente no checkpoint liberado: para difusão uniforme, o denoiser e a posterior leave-one-out são intercambiáveis em forma fechada, de modo que o corretor recupera as quantidades leave-one-out de que precisa sem qualquer re-treinamento.

As varreduras do corretor são incorporadas ao orçamento de num_inference_steps em vez de serem adicionadas por cima: o pipeline executa menos etapas do preditor e gasta os forwards liberados nos corretores, de modo que o número total de forwards do modelo permanece num_inference_steps e o custo do preditor-corretor é o mesmo que o da amostragem ancestral simples.

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])

Adaptadores PEFT

O denoiser é um modelo do 🤗 Transformers, portanto os adaptadores são carregados por meio de sua integração nativa PEFT, em vez da API load_lora_weights do diffusers. Como essa integração é agnóstica em relação ao tipo de adaptador, as mesmas chamadas carregam LoRA, DoRA ou qualquer outro adaptador PEFT (por exemplo, a saída do SFTTrainer do TRL). Gerencie os adaptadores diretamente no componente do modelo:

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")

Os adaptadores permanecem ativos e não mesclados: o DiffusionGemma vincula os pesos base do codificador e do decodificador, portanto, fundir um adaptador neles corromperia ambos os ramos.

Cache estático e compilação

O pipeline pré-preenche o codificador uma vez por bloco em um cache reutilizável (um DynamicCache por padrão). Passar cache_implementation="static" usa um StaticCache de forma fixa em vez disso, cujas formas permitem torch.compile o decodificador com cudagraphs para um aumento adicional de velocidade (o pipeline marca cada etapa e clona os logits para que a memória dos cudagraphs não seja sobrescrita); este é o setup mostrado em Uso. Remova tanto a chamada torch.compile quanto cache_implementation="static" para uma execução mais simples com cache dinâmico.

Parada adaptativa

Um bloco geralmente converge antes que todos os num_inference_steps sejam utilizados, portanto, por padrão, o pipeline congela cada exemplo do lote uma vez que sua previsão de argmax está estável por stability_threshold passos e sua entropia média por token fica abaixo de confidence_threshold (0.005, o valor usado pelo checkpoint liberado). O loop de desruidamento termina quando todos os exemplos são congelados. Isso reduz aproximadamente pela metade o número de forwards do decodificador com qualidade equivalente e é o principal fator de throughput. Passe confidence_threshold=None para sempre executar todos os num_inference_steps:

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

Callbacks

Os callbacks são executados após cada etapa de remoção de ruído. Passe callback_on_step_end_tensor_inputs para selecionar quais tensores são incluídos em callback_kwargs; canvas (os tokens do bloco atual) e logits estão disponíveis. Retorne {"canvas": ...} do callback para substituir o 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

>>> 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