Skip to content

Commit d65600b

Browse files
committed
style: break url into base and file and factorize
1 parent 11ebe19 commit d65600b

File tree

1 file changed

+4
-5
lines changed

1 file changed

+4
-5
lines changed

src/pruna/evaluation/metrics/aesthetic_laion.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from pruna.logging.logger import pruna_logger
3131

3232
METRIC_AESTHETIC_LAION = "aesthetic_laion"
33+
LAION_AESTHETIC_BASE_URL = "https://raw.githubusercontent.com/LAION-AI/aesthetic-predictor/refs/heads/main"
3334

3435

3536
@MetricRegistry.register(METRIC_AESTHETIC_LAION)
@@ -176,13 +177,11 @@ def _get_aesthetic_model(self, clip_model="openai/clip-vit-large-patch14"):
176177
pruna_logger.error(f"Model {clip_model} is not supported by aesthetic predictor.")
177178
raise ValueError(f"Model {clip_model} is not supported by aesthetic predictor.")
178179

179-
path_to_model = cache_folder / f"sa_0_4_{clip_model}_linear.pth"
180+
model_file = f"sa_0_4_{clip_model}_linear.pth"
181+
path_to_model = cache_folder / model_file
180182
if not path_to_model.exists():
181183
cache_folder.mkdir(exist_ok=True, parents=True)
182-
url_model = (
183-
f"https://raw.githubusercontent.com/LAION-AI/aesthetic-predictor/refs/heads/main/sa_0_4_{clip_model}_linear.pth"
184-
)
185-
urlretrieve(url_model, path_to_model)
184+
urlretrieve(f"{LAION_AESTHETIC_BASE_URL}/{model_file}", path_to_model)
186185

187186
aesthetic_linear_head = nn.Linear(hidden_dim, 1)
188187
aesthetic_linear_head.load_state_dict(torch.load(path_to_model, map_location=self.device))

0 commit comments

Comments
 (0)