Update README

README.md

@@ -8,13 +8,8 @@ language:
 
 ## Model Summary
 
-**HistAug** is a lightweight transformer-based generator for **controllable latent-space augmentations** in the feature space of the [H-optimus-1 foundation model](https://huggingface.co/bioptimus/H-optimus-1).  
 
-Instead of applying costly image-space augmentations on millions of WSI patches, HistAug operates **directly on patch embeddings** extracted from H-optimus-1.  
-
-By conditioning on explicit transformation parameters (e.g., hue shift, erosion, HED color transform), HistAug generates realistic augmented embeddings while preserving semantic content.  
-
-In practice, the H-optimus-1 variant of HistAug can reconstruct the corresponding “true” augmented features with an average cosine similarity of **about 81%** at **10X, 20X, and 40X magnification**.
+**HistAug** is a lightweight transformer-based generator for **controllable latent-space augmentations** in the feature space of the [H-optimus-1 foundation model](https://huggingface.co/bioptimus/H-optimus-1). Instead of applying costly image-space augmentations on millions of WSI patches, HistAug operates **directly on patch embeddings** extracted from a given foundation model(here H-optimus-1). By conditioning on explicit transformation parameters (e.g., hue shift, erosion, HED color transform), HistAug generates realistic augmented embeddings while preserving semantic content. In practice, the H-optimus-1 variant of HistAug can reconstruct the corresponding ground-truth augmented embeddings with an average cosine similarity of **about 90%** at **10X, 20X, and 40X magnification**.
 
 This enables training of Multiple Instance Learning (MIL) models with:  
 - ⚡ **Fast augmentation** 
@@ -24,7 +19,8 @@ This enables training of Multiple Instance Learning (MIL) models with:
 Need HistAug for a different foundation model? Explore the full collection: [**HistAug models collection**](https://huggingface.co/collections/sofieneb/histaug-models-68a334437f71d35c7037a54e).
 
 
-📄 **Paper**: *Controllable Latent Space Augmentation for Digital Pathology* (Boutaj *et al.*, 2025)  
+📄 **Paper**: [*Controllable Latent Space Augmentation for Digital Pathology* (Boutaj *et al.*, 2025)](https://arxiv.org/abs/2508.14588)
+
 
 
 
@@ -38,14 +34,16 @@ You can load the model from the Hub with Hugging Face’s `transformers`:
 import torch
 from transformers import AutoModel
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Load HistAug (H-optimus-1 latent augmentation model)
 model_id = "sofieneb/histaug-hoptimus1"
-model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device)
 
 # Example: patch embeddings from H-optimus-1
 num_patches = 50000
 embedding_dim = 1536
-patch_embeddings = torch.randn((num_patches, embedding_dim), device="cuda")
+patch_embeddings = torch.randn((num_patches, embedding_dim), device=device)
 
 # Sample augmentation parameters
 # mode="wsi_wise" applies the same transformation across the whole slide
@@ -197,22 +195,54 @@ for bag_features, label in loader:  # bag_features: (num_patches, embed_dim)
 
 ## Offline usage (HPC clusters without internet)
 
-If your cluster doesn’t allow internet access on compute nodes, **pre-download the model on the front-end/login node** (which has internet), so it’s cached locally, then run jobs offline:
+If your cluster doesn’t allow internet access on compute nodes, you have two ways to use HistAug offline:
+
+1. **Rely on the cache**: pre-download the model on the front-end/login node (with internet), so it’s cached locally, then run jobs offline:
 
 ```bash
 # On the front-end/login node (with internet):
-python -c "from transformers import AutoModel; AutoModel.from_pretrained('sofieneb/histaug-hoptimus1', trust_remote_code=True)"
+python -c "from transformers import AutoModel; AutoModel.from_pretrained('sofieneb/histaug-uni', trust_remote_code=True)"
 
 # On your compute job (no internet):
 export HF_HUB_OFFLINE=1
 export TRANSFORMERS_OFFLINE=1
-```
+````
 
 This prevents unnecessary network calls and ensures `transformers` loads HistAug from the local cache.
 
+2. **Explicitly download with `hf download`**: instead of relying on the cache variables above, you can download the model files manually on the front-end/login node (with internet) and always point to the local folder:
+
+```bash
+# On the front-end/login node (with internet):
+hf download sofieneb/histaug-virchow2 --local-dir ./histaug-virchow2
+```
+
+Then load the model from that directory in your script:
+
+```python
+from transformers import AutoModel
+
+cross_transformer = AutoModel.from_pretrained(
+    "./histaug-virchow2",   # local path instead of hub ID
+    trust_remote_code=True,
+    local_files_only=True
+)
+```
 ---
-## Paper
+## Citation
+If our work contributes to your research, or if you incorporate part of this code, please consider citing our paper:
+
+```bibtex
+@misc{boutaj2025controllablelatentspaceaugmentation,
+      title={Controllable Latent Space Augmentation for Digital Pathology}, 
+      author={Sofiène Boutaj and Marin Scalbert and Pierre Marza and Florent Couzinie-Devy and Maria Vakalopoulou and Stergios Christodoulidis},
+      year={2025},
+      eprint={2508.14588},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2508.14588}, 
+}
+```
 
-*Controllable Latent Space Augmentation for Digital Pathology* (Boutaj *et al.*, 2025).