fix weight path issue

wrapper.py

@@ -183,7 +183,7 @@ def peca(pipeline, save_path = "./weight", n_layer = 10):
     return encoder, feature_encoder.get_text_feature, size, num_inference_steps, skip
 
 class DrUM(DiffusionPipeline):
-    def __init__(self, pipeline, repo_id = "Burf/DrUM"):
+    def __init__(self, pipeline, repo_id = "Burf/DrUM", weight = None):
         """
         DrUM for various diffusion models
         
@@ -194,7 +194,7 @@ class DrUM(DiffusionPipeline):
         self.pipeline = pipeline
         self.repo_id = repo_id
         
-        self.adapter, self.feature_encoder, self.size, self.num_inference_steps, self.skip = self.load_peca(pipeline, repo_id)
+        self.adapter, self.feature_encoder, self.size, self.num_inference_steps, self.skip = self.load_peca(pipeline, repo_id, weight)
     
     @classmethod
     def from_pretrained(cls, model_id, repo_id = "Burf/DrUM", torch_dtype = torch.bfloat16, device = "cuda"):
@@ -220,7 +220,7 @@ class DrUM(DiffusionPipeline):
         #pipeline.safety_checker = lambda images, clip_input: (images, [False] * len(images))
         return cls(pipeline, repo_id)
     
-    def load_weight(self, pipeline, repo_id = "Burf/DrUM"):
+    def load_weight(self, pipeline, repo_id = "Burf/DrUM", weight = None):
         name = pipeline.config._name_or_path.split("/")[-1].lower()
         
         weights = []
@@ -236,12 +236,16 @@ class DrUM(DiffusionPipeline):
             weights = ["L.safetensors"]
         
         for weight_file in weights:
-            safetensor_path = hf_hub_download(repo_id = repo_id, filename = weight_file)
-            weight_path = os.path.dirname(safetensor_path)
+            if isinstance(weight, str) and os.path.exists(os.path.join(weight, weight_file)):
+                weight_path = weight
+                break
+            else:
+                safetensor_path = hf_hub_download(repo_id = repo_id, filename = "weight/" + weight_file)
+                weight_path = os.path.dirname(safetensor_path)
         return weight_path
     
-    def load_peca(self, pipeline, repo_id = "Burf/DrUM"):
-        adapter, feature_encoder, size, num_inference_steps, skip = peca(pipeline, save_path = self.load_weight(pipeline, repo_id))
+    def load_peca(self, pipeline, repo_id = "Burf/DrUM", weight = None):
+        adapter, feature_encoder, size, num_inference_steps, skip = peca(pipeline, save_path = self.load_weight(pipeline, repo_id, weight))
         return adapter, feature_encoder, size, num_inference_steps, skip
     
     def __call__(self, prompt, ref = None, weight = None, alpha = 0.3, skip = None, sampling = False, seed = 42,