feat: download custom weights from HF Hub at runtime using hf_hub_download

infer_full.py

@@ -13,6 +13,39 @@ from ref_encoder.reference_unet import ref_unet
 from utils.pipeline import StableHairPipeline
 from utils.pipeline_cn import StableDiffusionControlNetPipeline
 
+def _resolve_weight(prefix_path: str, filename: str) -> str:
+    """Resolve a weight path, downloading from Hugging Face Hub if needed.
+
+    prefix_path can be either a local directory (e.g., ./models/stage2)
+    or a hub path like Org/Repo/subfolder. When it looks like a hub path,
+    we download the file via hf_hub_download using repo_id Org/Repo and
+    subfolder the remaining segments.
+    """
+    # Try local first
+    local_path = os.path.join(prefix_path, filename)
+    if os.path.exists(local_path):
+        return local_path
+
+    # Attempt Hub download
+    try:
+        from huggingface_hub import hf_hub_download
+
+        parts = prefix_path.strip("/").split("/")
+        if len(parts) >= 2:
+            repo_id = "/".join(parts[:2])
+            subfolder = "/".join(parts[2:]) if len(parts) > 2 else None
+            downloaded = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                subfolder=subfolder,
+                token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
+            )
+            return downloaded
+    except Exception as exc:  # noqa: WPS440
+        raise RuntimeError(f"Failed to fetch {filename} from hub ({prefix_path}): {exc}")
+
+    raise FileNotFoundError(f"Weight not found locally and not a valid hub path: {prefix_path}/{filename}")
+
 def concatenate_images(image_files, output_file, type="pil"):
     if type == "np":
         image_files = [Image.fromarray(img) for img in image_files]
@@ -36,7 +69,8 @@ class StableHair:
         ### Load controlnet
         unet = UNet2DConditionModel.from_pretrained(self.config.pretrained_model_path, subfolder="unet").to(device)
         controlnet = ControlNetModel.from_unet(unet).to(device)
-        _state_dict = torch.load(os.path.join(self.config.pretrained_folder, self.config.controlnet_path))
+        controlnet_weight_path = _resolve_weight(self.config.pretrained_folder, self.config.controlnet_path)
+        _state_dict = torch.load(controlnet_weight_path, map_location="cpu")
         controlnet.load_state_dict(_state_dict, strict=False)
         controlnet.to(weight_dtype)
 
@@ -51,15 +85,24 @@ class StableHair:
 
         ### load Hair encoder/adapter
         self.hair_encoder = ref_unet.from_pretrained(self.config.pretrained_model_path, subfolder="unet").to(device)
-        _state_dict = torch.load(os.path.join(self.config.pretrained_folder, self.config.encoder_path))
+        encoder_weight_path = _resolve_weight(self.config.pretrained_folder, self.config.encoder_path)
+        _state_dict = torch.load(encoder_weight_path, map_location="cpu")
         self.hair_encoder.load_state_dict(_state_dict, strict=False)
         self.hair_adapter = adapter_injection(self.pipeline.unet, device=self.device, dtype=torch.float16, use_resampler=False)
-        _state_dict = torch.load(os.path.join(self.config.pretrained_folder, self.config.adapter_path))
+        adapter_weight_path = _resolve_weight(self.config.pretrained_folder, self.config.adapter_path)
+        _state_dict = torch.load(adapter_weight_path, map_location="cpu")
         self.hair_adapter.load_state_dict(_state_dict, strict=False)
 
         ### load bald converter
         bald_converter = ControlNetModel.from_unet(unet).to(device)
-        _state_dict = torch.load(self.config.bald_converter_path)
+        # bald_converter_path may be a local full path or a hub-like path
+        if os.path.exists(self.config.bald_converter_path):
+            bald_weight_path = self.config.bald_converter_path
+        else:
+            prefix, filename = os.path.split(self.config.bald_converter_path)
+            bald_weight_path = _resolve_weight(prefix, filename)
+
+        _state_dict = torch.load(bald_weight_path, map_location="cpu")
         bald_converter.load_state_dict(_state_dict, strict=False)
         bald_converter.to(dtype=weight_dtype)
         del unet