| 111 | |
| 112 | |
| 113 | def log_validation( |
| 114 | base_model_path, |
| 115 | transformer, |
| 116 | vae, |
| 117 | text_encoder, |
| 118 | accelerator, |
| 119 | val_dataloader, |
| 120 | validation_store_path, |
| 121 | num_inference_steps, |
| 122 | ): |
| 123 | |
| 124 | |
| 125 | ################################################# The following should be the same as test_code section ############################################## |
| 126 | |
| 127 | # Create pipeline and run inference |
| 128 | pipe = CogVideoXImageToVideoPipeline.from_pretrained( |
| 129 | base_model_path, |
| 130 | text_encoder = text_encoder, |
| 131 | transformer = accelerator.unwrap_model(transformer), |
| 132 | vae = accelerator.unwrap_model(vae), |
| 133 | torch_dtype = torch.float16, |
| 134 | ) |
| 135 | pipe = pipe.to(accelerator.device) |
| 136 | # pipe.set_progress_bar_config(disable=True) |
| 137 | |
| 138 | |
| 139 | |
| 140 | # Iterate the validation dataset |
| 141 | for idx, batch in enumerate(val_dataloader): |
| 142 | |
| 143 | if idx != accelerator.process_index: # We want each process to be different index to do inference |
| 144 | continue |
| 145 | print("This Process Idx is", accelerator.process_index) |
| 146 | |
| 147 | |
| 148 | # Prepare the store folder here such that only one process come here each time |
| 149 | store_folder_path = os.path.join(validation_store_path, "Process"+str(idx)) |
| 150 | if os.path.exists(store_folder_path): # Should not have the path exists (unless there is a conflict) |
| 151 | os.system("rm -rf ", store_folder_path) |
| 152 | os.makedirs(store_folder_path, exist_ok=True) |
| 153 | |
| 154 | |
| 155 | # Fetch and Only the first batch size |
| 156 | first_frame_np = np.asarray(batch["first_frame_np"][0]) |
| 157 | text_prompt = batch["text_prompt"][0] |
| 158 | traj_tensor = batch["traj_tensor"][0] |
| 159 | traj_imgs_np = np.asarray(batch["traj_imgs_np"][0]) |
| 160 | gt_video_path = batch["gt_video_path"][0] |
| 161 | merge_frames = batch["merge_frames"][0] |
| 162 | |
| 163 | |
| 164 | # Fetch Generation Resulotion |
| 165 | gen_height, gen_width, _ = first_frame_np.shape |
| 166 | |
| 167 | |
| 168 | # Save the GT video |
| 169 | shutil.copyfile(gt_video_path, os.path.join(store_folder_path, "gt_video.mp4")) |
| 170 | |