Spaces:
Running
Running
File size: 2,998 Bytes
756b108 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | #!/usr/bin/env python3
"""
Download all model weights required for FASHN VTON.
Usage:
python scripts/download_weights.py --weights-dir ./weights
This will download:
- TryOnModel weights (model.safetensors) from HuggingFace
- DWPose ONNX models (yolox_l.onnx, dw-ll_ucoco_384.onnx)
- FashnHumanParser weights (auto-cached by HuggingFace)
"""
import argparse
import os
from huggingface_hub import hf_hub_download
def download_tryon_model(weights_dir: str) -> str:
"""Download TryOnModel weights from HuggingFace."""
print("Downloading TryOnModel weights...")
path = hf_hub_download(
repo_id="fashn-ai/fashn-vton-1.5",
filename="model.safetensors",
local_dir=weights_dir,
)
print(f" Saved to: {path}")
return path
def download_dwpose_models(weights_dir: str) -> str:
"""Download DWPose ONNX models from HuggingFace."""
dwpose_dir = os.path.join(weights_dir, "dwpose")
os.makedirs(dwpose_dir, exist_ok=True)
repo_id = "fashn-ai/DWPose"
filenames = ["yolox_l.onnx", "dw-ll_ucoco_384.onnx"]
for filename in filenames:
print(f"Downloading DWPose/{filename}...")
path = hf_hub_download(
repo_id=repo_id,
filename=filename,
local_dir=dwpose_dir,
)
print(f" Saved to: {path}")
return dwpose_dir
def download_human_parser() -> None:
"""Initialize FashnHumanParser to trigger weight download."""
print("Downloading FashnHumanParser weights...")
from fashn_human_parser import FashnHumanParser
# This will auto-download weights to HuggingFace cache if not present
_ = FashnHumanParser(device="cpu")
print(" Cached in HuggingFace hub cache")
def main():
parser = argparse.ArgumentParser(
description="Download all model weights for FASHN VTON",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Example:
python scripts/download_weights.py --weights-dir ./weights
After downloading, use the pipeline:
from fashn_vton import TryOnPipeline
pipeline = TryOnPipeline(weights_dir="./weights")
""",
)
parser.add_argument(
"--weights-dir",
type=str,
required=True,
help="Directory to save model weights",
)
args = parser.parse_args()
weights_dir = os.path.abspath(args.weights_dir)
os.makedirs(weights_dir, exist_ok=True)
print(f"\nDownloading weights to: {weights_dir}\n")
# Download all models
download_tryon_model(weights_dir)
print()
download_dwpose_models(weights_dir)
print()
download_human_parser()
print(f"""
Download complete!
Weights directory structure:
{weights_dir}/
├── model.safetensors
└── dwpose/
├── yolox_l.onnx
└── dw-ll_ucoco_384.onnx
Usage:
from fashn_vton import TryOnPipeline
pipeline = TryOnPipeline(weights_dir="{weights_dir}")
""")
if __name__ == "__main__":
main()
|