Use DiffusionLM pipeline in usage example

README.md

@@ -52,24 +52,27 @@ weight `1/σ(t)`, is what turned word-salad into coherent stories.
 
 ## Usage
 
-Install the model code from the GitHub repo, then:
+Install the model code from the
+[GitHub repo](https://github.com/tchauffi/diffusionlm-from-scratch), then
+generate stories in two lines — `DiffusionLM` bundles the model, tokenizer, and
+absorbing-state scheduler:
 
 ```python
-import torch
-from huggingface_hub import hf_hub_download
-from transformers import PreTrainedTokenizerFast
-from diffusionlm_from_scratch.model import DiT, DiTConfig
+from diffusionlm_from_scratch import DiffusionLM
 
-repo = "tchauffi/diffusionlm-from-scratch"
+lm = DiffusionLM.from_pretrained("tchauffi/diffusionlm-from-scratch")
+for story in lm.generate(n=4, seq_len=80, temperature=0.9):
+    print(story)
+```
 
-ckpt_path = hf_hub_download(repo, "final.pt")
-ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+`generate` exposes the sampler knobs (`order`, `steps`, `corrector_frac`,
+`confidence_threshold`, …). For lower-level access, load just the model:
 
-model = DiT(DiTConfig(**ck["config"]))
-model.load_state_dict(ck["model"])   # EMA weights ("raw" also available)
-model.eval()
+```python
+from diffusionlm_from_scratch.model import DiT
 
-tokenizer = PreTrainedTokenizerFast.from_pretrained(repo)
+model = DiT.from_pretrained("tchauffi/diffusionlm-from-scratch")  # downloads final.pt
+# the raw checkpoint carries ck["config"], ck["model"] (EMA), and ck["raw"].
 ```
 
 See [`scripts/capture_trajectories.py`](https://github.com/tchauffi/diffusionlm-from-scratch/blob/main/scripts/capture_trajectories.py)