theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-22 16:50:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

add new quality metric

Browse Source
This commit is contained in:
YunhongLu-ZJU
2025-02-17 14:42:20 +08:00
parent 77d0f4d297
commit 991ba162bd
69 changed files with 88 additions and 1461 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
diffsynth/extensions/QualityMetric/pickscore.py
View File

@@ -6,15 +6,15 @@ import os
from .config import MODEL_PATHS
class PickScore:
def __init__(self, device: Union[str, torch.device]):
def __init__(self, device: Union[str, torch.device], path: str = MODEL_PATHS):
"""Initialize the Selector with a processor and model.
Args:
device (Union[str, torch.device]): The device to load the model on.
"""
self.device = device if isinstance(device, torch.device) else torch.device(device)
processor_name_or_path = MODEL_PATHS.get("clip")
model_pretrained_name_or_path = MODEL_PATHS.get("pickscore")
processor_name_or_path = path.get("clip")
model_pretrained_name_or_path = path.get("pickscore")
self.processor = AutoProcessor.from_pretrained(processor_name_or_path)
self.model = AutoModel.from_pretrained(model_pretrained_name_or_path).eval().to(self.device)