(args, device)
| 39 | |
| 40 | |
| 41 | def get_model(args, device): |
| 42 | if args.is_latent: |
| 43 | # Create model: |
| 44 | assert ( |
| 45 | args.data.image_size % 8 == 0 |
| 46 | ), "Image size must be divisible by 8 (for the VAE encoder)." |
| 47 | in_channels = 4 |
| 48 | input_size = args.data.image_size // 8 |
| 49 | else: |
| 50 | in_channels = 3 |
| 51 | input_size = args.data.image_size |
| 52 | logging.info(f"input_size: {input_size}, in_channels: {in_channels}") |
| 53 | |
| 54 | model = instantiate_from_config(args.model).to(device) |
| 55 | |
| 56 | return model, in_channels, input_size |
| 57 | |
| 58 | |
| 59 | def create_logger(rank, logging_dir=None): |
no test coverage detected