(original_config, image_size: int)
| 369 | |
| 370 | |
| 371 | def create_vae_config(original_config, image_size: int): |
| 372 | vae_params = original_config["autoencoder"]["params"]["ddconfig"] |
| 373 | _ = original_config["autoencoder"]["params"]["embed_dim"] |
| 374 | |
| 375 | block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] |
| 376 | down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) |
| 377 | up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) |
| 378 | |
| 379 | config = { |
| 380 | "sample_size": image_size, |
| 381 | "in_channels": vae_params["in_channels"], |
| 382 | "out_channels": vae_params["out_ch"], |
| 383 | "down_block_types": tuple(down_block_types), |
| 384 | "up_block_types": tuple(up_block_types), |
| 385 | "block_out_channels": tuple(block_out_channels), |
| 386 | "latent_channels": vae_params["z_channels"], |
| 387 | "layers_per_block": vae_params["num_res_blocks"], |
| 388 | } |
| 389 | |
| 390 | return config |
| 391 | |
| 392 | |
| 393 | def create_unet_config(original_config, image_size: int, attention_type): |
no outgoing calls
no test coverage detected
searching dependent graphs…