135 lines
5.1 KiB
Python
135 lines
5.1 KiB
Python
import argparse
|
|
import os
|
|
import shutil
|
|
import glob
|
|
from huggingface_hub import snapshot_download, upload_folder, create_repo
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Download a HF Diffusers repo and replace its transformer weights.")
|
|
parser.add_argument(
|
|
"--repo_id",
|
|
type=str,
|
|
required=True,
|
|
help="The Hugging Face repository ID to download (e.g., 'stabilityai/stable-diffusion-xl-base-1.0')."
|
|
)
|
|
parser.add_argument(
|
|
"--local_dir",
|
|
type=str,
|
|
required=True,
|
|
help="The local directory where the repository will be downloaded and modified."
|
|
)
|
|
parser.add_argument(
|
|
"--checkpoint_dir",
|
|
type=str,
|
|
required=True,
|
|
help="The directory containing the new transformer model weights (checkpoints) to inject."
|
|
)
|
|
parser.add_argument(
|
|
"--component_name",
|
|
type=str,
|
|
default="transformer",
|
|
help="The name of the component subfolder in the diffusers repo to replace weights for (default: 'transformer')."
|
|
)
|
|
parser.add_argument(
|
|
"--weight_file_name",
|
|
type=str,
|
|
default="diffusion_pytorch_model.safetensors",
|
|
help="The target weight filename in the component directory (default: 'diffusion_pytorch_model.safetensors')."
|
|
)
|
|
parser.add_argument(
|
|
"--ignore_patterns",
|
|
nargs="+",
|
|
default=[],
|
|
help="Patterns to ignore when downloading the repo (passed to snapshot_download)."
|
|
)
|
|
parser.add_argument(
|
|
"--push_to_hub",
|
|
action="store_true",
|
|
help="If set, upload the modified local directory back to the Hugging Face Hub (to --repo_id)."
|
|
)
|
|
parser.add_argument(
|
|
"--hub_commit_message",
|
|
type=str,
|
|
default="Update transformer weights",
|
|
help="Commit message for the Hub upload."
|
|
)
|
|
parser.add_argument(
|
|
"--upload_repo_id",
|
|
type=str,
|
|
default=None,
|
|
help="The Hugging Face repository ID to upload to. If not provided, defaults to --repo_id. Creates the repo if it doesn't exist."
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f"Downloading repo '{args.repo_id}' to '{args.local_dir}'...")
|
|
snapshot_download(
|
|
repo_id=args.repo_id,
|
|
local_dir=args.local_dir,
|
|
ignore_patterns=args.ignore_patterns,
|
|
local_dir_use_symlinks=False # We want actual files to modify them
|
|
)
|
|
print("Download complete.")
|
|
|
|
target_component_dir = os.path.join(args.local_dir, args.component_name)
|
|
if not os.path.exists(target_component_dir):
|
|
print(f"Warning: Target component directory '{target_component_dir}' does not exist in the downloaded repo.")
|
|
# We might still want to create it if it's a new component, but usually it should exist for replacement.
|
|
# Proceeding assuming the user knows what they are doing.
|
|
os.makedirs(target_component_dir, exist_ok=True)
|
|
|
|
# Find weight files in checkpoint_dir
|
|
source_weights = glob.glob(os.path.join(args.checkpoint_dir, "*.safetensors")) + \
|
|
glob.glob(os.path.join(args.checkpoint_dir, "*.bin"))
|
|
|
|
if not source_weights:
|
|
print(f"Error: No .safetensors or .bin weights found in '{args.checkpoint_dir}'.")
|
|
return
|
|
|
|
print(f"Found weights in checkpoint dir: {[os.path.basename(x) for x in source_weights]}")
|
|
|
|
# Remove existing weights in target component dir to avoid mixing shards/files
|
|
existing_weights = glob.glob(os.path.join(target_component_dir, "*.safetensors")) + \
|
|
glob.glob(os.path.join(target_component_dir, "*.bin")) + \
|
|
glob.glob(os.path.join(target_component_dir, "*.index.json")) # Remove index if present
|
|
|
|
for ew in existing_weights:
|
|
print(f"Removing existing weight file: {ew}")
|
|
os.remove(ew)
|
|
|
|
# Copy new weights
|
|
for src in source_weights:
|
|
filename = os.path.basename(src)
|
|
dst = os.path.join(target_component_dir, filename)
|
|
|
|
# If there is only one file and it doesn't match the standard name, maybe we should rename it?
|
|
if len(source_weights) == 1 and filename != args.weight_file_name:
|
|
print(f"Renaming '{filename}' to '{args.weight_file_name}'...")
|
|
dst = os.path.join(target_component_dir, args.weight_file_name)
|
|
|
|
print(f"Copying '{src}' to '{dst}'...")
|
|
shutil.copy2(src, dst)
|
|
|
|
print("Replacement complete.")
|
|
print(f"New repo ready at: {args.local_dir}")
|
|
|
|
if args.push_to_hub:
|
|
target_repo_id = args.upload_repo_id if args.upload_repo_id else args.repo_id
|
|
print(f"Uploading to Hugging Face Hub: {target_repo_id}...")
|
|
|
|
try:
|
|
# Ensure the repo exists
|
|
create_repo(target_repo_id, exist_ok=True)
|
|
|
|
upload_folder(
|
|
repo_id=target_repo_id,
|
|
folder_path=args.local_dir,
|
|
commit_message=args.hub_commit_message
|
|
)
|
|
print("Upload complete!")
|
|
except Exception as e:
|
|
print(f"Error uploading to Hub: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|