theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-23 17:38:10 +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/imagereward.py
View File

@@ -188,15 +188,15 @@ class ImageReward(torch.nn.Module):
class ImageRewardScore:
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)
model_path = MODEL_PATHS.get("imagereward")
med_config = MODEL_PATHS.get("med_config")
model_path = path.get("imagereward")
med_config = path.get("med_config")
state_dict = load_file(model_path)
self.model = ImageReward(device=self.device, med_config=med_config).to(self.device)
self.model.load_state_dict(state_dict, strict=False)