| from modelscope import snapshot_download | |
| from typing_extensions import Literal, TypeAlias | |
| import os | |
| from diffsynth.extensions.ImageQualityMetric.aesthetic import AestheticScore | |
| from diffsynth.extensions.ImageQualityMetric.imagereward import ImageRewardScore | |
| from diffsynth.extensions.ImageQualityMetric.pickscore import PickScore | |
| from diffsynth.extensions.ImageQualityMetric.clip import CLIPScore | |
| from diffsynth.extensions.ImageQualityMetric.hps import HPScore_v2 | |
| from diffsynth.extensions.ImageQualityMetric.mps import MPScore | |
| preference_model_id: TypeAlias = Literal[ | |
| "ImageReward", | |
| "Aesthetic", | |
| "PickScore", | |
| "CLIP", | |
| "HPSv2", | |
| "HPSv2.1", | |
| "MPS", | |
| ] | |
| model_dict = { | |
| "ImageReward": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "ImageReward/ImageReward.safetensors", | |
| "ImageReward/med_config.json", | |
| "bert-base-uncased/config.json", | |
| "bert-base-uncased/model.safetensors", | |
| "bert-base-uncased/tokenizer.json", | |
| "bert-base-uncased/tokenizer_config.json", | |
| "bert-base-uncased/vocab.txt", | |
| ], | |
| "load_path": { | |
| "imagereward": "ImageReward/ImageReward.safetensors", | |
| "med_config": "ImageReward/med_config.json", | |
| "bert_model_path": "bert-base-uncased", | |
| }, | |
| "model_class": ImageRewardScore | |
| }, | |
| "Aesthetic": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "aesthetic-predictor/sac+logos+ava1-l14-linearMSE.safetensors", | |
| "clip-vit-large-patch14/config.json", | |
| "clip-vit-large-patch14/merges.txt", | |
| "clip-vit-large-patch14/model.safetensors", | |
| "clip-vit-large-patch14/preprocessor_config.json", | |
| "clip-vit-large-patch14/special_tokens_map.json", | |
| "clip-vit-large-patch14/tokenizer.json", | |
| "clip-vit-large-patch14/tokenizer_config.json", | |
| "clip-vit-large-patch14/vocab.json", | |
| ], | |
| "load_path": { | |
| "aesthetic_predictor": "aesthetic-predictor/sac+logos+ava1-l14-linearMSE.safetensors", | |
| "clip-large": "clip-vit-large-patch14", | |
| }, | |
| "model_class": AestheticScore | |
| }, | |
| "PickScore": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "PickScore_v1/*", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/config.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/merges.txt", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/preprocessor_config.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/special_tokens_map.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/tokenizer.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/tokenizer_config.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/vocab.json", | |
| ], | |
| "load_path": { | |
| "pickscore": "PickScore_v1", | |
| "clip": "CLIP-ViT-H-14-laion2B-s32B-b79K", | |
| }, | |
| "model_class": PickScore | |
| }, | |
| "CLIP": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin", | |
| "bpe_simple_vocab_16e6.txt.gz", | |
| ], | |
| "load_path": { | |
| "open_clip": "CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin", | |
| "open_clip_bpe": "bpe_simple_vocab_16e6.txt.gz", | |
| }, | |
| "model_class": CLIPScore | |
| }, | |
| "HPSv2": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "HPS_v2/HPS_v2_compressed.safetensors", | |
| "bpe_simple_vocab_16e6.txt.gz", | |
| ], | |
| "load_path": { | |
| "hpsv2": "HPS_v2/HPS_v2_compressed.safetensors", | |
| "open_clip_bpe": "bpe_simple_vocab_16e6.txt.gz", | |
| }, | |
| "model_class": HPScore_v2, | |
| "extra_kwargs": {"model_version": "v2"} | |
| }, | |
| "HPSv2.1": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "HPS_v2/HPS_v2.1_compressed.safetensors", | |
| "bpe_simple_vocab_16e6.txt.gz", | |
| ], | |
| "load_path": { | |
| "hpsv2.1": "HPS_v2/HPS_v2.1_compressed.safetensors", | |
| "open_clip_bpe": "bpe_simple_vocab_16e6.txt.gz", | |
| }, | |
| "model_class": HPScore_v2, | |
| "extra_kwargs": {"model_version": "v21"} | |
| }, | |
| "MPS": { | |
| "model_id": "DiffSynth-Studio/QualityMetric_reward_pretrained", | |
| "allow_file_pattern": [ | |
| "MPS_overall_checkpoint/MPS_overall_checkpoint_diffsynth.safetensors", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/config.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/merges.txt", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/preprocessor_config.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/special_tokens_map.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/tokenizer.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/tokenizer_config.json", | |
| "CLIP-ViT-H-14-laion2B-s32B-b79K/vocab.json", | |
| ], | |
| "load_path": { | |
| "mps": "MPS_overall_checkpoint/MPS_overall_checkpoint_diffsynth.safetensors", | |
| "clip": "CLIP-ViT-H-14-laion2B-s32B-b79K", | |
| }, | |
| "model_class": MPScore | |
| }, | |
| } | |
| def download_preference_model(model_name: preference_model_id, cache_dir="models"): | |
| metadata = model_dict[model_name] | |
| snapshot_download(model_id=metadata["model_id"], allow_file_pattern=metadata["allow_file_pattern"], cache_dir=cache_dir) | |
| load_path = metadata["load_path"] | |
| load_path = {key: os.path.join(cache_dir, metadata["model_id"], path) for key, path in load_path.items()} | |
| return load_path | |
| def load_preference_model(model_name: preference_model_id, device = "cuda", path = None): | |
| model_class = model_dict[model_name]["model_class"] | |
| extra_kwargs = model_dict[model_name].get("extra_kwargs", {}) | |
| preference_model = model_class(device=device, path=path, **extra_kwargs) | |
| return preference_model | |