Update README.md

README.md

@@ -27,19 +27,22 @@ pip install -r requirements.txt
 
 ## 🚀 Usage
 For each HiT-SR model, we provide 2x, 3x, 4x upscaling versions:
-| Repo Name         |   | Model   |   | Upscale |
-|-------------------|---|---------|---|---------|
-| `XiangZ/hit-sir-2x` |   | HiT-SIR |   | 2x      |
-| `XiangZ/hit-sir-3x` |   | HiT-SIR |   | 3x      |
-| `XiangZ/hit-sir-4x` |   | HiT-SIR |   | 4x      |
-| `XiangZ/hit-sng-2x` |   | HiT-SNG |   | 2x      |
-| `XiangZ/hit-sng-3x` |   | HiT-SNG |   | 3x      |
-| `XiangZ/hit-sng-4x` |   | HiT-SNG |   | 4x      |
-| `XiangZ/hit-srf-2x` |   | HiT-SRF |   | 2x      |
-| `XiangZ/hit-srf-3x` |   | HiT-SRF |   | 3x      |
-| `XiangZ/hit-srf-4x` |   | HiT-SRF |   | 4x      |
+| Repo Name         |   | Model   |   | Training Set   |   | Upscale |
+|-------------------|---|---------|---|---------|---|---------|
+| `XiangZ/hit-sir-2x` |   | HiT-SIR |   | DIV2K      |   | 2x      |
+| `XiangZ/hit-sir-3x` |   | HiT-SIR |   | DIV2K      |   | 3x      |
+| `XiangZ/hit-sir-4x` |   | HiT-SIR |   | DIV2K      |   | 4x      |
+| `XiangZ/hit-sng-2x` |   | HiT-SNG |   | DIV2K      |   | 2x      |
+| `XiangZ/hit-sng-3x` |   | HiT-SNG |   | DIV2K      |   | 3x      |
+| `XiangZ/hit-sng-4x` |   | HiT-SNG |   | DIV2K      |   | 4x      |
+| `XiangZ/hit-srf-2x` |   | HiT-SRF |   | DIV2K      |   | 2x      |
+| `XiangZ/hit-srf-3x` |   | HiT-SRF |   | DIV2K      |   | 3x      |
+| `XiangZ/hit-srf-4x` |   | HiT-SRF |   | DIV2K      |   | 4x      |
+| `XiangZ/hit-srf-2x-df2k` |   | HiT-SRF |   | DF2K      |   | 2x      |
+| `XiangZ/hit-srf-3x-df2k` |   | HiT-SRF |   | DF2K      |   | 3x      |
+| `XiangZ/hit-srf-4x-df2k` |   | HiT-SRF |   | DF2K      |   | 4x      |
 
-To test the model (use hit-srf-4x as an example):
+To test the model (use hit-srf-4x-df2k as an example):
 ```
 from hit_sir_arch import HiT_SIR
 from hit_sng_arch import HiT_SNG
@@ -54,7 +57,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.ba
 model = HiT_SRF(upscale=4) 
 
 # load model (change repo_name according to your setting)
-repo_name = "XiangZ/hit-srf-4x"
+repo_name = "XiangZ/hit-srf-4x-df2k"
 model = model.from_pretrained(repo_name)
 model.to(device)