()
| 23 | pass |
| 24 | |
| 25 | def main(): |
| 26 | parser = argparse.ArgumentParser(description="Merge Stable Diffusion weights with ControlNet training weights") |
| 27 | parser.add_argument("checkpoint_path", type=str, help="Path to the trained model checkpoint, e.g.: ./checkpoints/checkpoints_DIOR_train/model-step=10000.ckpt") |
| 28 | parser.add_argument("--sd15_path", type=str, default="./models/control_sd15_ini.ckpt", help="Path to SD15 initialization model") |
| 29 | parser.add_argument("--output_dir", type=str, help="Output directory, defaults to a 'merged' folder in the same directory as the checkpoint") |
| 30 | parser.add_argument("--use_direct_load", action="store_true", help="Directly load model files instead of using zero_to_fp32.py") |
| 31 | args = parser.parse_args() |
| 32 | |
| 33 | # Set paths |
| 34 | checkpoint_path = args.checkpoint_path |
| 35 | sd15_path = args.sd15_path |
| 36 | |
| 37 | # If no output directory is specified, create a 'merged' folder in the same directory as the checkpoint |
| 38 | if args.output_dir: |
| 39 | output_dir = args.output_dir |
| 40 | else: |
| 41 | checkpoint_dir = os.path.dirname(checkpoint_path) |
| 42 | output_dir = os.path.join(checkpoint_dir, "merged") |
| 43 | |
| 44 | # Create necessary directories |
| 45 | os.makedirs(output_dir, exist_ok=True) |
| 46 | |
| 47 | # Output file path |
| 48 | output_path = os.path.join(output_dir, "merged_pytorch_model.pth") |
| 49 | |
| 50 | # Find checkpoint subdirectory |
| 51 | checkpoint_subdir = os.path.join(checkpoint_path, "checkpoint") |
| 52 | if not os.path.isdir(checkpoint_subdir): |
| 53 | print(f"Warning: checkpoint subdirectory not found in {checkpoint_path}") |
| 54 | # Try using checkpoint_path directly |
| 55 | checkpoint_subdir = checkpoint_path |
| 56 | |
| 57 | # Check if latest file exists |
| 58 | latest_file = os.path.join(checkpoint_subdir, "latest") |
| 59 | |
| 60 | # Two modes: using zero_to_fp32.py or direct loading |
| 61 | if args.use_direct_load or not os.path.exists(latest_file): |
| 62 | print(f"Step 1: Preparing to load model files directly...") |
| 63 | |
| 64 | # Find all mp_rank files |
| 65 | mp_rank_files = glob.glob(os.path.join(checkpoint_subdir, "mp_rank_*_model_states.pt")) |
| 66 | |
| 67 | if not mp_rank_files: |
| 68 | print(f"Warning: No mp_rank files found in {checkpoint_subdir}") |
| 69 | raise FileNotFoundError(f"Could not find model files in {checkpoint_subdir}") |
| 70 | |
| 71 | print(f"Found model file: {mp_rank_files[0]}") |
| 72 | |
| 73 | # Load model file directly |
| 74 | print(f"Loading model file: {mp_rank_files[0]}") |
| 75 | model_state_dict = torch.load(mp_rank_files[0], map_location="cpu") |
| 76 | |
| 77 | # Extract model weights |
| 78 | if "module" in model_state_dict: |
| 79 | sd15_with_control_state_dict = model_state_dict["module"] |
| 80 | else: |
| 81 | sd15_with_control_state_dict = model_state_dict |
| 82 | else: |
no test coverage detected