()
| 276 | |
| 277 | |
| 278 | def main(): |
| 279 | args = parse_args() |
| 280 | logging_dir = Path(args.output_dir, args.logging_dir) |
| 281 | |
| 282 | accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) |
| 283 | |
| 284 | accelerator = Accelerator( |
| 285 | mixed_precision=args.mixed_precision, |
| 286 | log_with=args.report_to, |
| 287 | project_config=accelerator_project_config, |
| 288 | ) |
| 289 | |
| 290 | if accelerator.is_main_process: |
| 291 | if args.output_dir is not None: |
| 292 | os.makedirs(args.output_dir, exist_ok=True) |
| 293 | |
| 294 | # Load scheduler, tokenizer and models. |
| 295 | noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") |
| 296 | tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer") |
| 297 | text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder") |
| 298 | vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae") |
| 299 | unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet") |
| 300 | image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder_path) |
| 301 | # freeze parameters of models to save more memory |
| 302 | unet.requires_grad_(False) |
| 303 | vae.requires_grad_(False) |
| 304 | text_encoder.requires_grad_(False) |
| 305 | image_encoder.requires_grad_(False) |
| 306 | |
| 307 | # ip-adapter-plus |
| 308 | image_proj_model = Resampler( |
| 309 | dim=unet.config.cross_attention_dim, |
| 310 | depth=4, |
| 311 | dim_head=64, |
| 312 | heads=12, |
| 313 | num_queries=args.num_tokens, |
| 314 | embedding_dim=image_encoder.config.hidden_size, |
| 315 | output_dim=unet.config.cross_attention_dim, |
| 316 | ff_mult=4, |
| 317 | ) |
| 318 | # init adapter modules |
| 319 | attn_procs = {} |
| 320 | unet_sd = unet.state_dict() |
| 321 | for name in unet.attn_processors.keys(): |
| 322 | cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim |
| 323 | if name.startswith("mid_block"): |
| 324 | hidden_size = unet.config.block_out_channels[-1] |
| 325 | elif name.startswith("up_blocks"): |
| 326 | block_id = int(name[len("up_blocks.")]) |
| 327 | hidden_size = list(reversed(unet.config.block_out_channels))[block_id] |
| 328 | elif name.startswith("down_blocks"): |
| 329 | block_id = int(name[len("down_blocks.")]) |
| 330 | hidden_size = unet.config.block_out_channels[block_id] |
| 331 | if cross_attention_dim is None: |
| 332 | attn_procs[name] = AttnProcessor() |
| 333 | else: |
| 334 | layer_name = name.split(".processor")[0] |
| 335 | weights = { |
no test coverage detected
searching dependent graphs…