()
| 67 | |
| 68 | |
| 69 | def main(): |
| 70 | parser = argparse.ArgumentParser() |
| 71 | parser.add_argument('--config', type=str, required=True) |
| 72 | parser.add_argument('--demo', action='store_true', default=False) |
| 73 | args = parser.parse_args() |
| 74 | |
| 75 | accelerator = Accelerator() |
| 76 | train_config = load_config(args.config) |
| 77 | |
| 78 | # Setup logger - log to sample output directory |
| 79 | exp_name = train_config['train']['exp_name'] |
| 80 | output_dir = train_config['train']['output_dir'] |
| 81 | log_dir = os.path.join(output_dir, exp_name, 'logs') |
| 82 | logger = setup_logger(log_dir, accelerator.process_index) |
| 83 | |
| 84 | # Log sampling config |
| 85 | if accelerator.process_index == 0: |
| 86 | sample_cfg = train_config.get('sample', {}) |
| 87 | logger.info(f"Sampling config: " |
| 88 | f"method={sample_cfg.get('sampling_method')}, " |
| 89 | f"steps={sample_cfg.get('num_sampling_steps')}, " |
| 90 | f"shift={sample_cfg.get('timestep_shift')}, " |
| 91 | f"cfg={sample_cfg.get('cfg_scale')}, " |
| 92 | f"fid_num={sample_cfg.get('fid_num')}, " |
| 93 | f"seed={train_config.get('train', {}).get('global_seed')}" |
| 94 | ) |
| 95 | |
| 96 | train_config['vae']['model_name'] = 'vtp' |
| 97 | |
| 98 | # Load HuggingFace VTP model config |
| 99 | hf_model_path = train_config['vae'].get('hf_model_path', None) |
| 100 | if hf_model_path is None: |
| 101 | raise ValueError("vae.hf_model_path must be specified") |
| 102 | |
| 103 | from vtp.models.vtp_hf import VTPConfig |
| 104 | hf_config = VTPConfig.from_pretrained(hf_model_path) |
| 105 | patch_size = hf_config.vision_patch_size |
| 106 | in_chans = hf_config.vision_feature_bottleneck |
| 107 | train_config['vae']['downsample_ratio'] = patch_size |
| 108 | if accelerator.process_index == 0: |
| 109 | logger.info(f'Using VTP model: {hf_model_path}') |
| 110 | |
| 111 | # Get checkpoint path from config |
| 112 | ckpt_path = train_config.get('ckpt_path') |
| 113 | if ckpt_path is None: |
| 114 | raise ValueError("ckpt_path must be specified in config") |
| 115 | |
| 116 | if accelerator.process_index == 0: |
| 117 | logger.info(f'Using ckpt: {ckpt_path}') |
| 118 | |
| 119 | latent_size = train_config['data']['image_size'] // train_config['vae']['downsample_ratio'] |
| 120 | |
| 121 | model = LightningDiT_models[train_config['model']['model_type']]( |
| 122 | input_size=latent_size, |
| 123 | num_classes=train_config['data']['num_classes'], |
| 124 | use_qknorm=train_config['model']['use_qknorm'], |
| 125 | use_swiglu=train_config['model'].get('use_swiglu', False), |
| 126 | use_rope=train_config['model'].get('use_rope', False), |
no test coverage detected