Spaces:
Running
Running
File size: 1,680 Bytes
2c50826 4f41410 2c50826 4f41410 2c50826 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import json
import os
from pathlib import Path
from typing import Dict, Iterator, List, Tuple
import huggingface_hub
class HPSPrompts:
def __init__(self):
super().__init__()
self.hps_prompt_files = ['anime.json', 'concept-art.json', 'paintings.json', 'photo.json']
self._download_benchmark_prompts()
self.prompts: Dict[str, str] = {}
self._size = 0
for file in self.hps_prompt_files:
category = file.replace('.json', '')
with open(os.path.join('downloads/hps', file), 'r') as f:
prompts = json.load(f)
for i, prompt in enumerate(prompts):
if i == 100:
break
filename = f"{category}_{i:03d}.png"
self.prompts[filename] = prompt
self._size += 1
def __iter__(self) -> Iterator[Tuple[str, Path]]:
for filename, prompt in self.prompts.items():
yield prompt, Path(filename)
@property
def name(self) -> str:
return "hps"
@property
def size(self) -> int:
return self._size
def _download_benchmark_prompts(self) -> None:
folder_name = Path('downloads/hps')
folder_name.mkdir(parents=True, exist_ok=True)
for file in self.hps_prompt_files:
file_name = huggingface_hub.hf_hub_download("zhwang/HPDv2", file, subfolder="benchmark", repo_type="dataset")
if not os.path.exists(os.path.join(folder_name, file)):
os.symlink(file_name, os.path.join(folder_name, file))
@property
def metrics(self) -> List[str]:
return ["hps"]
|