()
| 29 | import kiui |
| 30 | |
| 31 | def main(): |
| 32 | batch_size = 1 |
| 33 | gradient_accumulation_steps = 1 |
| 34 | |
| 35 | opt = Options( |
| 36 | input_size=256, |
| 37 | up_channels=(1024, 1024, 512, 256, 128), |
| 38 | up_attention=(True, True, True, False, False), |
| 39 | splat_size=128, |
| 40 | output_size=512, |
| 41 | batch_size=batch_size, |
| 42 | data_mode='imagedream', |
| 43 | num_views=12, |
| 44 | num_epochs=200, |
| 45 | workspace='mvr_pretrain_512_vae_xt_continue', |
| 46 | resume='./checkpoints/model_1.safetensors', |
| 47 | gradient_accumulation_steps=gradient_accumulation_steps, |
| 48 | mixed_precision='bf16', |
| 49 | lr=1e-4, |
| 50 | cam_radius=1.5, |
| 51 | tracker_project_name='mvr_pretrain_512_bf16', |
| 52 | lambda_lpips=1.0, |
| 53 | lambda_distortion=100.0, |
| 54 | lambda_normal=0.0, |
| 55 | lambda_depth=0.0 |
| 56 | ) |
| 57 | |
| 58 | accelerator = Accelerator( |
| 59 | mixed_precision=opt.mixed_precision, |
| 60 | gradient_accumulation_steps=opt.gradient_accumulation_steps, |
| 61 | log_with='wandb' |
| 62 | ) |
| 63 | |
| 64 | logging.basicConfig( |
| 65 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 66 | datefmt="%m/%d/%Y %H:%M:%S", |
| 67 | level=logging.INFO, |
| 68 | ) |
| 69 | logger.info(accelerator.state, main_process_only=False) |
| 70 | |
| 71 | model = LGM_timeimagecond_noise_gof(opt) |
| 72 | model.train() |
| 73 | |
| 74 | from diffusers import AutoencoderKL |
| 75 | vae = AutoencoderKL.from_pretrained("ashawkey/imagedream-ipmv-diffusers", subfolder="vae", revision=None) |
| 76 | vae.eval() |
| 77 | vae.requires_grad_(False) |
| 78 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| 79 | vae.to(device) |
| 80 | |
| 81 | noise_scheduler = DDPMScheduler.from_pretrained("ashawkey/imagedream-ipmv-diffusers", subfolder="scheduler", revision=None) |
| 82 | |
| 83 | def print_model_info(model): |
| 84 | print("="*20) |
| 85 | print("model name: ", type(model).__name__) |
| 86 | print("learnable parameters(M): ", sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6) |
| 87 | print("non-learnable parameters(M): ", sum(p.numel() for p in model.parameters() if not p.requires_grad) / 1e6) |
| 88 | print("total parameters(M): ", sum(p.numel() for p in model.parameters()) / 1e6) |
no test coverage detected