Update Garage/models/PowerPaintModel.py

Garage/models/PowerPaintModel.py

@@ -54,7 +54,7 @@ class PowerPaintModel:
         brushnet = BrushNetModel.from_unet(unet)
         base_model_path = os.path.join(self.model_name, "realisticVisionV60B1_v51VAE")
         pipe = StableDiffusionPowerPaintBrushNetPipeline.from_pretrained(
-            self.model_name,
+            "JunhaoZhuang/PowerPaint-v2-1",
             subfolder="realisticVisionV60B1_v51VAE",
             brushnet=brushnet,
             text_encoder_brushnet=text_encoder_brushnet,
@@ -64,15 +64,15 @@ class PowerPaintModel:
         )
         
         pipe.unet = UNet2DConditionModel.from_pretrained(
-            base_model_path,
-            subfolder="unet",
+             "JunhaoZhuang/PowerPaint-v2-1",
+            subfolder="realisticVisionV60B1_v51VAE/unet",
             revision=None,
             torch_dtype=self.weight_dtype
         )
         
         pipe.tokenizer = TokenizerWrapper(
-            from_pretrained=base_model_path,
-            subfolder="tokenizer",
+            from_pretrained= "JunhaoZhuang/PowerPaint-v2-1",
+            subfolder="realisticVisionV60B1_v51VAE/tokenizer",
             revision=None,
             torch_type=self.weight_dtype
         )
@@ -87,11 +87,11 @@ class PowerPaintModel:
         
         load_model(
             pipe.brushnet,
-            os.path.join(base_model_path, "PowerPaint_Brushnet/diffusion_pytorch_model.safetensors")
+            os.path.join( "JunhaoZhuang/PowerPaint-v2-1", "PowerPaint_Brushnet/diffusion_pytorch_model.safetensors")
         )
         
         pipe.text_encoder_brushnet.load_state_dict(
-            torch.load(os.path.join(base_model_path, "PowerPaint_Brushnet/pytorch_model.bin")), strict=False
+            torch.load(os.path.join("JunhaoZhuang/PowerPaint-v2-1", "PowerPaint_Brushnet/pytorch_model.bin")), strict=False
         )
 
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)