| 163 | |
| 164 | |
| 165 | def log_validation( |
| 166 | base_model_path, |
| 167 | transformer, |
| 168 | vae, |
| 169 | text_encoder, |
| 170 | accelerator, |
| 171 | val_dataloader, |
| 172 | validation_store_path, |
| 173 | num_inference_steps, |
| 174 | ): |
| 175 | |
| 176 | |
| 177 | ################################################# The following should be the same as test_code section ############################################## |
| 178 | |
| 179 | # Create pipeline and run inference |
| 180 | pipe = WanImageToVideoPipeline.from_pretrained( |
| 181 | base_model_path, |
| 182 | text_encoder = text_encoder, |
| 183 | transformer = accelerator.unwrap_model(transformer), |
| 184 | vae = accelerator.unwrap_model(vae), |
| 185 | torch_dtype = torch.float16, |
| 186 | ) |
| 187 | pipe = pipe.to(accelerator.device) |
| 188 | # pipe.set_progress_bar_config(disable=True) |
| 189 | |
| 190 | |
| 191 | # Prepare the store folder |
| 192 | store_folder_path = os.path.join(validation_store_path, "Process"+str(accelerator.process_index)) |
| 193 | os.makedirs(store_folder_path, exist_ok=True) # Should not have the path exists (unless there is a conflict) |
| 194 | |
| 195 | |
| 196 | |
| 197 | # Iterate the validation dataset |
| 198 | for idx, batch in enumerate(val_dataloader): |
| 199 | |
| 200 | if idx != accelerator.process_index: # We want each process to be different index to do inference |
| 201 | continue |
| 202 | |
| 203 | print("This Process Idx is", accelerator.process_index) |
| 204 | # print("Node Rank is ", int(os.environ.get("NODE_RANK", 0)) ) |
| 205 | |
| 206 | # Fetch and Only the first batch size |
| 207 | first_frame_np = np.asarray(batch["first_frame_np"][0]) |
| 208 | text_prompt = batch["text_prompt"][0] |
| 209 | traj_tensor = batch["traj_tensor"][0] |
| 210 | traj_imgs_np = np.asarray(batch["traj_imgs_np"][0]) |
| 211 | gt_video_path = batch["gt_video_path"][0] |
| 212 | merge_frames = batch["merge_frames"][0] |
| 213 | |
| 214 | |
| 215 | # Chcek |
| 216 | gen_height, gen_width, _ = first_frame_np.shape |
| 217 | |
| 218 | |
| 219 | # Save the GT video |
| 220 | shutil.copyfile(gt_video_path, os.path.join(store_folder_path, "gt_video.mp4")) |
| 221 | |
| 222 | # Convert first frame np to image form |