import argparse import os.path as osp import itertools from omegaconf import OmegaConf from paintmind.engine.util import instantiate_from_config from paintmind.utils.device_utils import configure_compute_backend def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser("Test a model") # Model and testing configuration parser.add_argument('--model', type=str, nargs='+', default=[None], help="Path to model directory") parser.add_argument('--step', type=int, nargs='+', default=[250000], help="Step number to test") parser.add_argument('--cfg', type=str, default=None, help="Path to config file") parser.add_argument('--dataset', type=str, default='imagenet', help="Dataset to use") # Legacy parameter (preserved for backward compatibility) parser.add_argument('--cfg_value', type=float, nargs='+', default=[None], help='Legacy parameter for GPT classifier-free guidance scale') parser.add_argument('--ae_cfg', type=float, nargs='+', default=[None], help="Autoencoder classifier-free guidance scale") parser.add_argument('--diff_cfg', type=float, nargs='+', default=[None], help="Diffusion classifier-free guidance scale") parser.add_argument('--cfg_schedule', type=str, nargs='+', default=[None], help="CFG schedule type (e.g., constant, linear)") parser.add_argument('--diff_cfg_schedule', type=str, nargs='+', default=[None], help="Diffusion CFG schedule type (e.g., constant, inv_linear)") parser.add_argument('--test_num_slots', type=int, nargs='+', default=[None], help="Number of slots to use for inference") parser.add_argument('--temperature', type=float, nargs='+', default=[None], help="Temperature for sampling") return parser.parse_args() def load_config(model_path, cfg_path=None): """Load configuration from file or model directory.""" if cfg_path is not None and osp.exists(cfg_path): config_path = cfg_path elif model_path and osp.exists(osp.join(model_path, 'config.yaml')): config_path = osp.join(model_path, 'config.yaml') else: raise ValueError(f"No config file found at {model_path} or {cfg_path}") return OmegaConf.load(config_path) def setup_checkpoint_path(model_path, step, config): """Set up the checkpoint path based on model and step.""" if model_path: ckpt_path = osp.join(model_path, 'models', f'step{step}') if not osp.exists(ckpt_path): print(f"Skipping non-existent checkpoint: {ckpt_path}") return None if hasattr(config.trainer.params, 'model'): config.trainer.params.model.params.ckpt_path = ckpt_path else: config.trainer.params.gpt_model.params.ckpt_path = ckpt_path else: result_folder = config.trainer.params.result_folder ckpt_path = osp.join(result_folder, 'models', f'step{step}') if hasattr(config.trainer.params, 'model'): config.trainer.params.model.params.ckpt_path = ckpt_path else: config.trainer.params.gpt_model.params.ckpt_path = ckpt_path return ckpt_path def setup_test_config(config, use_coco=False): """Set up common test configuration parameters.""" config.trainer.params.test_dataset = config.trainer.params.dataset if not use_coco: config.trainer.params.test_dataset.params.split = 'val' else: config.trainer.params.test_dataset.target = 'paintmind.utils.datasets.COCO' config.trainer.params.test_dataset.params.root = './dataset/coco' config.trainer.params.test_dataset.params.split = 'val2017' config.trainer.params.test_only = True config.trainer.params.compile = False config.trainer.params.eval_fid = True config.trainer.params.fid_stats = 'fid_stats/adm_in256_stats.npz' if hasattr(config.trainer.params, 'model'): config.trainer.params.model.params.num_sampling_steps = '250' else: config.trainer.params.ae_model.params.num_sampling_steps = '250' def apply_cfg_params(config, param_dict): """Apply CFG-related parameters to the config.""" # Apply each parameter if it's not None if param_dict.get('cfg_value') is not None: config.trainer.params.cfg = param_dict['cfg_value'] print(f"Setting cfg to {param_dict['cfg_value']}") if param_dict.get('ae_cfg') is not None: config.trainer.params.ae_cfg = param_dict['ae_cfg'] print(f"Setting ae_cfg to {param_dict['ae_cfg']}") if param_dict.get('diff_cfg') is not None: config.trainer.params.diff_cfg = param_dict['diff_cfg'] print(f"Setting diff_cfg to {param_dict['diff_cfg']}") if param_dict.get('cfg_schedule') is not None: config.trainer.params.cfg_schedule = param_dict['cfg_schedule'] print(f"Setting cfg_schedule to {param_dict['cfg_schedule']}") if param_dict.get('diff_cfg_schedule') is not None: config.trainer.params.diff_cfg_schedule = param_dict['diff_cfg_schedule'] print(f"Setting diff_cfg_schedule to {param_dict['diff_cfg_schedule']}") if param_dict.get('test_num_slots') is not None: config.trainer.params.test_num_slots = param_dict['test_num_slots'] print(f"Setting test_num_slots to {param_dict['test_num_slots']}") if param_dict.get('temperature') is not None: config.trainer.params.temperature = param_dict['temperature'] print(f"Setting temperature to {param_dict['temperature']}") def run_test(config): """Instantiate trainer and run test.""" trainer = instantiate_from_config(config.trainer) trainer.train() def generate_param_combinations(args): """Generate all combinations of parameters from the provided arguments.""" # Create parameter grid for all combinations param_grid = { 'cfg_value': [None] if args.cfg_value == [None] else args.cfg_value, 'ae_cfg': [None] if args.ae_cfg == [None] else args.ae_cfg, 'diff_cfg': [None] if args.diff_cfg == [None] else args.diff_cfg, 'cfg_schedule': [None] if args.cfg_schedule == [None] else args.cfg_schedule, 'diff_cfg_schedule': [None] if args.diff_cfg_schedule == [None] else args.diff_cfg_schedule, 'test_num_slots': [None] if args.test_num_slots == [None] else args.test_num_slots, 'temperature': [None] if args.temperature == [None] else args.temperature } # Get all parameter names that have non-None values active_params = [k for k, v in param_grid.items() if v != [None]] if not active_params: # If no parameters are specified, yield a dict with all None values yield {k: None for k in param_grid.keys()} return # Generate all combinations of active parameters active_values = [param_grid[k] for k in active_params] for combination in itertools.product(*active_values): param_dict = {k: None for k in param_grid.keys()} # Start with all None for i, param_name in enumerate(active_params): param_dict[param_name] = combination[i] yield param_dict def test(args): """Main test function that processes arguments and runs tests.""" # Iterate through all model and step combinations for model in args.model: for step in args.step: print(f"Testing model: {model} at step: {step}") # Load configuration config = load_config(model, args.cfg) # Setup checkpoint path ckpt_path = setup_checkpoint_path(model, step, config) if ckpt_path is None: continue use_coco = args.dataset == 'coco' or args.dataset == 'COCO' # Setup test configuration setup_test_config(config, use_coco) # Generate and apply all parameter combinations for param_dict in generate_param_combinations(args): # Create a copy of the config for each parameter combination current_config = OmegaConf.create(OmegaConf.to_container(config, resolve=True)) # Print parameter combination param_str = ", ".join([f"{k}={v}" for k, v in param_dict.items() if v is not None]) print(f"Testing with parameters: {param_str}") # Apply parameters and run test apply_cfg_params(current_config, param_dict) run_test(current_config) def main(): """Main entry point for the script.""" args = parse_args() configure_compute_backend() test(args) if __name__ == "__main__": main()