Update README

README.md

@@ -8,7 +8,7 @@ language:
 
 ## Model Summary
 
-**HistAug** is a lightweight transformer-based generator for **controllable latent-space augmentations** in the feature space of the [CONCH foundation model](https://www.nature.com/articles/s41591-024-02856-4). 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 CONCH). 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 CONCH 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**.
+**HistAug** is a lightweight transformer-based generator for **controllable latent-space augmentations** in the feature space of the [CONCH foundation model](https://www.nature.com/articles/s41591-024-02856-4). 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 CONCH). 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 CONCH variant of HistAug can reconstruct the corresponding ground-truth augmented embeddings with an average cosine similarity of **about 93%** at **10X, 20X, and 40X magnification**.
 
 This enables training of Multiple Instance Learning (MIL) models with:  
 - ⚡ **Fast augmentation** 
@@ -194,39 +194,52 @@ 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, you have two ways to use HistAug offline:
+If compute nodes don’t have internet, **always** run jobs with the offline flags to **prevent unnecessary network calls** and force local loads:
 
-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 your compute job (no internet):
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+```
+
+Prepare the model **in advance** on a front-end/login node (with internet), then choose **either** approach below.
+
+### Option — Warm the cache (simplest)
 
 ```bash
 # On the front-end/login node (with internet):
 python -c "from transformers import AutoModel; AutoModel.from_pretrained('sofieneb/histaug-conch', trust_remote_code=True)"
+```
 
-# On your compute job (no internet):
-export HF_HUB_OFFLINE=1
-export TRANSFORMERS_OFFLINE=1
-````
+Then in your offline job/script:
 
-This prevents unnecessary network calls and ensures `transformers` loads HistAug from the local cache.
+```python
+from transformers import AutoModel
+model = AutoModel.from_pretrained(
+    "sofieneb/histaug-conch",
+    trust_remote_code=True,
+    local_files_only=True,  # uses local cache only
+)
+```
 
-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:
+### Option — Download to a local folder with `hf download`
 
 ```bash
 # On the front-end/login node (with internet):
 hf download sofieneb/histaug-conch --local-dir ./histaug-conch
 ```
 
-Then load the model from that directory in your script:
+Then in your offline job/script:
 
 ```python
 from transformers import AutoModel
-
 cross_transformer = AutoModel.from_pretrained(
     "./histaug-conch",   # local path instead of hub ID
     trust_remote_code=True,
-    local_files_only=True
+    local_files_only=True,  # uses local files only
 )
 ```
+
 ---
 ## Citation
 If our work contributes to your research, or if you incorporate part of this code, please consider citing our paper: