Update README.md

README.md

@@ -32,16 +32,27 @@ This model is intended for:
 ## 🧪 Example
 
 ```python
-import torch
-from custom_rcan import RcanModelM  # Replace with your actual module
+from super_image import RcanModel, RcanConfig
 from huggingface_hub import hf_hub_download
+import torch
+
+path = hf_hub_download(repo_id="lschmidt/rcan-dsc", filename="rcan_model.py")
+
+# Dynamically load it
+exec(open(path).read())
 
-# Load model
-model = RcanModelM(in_channels=2, scaling_factor=4)
-state_dict = torch.load(hf_hub_download(repo_id="your-username/rcan-dsc", filename="pytorch_model_4x.pt"))
+
+# Now you can call the function directly
+model = load_rcan()
+
+# load config
+config, _ = RcanConfig.from_pretrained("lschmidt/rcan-dsc")
+
+# load pretrained weights
+state_dict_path = hf_hub_download(repo_id="lschmidt/rcan-dsc", filename="pytorch_model_4x.pt")
+state_dict = torch.load(state_dict_path, map_location="cpu")
 model.load_state_dict(state_dict, strict=False)
 
-# Dummy input: shape (B, C=2, H, W)
-x = torch.randn(1, 2, 64, 64)
-y = model(x)
-print(y.shape)  # (1, 2, 256, 256) for scale=4
+inputs = torch.randn(1, 2, 10, 10)
+output = model(inputs)
+