init

README.md

@@ -8,6 +8,11 @@ sdk_version: 5.42.0
 app_file: app.py
 pinned: false
 short_description: HPSv3 demo
+preload_from_hub:
+  - MizzenAI/HPSv3 HPSv3.safetensors
+  - xswu/HPSv2 HPS_v2.1_compressed.pt
+  - yuvalkirstain/PickScore_v1
+  - laion/CLIP-ViT-H-14-laion2B-s32B-b79K
 ---
 
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -4,6 +4,7 @@ import os
 import sys
 from PIL import Image
 import uuid
+import huggingface_hub
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 from hpsv3.inference import HPSv3RewardInferencer
@@ -30,7 +31,7 @@ MODEL_CONFIGS = {
     },
     "HPSv2": {
         "name": "HPSv2",
-        "checkpoint_path": "your_path_to_HPS_v2_compressed.pt",
+        "checkpoint_path": "xswu/HPSv2/HPS_v2.1_compressed.pt",
         "type": "hpsv2"
     },
     "ImageReward": {
@@ -40,12 +41,12 @@ MODEL_CONFIGS = {
     },
     "PickScore": {
         "name": "PickScore",
-        "checkpoint_path": "your_path_to_pickscore",
+        "checkpoint_path": "yuvalkirstain/PickScore_v1",
         "type": "pickscore"
     },
     "CLIP": {
         "name": "CLIP ViT-H-14",
-        "checkpoint_path": "/preflab/models/CLIP-ViT-H-14-laion2B-s32B-b79K",
+        "checkpoint_path": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
         "type": "clip"
     }
 }
@@ -73,8 +74,10 @@ def load_model(model_key, update_status_fn=None):
     
     try:
         if config["type"] == "hpsv3":
+            checkpoint_path = huggingface_hub.hf_hub_download("MizzenAI/HPSv3", 'HPSv3.safetensors', repo_type='model')
             model = HPSv3RewardInferencer(
                 device=DEVICE, 
+                checkpoint_path=checkpoint_path
             )
         elif config["type"] == "hpsv2":
             model_obj, preprocess_train, preprocess_val = create_model_and_transforms(
@@ -96,7 +99,8 @@ def load_model(model_key, update_status_fn=None):
                 with_score_predictor=False,
                 with_region_predictor=False
             )
-            checkpoint = torch.load(config["checkpoint_path"], map_location=DEVICE, weights_only=False)
+            checkpoint_path = huggingface_hub.hf_hub_download("xswu/HPSv2", 'HPS_v2.1_compressed.pt', repo_type='model')
+            checkpoint = torch.load(checkpoint_path, map_location=DEVICE, weights_only=False)
             model_obj.load_state_dict(checkpoint['state_dict'])
             model_obj = model_obj.to(DEVICE).eval()
             tokenizer = get_tokenizer('ViT-H-14')

requirements.txt

@@ -191,4 +191,5 @@ xxhash==3.5.0
 yarl==1.20.1
 zipp==3.22.0
 # flash-attn==2.7.4.post1
-hpsv3==1.0.0
+hpsv3==1.0.0
+hpsv2