(args)
| 472 | weight_dict=weight_dict, |
| 473 | lora_weights=lora_weights, |
| 474 | alpha=alpha, |
| 475 | strength=strength, |
| 476 | ) |
| 477 | |
| 478 | |
| 479 | def convert_weights(args): |
| 480 | if os.path.isdir(args.source): |
| 481 | src_files = glob.glob(os.path.join(args.source, "*.safetensors"), recursive=True) |
| 482 | elif args.source.endswith((".pth", ".safetensors", "pt")): |
| 483 | src_files = [args.source] |
| 484 | else: |
| 485 | raise ValueError("Invalid input path") |
| 486 | |
| 487 | merged_weights = {} |
| 488 | logger.info(f"Processing source files: {src_files}") |
| 489 | |
| 490 | # Optimize loading for better memory usage |
| 491 | for file_path in tqdm(src_files, desc="Loading weights"): |
| 492 | logger.info(f"Loading weights from: {file_path}") |
| 493 | if file_path.endswith(".pt") or file_path.endswith(".pth"): |
| 494 | weights = torch.load(file_path, map_location=args.device, weights_only=True) |
| 495 | if args.model_type == "hunyuan_dit": |
| 496 | weights = weights["module"] |
| 497 | elif args.model_type == "self_forcing": |
| 498 | weights = weights["generator_ema"] |
| 499 | elif file_path.endswith(".safetensors"): |
| 500 | # Use lazy loading for safetensors to reduce memory usage |
| 501 | with safe_open(file_path, framework="pt", device=args.device) as f: |
| 502 | # Only load tensors when needed (lazy loading) |
| 503 | weights = {} |
| 504 | keys = f.keys() |
| 505 | |
| 506 | # For large files, show progress |
| 507 | if len(keys) > 100: |
| 508 | for k in tqdm(keys, desc=f"Loading {os.path.basename(file_path)}", leave=False): |
| 509 | weights[k] = f.get_tensor(k) |
| 510 | else: |
| 511 | weights = {k: f.get_tensor(k) for k in keys} |
| 512 | |
| 513 | duplicate_keys = set(weights.keys()) & set(merged_weights.keys()) |
| 514 | if duplicate_keys: |
| 515 | raise ValueError(f"Duplicate keys found: {duplicate_keys} in file {file_path}") |
| 516 | |
| 517 | # Update weights more efficiently |
| 518 | merged_weights.update(weights) |
| 519 | |
| 520 | # Clear weights dict to free memory |
| 521 | del weights |
| 522 | if len(src_files) > 1: |
| 523 | gc.collect() # Force garbage collection between files |
| 524 | |
| 525 | if args.direction is not None: |
| 526 | rules = get_key_mapping_rules(args.direction, args.model_type) |
| 527 | converted_weights = {} |
| 528 | logger.info("Converting keys...") |
| 529 | |
| 530 | # Pre-compile regex patterns for better performance |
| 531 | compiled_rules = [(re.compile(pattern), replacement) for pattern, replacement in rules] |
no test coverage detected