Update Garage/models/PowerPaintModel.py

Garage/models/PowerPaintModel.py

@@ -10,6 +10,7 @@ from .PowerPaint.utils import TokenizerWrapper, add_tokens
 from safetensors.torch import load_model
 from typing import Dict
 from PIL import Image
+from huggingface_hub import hf_hub_download
 
 class PowerPaintModel:
     def __init__(self, 
@@ -83,14 +84,18 @@ class PowerPaintModel:
             initialize_tokens=["a", "a", "a"],
             num_vectors_per_token=10
         )
-        
+
+        brushnet_weights_path = hf_hub_download(repo_id="anvilarth/PP_Brushnet", filename="diffusion_pytorch_model.safetensors")
+        brushnet_encoder_path = hf_hub_download(repo_id="anvilarth/PP_Brushnet", filename="pytorch_model.bin")
+
+
         load_model(
             pipe.brushnet,
-            os.path.join( "JunhaoZhuang/PowerPaint-v2-1", "PowerPaint_Brushnet/diffusion_pytorch_model.safetensors")
+            brushnet_weights_path
         )
         
         pipe.text_encoder_brushnet.load_state_dict(
-            torch.load(os.path.join("JunhaoZhuang/PowerPaint-v2-1", "PowerPaint_Brushnet/pytorch_model.bin")), strict=False
+            torch.load(brushnet_encoder_path), strict=False
         )
 
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)