dg845/DiffusionGemma-diffusers

diffusers-native checkpoint for the google/diffusiongemma-26B-A4B-it discrete diffusion LLM.

You can use the model as follows. Note that you need transformers>=5.12.0 for the underlying DiffusionGemmaForBlockDiffusion model.

import torch
from diffusers import DiffusionGemmaPipeline

pipe = DiffusionGemmaPipeline.from_pretrained(
    "dg845/DiffusionGemma-diffusers",
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")

# Compile the decoder model for faster inference
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead"),

output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
    generator=torch.Generator("cuda").manual_seed(42),
)
print(output.texts[0])