from transformers import AutoTokenizer, AutoModel
from dream.modeling_dream import DreamModel

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("{model_name}", trust_remote_code=True)
model = DreamModel.from_pretrained("{model_name}", trust_remote_code=True)

# Use for diffusion generation
inputs = tokenizer("Your prompt here", return_tensors="pt")
output = model.planned_diffusion_generate(inputs, max_length=512)
generated_text = tokenizer.decode(output.sequences[0], skip_special_tokens=True)
print(generated_text)

Citation

If you use this model, please cite the Planned Diffusion paper:

bibtex

@misc{israel2025planneddiffusion,
      title={Planned Diffusion}, 
      author={Daniel Israel and Tian Jin and Ellie Cheng and Guy Van den Broeck and Aditya Grover and Suvinay Subramanian and Michael Carbin},
      year={2025},
      eprint={2510.18087},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2510.18087}, 
}