(local_dir,
hf_model_path,
compiled_dir,
img_height=1024,
img_width=1024,
compression_factor=8,
model_dir="vae",
torch_dtype=torch.float32,
bs=1,
exhaustive_tune=False,
fp16=False,
bf16=True)
| 494 | |
| 495 | |
| 496 | def get_vae_model(local_dir, |
| 497 | hf_model_path, |
| 498 | compiled_dir, |
| 499 | img_height=1024, |
| 500 | img_width=1024, |
| 501 | compression_factor=8, |
| 502 | model_dir="vae", |
| 503 | torch_dtype=torch.float32, |
| 504 | bs=1, |
| 505 | exhaustive_tune=False, |
| 506 | fp16=False, |
| 507 | bf16=True): |
| 508 | |
| 509 | vae_local_dir = get_local_path(local_dir, model_dir) |
| 510 | onnx_file = "model.onnx" |
| 511 | onnx_path = os.path.join(vae_local_dir, onnx_file) |
| 512 | latent_h, latent_w = img_height // compression_factor, img_width // compression_factor |
| 513 | |
| 514 | def get_compiled_file_name(): |
| 515 | name = f"model_b{bs}" |
| 516 | name += f"_h{latent_h}_w{latent_w}" |
| 517 | if fp16: |
| 518 | name += "_fp16" |
| 519 | elif bf16: |
| 520 | name += "_bf16" |
| 521 | if exhaustive_tune: name += "_exh" |
| 522 | return name + ".mxr" |
| 523 | |
| 524 | vae_compiled_dir = get_local_path(compiled_dir, model_dir) |
| 525 | mxr_file = get_compiled_file_name() |
| 526 | mxr_path = os.path.join(vae_compiled_dir, mxr_file) |
| 527 | |
| 528 | config = AutoencoderKL.load_config(hf_model_path, subfolder=model_dir) |
| 529 | |
| 530 | if os.path.isfile(mxr_path): |
| 531 | print(f"found compiled model.. loading VAE decoder from {mxr_path}") |
| 532 | model = MGXModel(mxr_path, config=config) |
| 533 | return model |
| 534 | |
| 535 | sample_inputs = (torch.randn(bs, |
| 536 | config['latent_channels'], |
| 537 | latent_h, |
| 538 | latent_w, |
| 539 | dtype=torch_dtype), ) |
| 540 | input_names = ["latent"] |
| 541 | if not os.path.isfile(onnx_path): |
| 542 | model = AutoencoderKL.from_pretrained(hf_model_path, |
| 543 | subfolder=model_dir, |
| 544 | torch_dtype=torch_dtype) |
| 545 | model.forward = model.decode |
| 546 | |
| 547 | output_names = ["images"] |
| 548 | dynamic_axes = { |
| 549 | 'latent': { |
| 550 | 0: 'B', |
| 551 | 2: 'H', |
| 552 | 3: 'W' |
| 553 | }, |
no test coverage detected