()
| 54 | |
| 55 | |
| 56 | def main(): |
| 57 | args = parse_args() |
| 58 | config = OmegaConf.load(args.config) |
| 59 | |
| 60 | if args.seed is not None: |
| 61 | seed_everything(args.seed) |
| 62 | |
| 63 | if config.weight_dtype == "fp16": |
| 64 | weight_dtype = torch.float16 |
| 65 | else: |
| 66 | weight_dtype = torch.float32 |
| 67 | |
| 68 | vae = AutoencoderKL.from_pretrained(config.pretrained_vae_path, ).to("cuda", dtype=weight_dtype) |
| 69 | inference_config_path = config.inference_config |
| 70 | infer_config = OmegaConf.load(inference_config_path) |
| 71 | denoising_unet = UNet3DConditionModel.from_pretrained_2d(config.pretrained_base_model_path,config.motion_module_path, subfolder="unet",unet_additional_kwargs=infer_config.unet_additional_kwargs, ).to(dtype=weight_dtype, device="cuda") |
| 72 | estimator = Estimator.from_pretrained_2d(config.pretrained_base_model_path, |
| 73 | config.motion_module_path, subfolder="unet", |
| 74 | unet_additional_kwargs=infer_config.unet_additional_kwargs, ).to(dtype=weight_dtype, device="cuda") |
| 75 | |
| 76 | pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to(dtype=weight_dtype, device="cuda") |
| 77 | image_enc = CLIPVisionModelWithProjection.from_pretrained(config.image_encoder_path).to(dtype=weight_dtype, |
| 78 | device="cuda") |
| 79 | |
| 80 | sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs) |
| 81 | scheduler = DDIMScheduler(**sched_kwargs) |
| 82 | ref_scheduler = DDIMScheduler(**sched_kwargs) |
| 83 | |
| 84 | vae.requires_grad_(False) |
| 85 | image_enc.requires_grad_(False) |
| 86 | denoising_unet.requires_grad_(False) |
| 87 | pose_guider.requires_grad_(False) |
| 88 | |
| 89 | vae.eval() |
| 90 | image_enc.eval() |
| 91 | denoising_unet.eval() |
| 92 | pose_guider.eval() |
| 93 | |
| 94 | generator = torch.manual_seed(args.seed) |
| 95 | width, height = args.W, args.H |
| 96 | # load pretrained weights |
| 97 | denoising_unet.load_state_dict( |
| 98 | torch.load(config.denoising_unet_path, map_location="cpu"), |
| 99 | strict=False, |
| 100 | ) |
| 101 | pose_guider.load_state_dict( |
| 102 | torch.load(config.pose_guider_path, map_location="cpu"), |
| 103 | ) |
| 104 | |
| 105 | if args.enable_xformers_memory_efficient_attention: |
| 106 | if is_xformers_available(): |
| 107 | denoising_unet.enable_xformers_memory_efficient_attention() |
| 108 | else: |
| 109 | raise ValueError( |
| 110 | "xformers is not available. Make sure it is installed correctly" |
| 111 | ) |
| 112 | if args.gradient_checkpointing: |
| 113 | denoising_unet.enable_gradient_checkpointing() |
no test coverage detected