DiffusionGemma

DiffusionGemma es un modelo de lenguaje encoder-decoder de difusión por bloques. Un encoder causal lee el prompt limpio (y cualquier bloque generado previamente) en una caché KV, y un decoder bidireccional desruidiza un “lienzo” de tamaño fijo de canvas_length tokens mediante atención cruzada con esa caché. La generación alterna un bucle externo autoregresivo sobre lienzos con un bucle interno de desruidizado, donde cada paso muestrea tokens candidatos, confirma los de mayor confianza mediante BlockRefinementScheduler en modo de corrupción uniforme y vuelve a añadir ruido al resto. El modelo en sí se encuentra en transformers como DiffusionGemmaForBlockDiffusion; el checkpoint publicado es 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 es el número de pasos de desruido por lienzo (48 coincide con el punto de control publicado); menos pasos son más rápidos pero de menor calidad. cache_implementation="static" permite que el decodificador se compile con torch.compile y cudagraphs (ver Caché estática y compilación); omita ambos para una ejecución más simple con caché dinámica.

Para entradas de varios turnos o multimodales, pase una conversación cruda de messages en lugar de prompt. Es una lista de diccionarios {"role", "content"} en el formato de chat habitual, que el procesador ejecuta a través de su plantilla 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 un turno de un solo usuario, puedes omitir messages y pasar una image junto con el prompt; el procesador la convierte automáticamente en las entradas de imagen del modelo.

Planificadores

El scheduler es el sampler que desruidifica cada lienzo y es intercambiable: cámbialo para modificar la estrategia de muestreo sin tocar nada más. Hay tres schedulers disponibles:

• BlockRefinementScheduler (predeterminado): confirma en cada paso los tokens más seguros (los que superan threshold, más una cuota uniforme por paso) y vuelve a añadir ruido al resto. editing_threshold permite además reeditar tokens ya confirmados.
• DiscreteDDIMScheduler: muestrea cada posición de la posterior discreta exacta del proceso de corrupción uniforme (D3PM). No tiene parámetros y el paso final compromete de forma determinista los tokens predichos.
• EntropyBoundScheduler: confirma las posiciones de menor entropía cuya entropía conjunta se mantiene por debajo de entropy_bound, de modo que los tokens aproximadamente independientes se acepten juntos. Anela su temperatura de muestreo desde t_max (0.8) en el primer paso hasta t_min (0.4) en el último, coincidiendo con el muestreador del punto de control publicado.

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

Los parámetros de muestreo específicos del programador (el threshold/top_k de refinamiento de bloques, el límite de entropía, …) se configuran en la configuración del programador:

from diffusers import BlockRefinementScheduler

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

EntropyBoundScheduler reduce internamente su temperatura de muestreo (t_max/t_min) a lo largo de los pasos de desruido; DiscreteDDIMScheduler y BlockRefinementScheduler utilizan la temperature plana que se pasa a la pipeline (0.0 para greedy).

Muestreo predictor-corrección

DiscreteDDIMScheduler admite el predictor-corrección leave-one-out de Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. Refina el lienzo con corrector_steps barridos de Gibbs que muestrean de nuevo las posiciones de menor confianza a partir de la condicional de una coordenada de la marginal ruidosa, lo que deja esa marginal invariante y mejora la generación sin coste adicional de entrenamiento. Funciona directamente sobre el checkpoint publicado: para la difusión uniforme, el desruidador y la posterior leave-one-out son intercambiables en forma cerrada, por lo que el corrector recupera las cantidades leave-one-out que necesita sin necesidad de ningún reentrenamiento.

Los barridos del corrector se integran en el presupuesto de num_inference_steps en lugar de sumarse a él: la tubería ejecuta menos pasos del predictor y destina las pasadas liberadas a los correctores, de modo que el número total de pasadas del modelo se mantiene en num_inference_steps y el costo del predictor-corrector es el mismo que el de la muestreo ancestral simple.

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

El denoiser es un modelo de 🤗 Transformers, por lo que los adaptadores se cargan a través de su integración nativa de PEFT en lugar de la API load_lora_weights de diffusers. Dado que esa integración no depende del tipo de adaptador, las mismas llamadas cargan adaptadores LoRA, DoRA o cualquier otro adaptador PEFT (p. ej., la salida de SFTTrainer de TRL). Gestione los adaptadores directamente en el componente del 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")

Los adaptadores permanecen activos y sin fusionar: DiffusionGemma vincula los pesos base del codificador y del decodificador, por lo que fusionar un adaptador en ellos corruptaría ambas ramas.

Caché estática y compilación

El pipeline precarga el codificador una vez por bloque en una caché reutilizable (una DynamicCache por defecto). Pasar cache_implementation="static" utiliza en su lugar una StaticCache de forma fija, cuyas formas permiten torch.compile el decodificador con cudagraphs para una mayor aceleración (el pipeline marca cada paso y clona los logits para que la memoria de cudagraph no se sobrescriba); esta es la configuración mostrada en Uso. Elimina tanto la llamada a torch.compile como cache_implementation="static" para una ejecución más simple con caché dinámica.

Detención adaptativa

Un bloque suele converger antes de agotar todos los num_inference_steps, por lo que, de forma predeterminada, el pipeline congela cada ejemplo del lote una vez que su predicción de argmax es estable durante stability_threshold pasos y su entropía media por token cae por debajo de confidence_threshold (0.005, el valor utilizado por el checkpoint publicado). El bucle de desruido finaliza una vez que todos los ejemplos están congelados. Esto reduce aproximadamente a la mitad el número de pases del decodificador a una calidad equivalente y es el mayor factor individual de rendimiento. Pasa confidence_threshold=None para ejecutar siempre los num_inference_steps completos:

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

Callbacks

Los callbacks se ejecutan después de cada paso de desruido. Pasa callback_on_step_end_tensor_inputs para seleccionar qué tensores se incluyen en callback_kwargs; canvas (los tokens del bloque actual) y logits están disponibles. Devuelve {"canvas": ...} desde el callback para reemplazar el 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