Update README.md

README.md

@@ -95,31 +95,21 @@ pip install stylizing-vit
 
 ```python
 import torch
-from huggingface_hub import hf_hub_download
 from stylizing_vit import create_model
 
-# Option 1: Load using .pth weights
-# weights_path = hf_hub_download(
-#     repo_id="sdoerrich97/stylizing_vit_base_[DATASET_SLUG]",
-#     filename="model.pth"
-# )
-
-# Option 2: Load using .safetensors weights (Recommended)
-weights_path = hf_hub_download(
-    repo_id="sdoerrich97/stylizing_vit_base_[DATASET_SLUG]",
-    filename="model.safetensors"
-)
-
 # Initialize the model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-# Create model will automatically detect .safetensors if passed
-model = create_model(backbone="base", weights=weights_path, train=False).to(device)
+
+# Load the model with pretrained weights
+# This automatically downloads the weights from the Hugging Face Hub
+model = create_model(backbone="base", weights="camelyon17wilds", train=False).to(device)
 model.eval()
 
 # Apply Style Transfer
 # content_img and style_img should be normalized torch tensors of shape (1, 3, 224, 224)
 # with torch.no_grad():
 #     stylized_img = model(content_img, style_img)
+``` stylized_img = model(content_img, style_img)
 ```
 
 ## Training Details