from diffsynth import ModelManager, WanVideoRelitlivePipeline from torchvision.transforms import v2 from einops import rearrange from PIL import Image from tqdm import tqdm from natsort import natsorted from typing import List import torch, os, imageio, argparse, torchvision, pdb, cv2, pathlib, glob, pyexr import torch.nn as nn import numpy as np def stitch_frames(video_frames, rows=2, cols=5): num_frames = len(video_frames[0]) assert all(len(frames) == num_frames for frames in video_frames), "All video frames must have the same number of frames." frame_width, frame_height = video_frames[0][0].size stitched_images = [] for frame_idx in range(num_frames): stitched_image = Image.new('RGB', (cols * frame_width, rows * frame_height)) for i in range(rows): for j in range(cols): video_idx = i * cols + j if video_idx < len(video_frames): img = video_frames[video_idx][frame_idx] x = j * frame_width y = i * frame_height stitched_image.paste(img, (x, y)) stitched_images.append(stitched_image) return stitched_images def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None): writer = imageio.get_writer(save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params) for frame in tqdm(frames, desc="Saving video"): frame = np.array(frame) writer.append_data(frame) writer.close() def save_frames(frames, save_path): os.makedirs(save_path, exist_ok=True) for i, frame in enumerate(tqdm(frames, desc="Saving images")): frame.save(os.path.join(save_path, f"{i}.png")) def rotate_panorama_around_horizontal_axis( panorama_img, pitch_angle=10 ): if isinstance(panorama_img, Image.Image): panorama = np.array(panorama_img.convert('RGB')) elif isinstance(panorama_img, np.ndarray): panorama = panorama_img.copy() if len(panorama.shape) != 3 or panorama.shape[2] != 3: raise ValueError("The NumPy array for the panoramic image must be in RGB format with dimensions (H, W, 3).") else: raise TypeError("The panoramic image must be a PIL.Image or a numpy.ndarray.") pan_h, pan_w = panorama.shape[:2] pitch_rad = np.radians(pitch_angle) lon = np.linspace(0, 2 * np.pi, pan_w) lat = np.linspace(-np.pi/2, np.pi/2, pan_h) lon_grid, lat_grid = np.meshgrid(lon, lat) x = np.cos(lat_grid) * np.sin(lon_grid) y = np.sin(lat_grid) z = np.cos(lat_grid) * np.cos(lon_grid) coords = np.stack([x, y, z], axis=-1) R_pitch = np.array([ [1, 0, 0], [0, np.cos(pitch_rad), -np.sin(pitch_rad)], [0, np.sin(pitch_rad), np.cos(pitch_rad)] ]) coords_rot = np.dot(coords, R_pitch.T) lon_new = np.arctan2(coords_rot[..., 0], coords_rot[..., 2]) lat_new = np.arcsin(np.clip(coords_rot[..., 1], -1, 1)) u_new = (lon_new / (2 * np.pi) + 0.5) * pan_w v_new = (lat_new / np.pi + 0.5) * pan_h u_new = u_new.astype(np.float32) v_new = v_new.astype(np.float32) new_panorama = cv2.remap( panorama, u_new, v_new, interpolation=cv2.INTER_CUBIC, borderMode=cv2.BORDER_WRAP ) new_panorama = np.hstack([new_panorama[:, pan_w//2:, ...], new_panorama[:, :pan_w//2, ...]]) return new_panorama class PBRVideo_img_Dataset(torch.utils.data.Dataset): def __init__(self, base_path, max_num_frames=81, frame_interval=1, num_frames=81, height=480, width=832, env_map_path=None, dataset_type='relit-live', \ use_ref_image=False, full_resolution=False, padding_resolution=False, drop_mr=False, args=None): self.dataset_type = dataset_type self.base_path = base_path self.args = args self.num_frames = num_frames p = pathlib.Path(base_path) if not p.is_dir(): raise NotADirectoryError(f"{base_path}' is not a valid dir") if self.dataset_type == "relit-live": self.path: List[pathlib.Path] = natsorted([item for item in p.iterdir() if item.is_dir()]) self.frame_interval = frame_interval if self.num_frames!=1 else 0 print(f'============= Load {len(self.path)}seqs from {base_path} =============') self.max_num_frames = max_num_frames self.num_frames = num_frames self.height = height self.width = width self.frame_process = v2.Compose([ v2.CenterCrop(size=(height, width)), v2.ToTensor(), v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) self.env_map_path = env_map_path self.use_ref_image = use_ref_image self.full_resolution = full_resolution self.padding_resolution = padding_resolution self.drop_mr = drop_mr def crop_and_resize(self, image, std_shape=None, env_reshape=False, align=None): if env_reshape: image = torchvision.transforms.functional.resize( image, (self.height, self.width), interpolation=torchvision.transforms.InterpolationMode.BILINEAR ) else: if std_shape is not None: height, width = std_shape else: width, height = image.size if self.padding_resolution: scale = min(self.width / width, self.height / height) else: scale = max(self.width / width, self.height / height) if align is not None: if align == 'width': scale = self.width / width elif align == 'height': scale = self.height / height image = torchvision.transforms.functional.resize( image, (round(height*scale), round(width*scale)), interpolation=torchvision.transforms.InterpolationMode.BILINEAR ) return image def load_frames_using_imageio(self, file_path, start_frame_id, interval, num_frames, frame_process, std_shape=None, env_reshape=False, return_shape=False): reader = imageio.get_reader(file_path) image_num = reader.count_frames() frames = [] for frame_id in range(num_frames): sample_list = list(range(image_num)) + list(range(image_num-2, 0, -1)) sample_id = sample_list[((start_frame_id + frame_id * interval) % len(sample_list))] frame = reader.get_data(sample_id) img_shape = frame.shape[:2] frame = Image.fromarray(frame) frame = self.crop_and_resize(frame, std_shape, env_reshape) frame = frame_process(frame) frames.append(frame) reader.close() frames = torch.stack(frames, dim=0) frames = rearrange(frames, "T C H W -> C T H W") if return_shape: return frames, img_shape return frames def load_frames_using_imageio_from_imgdir(self, file_path, start_frame_id, interval, num_frames, frame_process, divided_max=False, std_shape=None, env_reshape=False, return_shape=False): image_paths = glob.glob(os.path.join(file_path, '*.png')) + glob.glob(os.path.join(file_path, '*.jpg')) + glob.glob(os.path.join(file_path, '*.exr')) if len(image_paths) == 0: return None sorted_image_paths = natsorted(image_paths) image_num = len(sorted_image_paths) frames = [] for frame_id in range(num_frames): sample_list = list(range(image_num)) + list(range(image_num-2, 0, -1)) sample_id = sample_list[((start_frame_id + frame_id * interval) % len(sample_list))] if ".exr" in sorted_image_paths[sample_id]: frame = pyexr.open(sorted_image_paths[sample_id]).get() else: frame = imageio.imread(sorted_image_paths[sample_id]) if frame.shape[-1] == 4: frame = frame[:,:,:3] elif frame.shape[-1] == 1: frame = np.repeat(frame, 3, axis=-1) img_shape = frame.shape[:2] if divided_max: bg_mask = (frame > 1000.0) if len(frame[bg_mask]) > 0: frame_max = float(np.percentile(frame, 99)) frame = np.clip(frame / frame_max, 0, 1) else: frame = frame / frame.max() if frame.min() < 0: frame = frame * 0.5 + 0.5 if frame.dtype != 'uint8': frame = (frame * 255).astype(np.uint8) frame = Image.fromarray(frame) frame = self.crop_and_resize(frame, std_shape, env_reshape) frame = frame_process(frame) frames.append(frame) frames = torch.stack(frames, dim=0) frames = rearrange(frames, "T C H W -> C T H W") if return_shape: return frames, img_shape return frames def load_frames_using_imageio_from_imgpath(self, file_path, start_frame_id, interval, num_frames, frame_process, divided_max=False, std_shape=None, env_reshape=False, return_shape=False, align=None): image_paths = [file_path] if len(image_paths) == 0: return None sorted_image_paths = natsorted(image_paths) image_num = len(sorted_image_paths) frames = [] for frame_id in range(num_frames): sample_list = list(range(image_num)) + list(range(image_num-2, 0, -1)) sample_id = sample_list[((start_frame_id + frame_id * interval) % len(sample_list))] if ".exr" in sorted_image_paths[sample_id]: frame = pyexr.open(sorted_image_paths[sample_id]).get() else: frame = imageio.imread(sorted_image_paths[sample_id]) if frame.shape[-1] == 4: frame = frame[:,:,:3] elif frame.shape[-1] == 1: frame = np.repeat(frame, 3, axis=-1) img_shape = frame.shape[:2] if divided_max: bg_mask = (frame > 1000.0) if len(frame[bg_mask]) > 0: frame_max = float(np.percentile(frame, 99)) frame = np.clip(frame / frame_max, 0, 1) else: frame = frame / frame.max() if frame.min() < 0: frame = frame * 0.5 + 0.5 if frame.dtype != 'uint8': frame = (frame * 255).astype(np.uint8) frame = Image.fromarray(frame) frame = self.crop_and_resize(frame, std_shape, env_reshape, align=align) frame = frame_process(frame) frames.append(frame) frames = torch.stack(frames, dim=0) frames = rearrange(frames, "T C H W -> C T H W") if return_shape: return frames, img_shape return frames def __getitem__(self, index): data_id = index % len(self.path) dir_path = self.path[data_id] RGB_path = os.path.join(dir_path, "images_4") basecolor_path = os.path.join(dir_path, "Base Color") depth_path = os.path.join(dir_path, "depth") metallic_path = os.path.join(dir_path, "Metallic") normal_path = os.path.join(dir_path, "normal") roughness_path = os.path.join(dir_path, "Roughness") fixed_env = True if self.env_map_path is not None: ldr_path = os.path.join(self.env_map_path, "ldr_video_fix_first_frame.mp4") hdr_log_path = os.path.join(self.env_map_path, "hdr_log_video_fix_first_frame.mp4") env_dir_path = os.path.join(self.env_map_path, "env_dir_video_fix_first_frame.mp4") else: ldr_path = os.path.join(dir_path, "env", "ldr_video_fix_first_frame.mp4") hdr_log_path = os.path.join(dir_path, "env", "hdr_log_video_fix_first_frame.mp4") env_dir_path = os.path.join(dir_path, "env", "env_dir_video_fix_first_frame.mp4") fixed_env = False start_frame_id = 0 if self.full_resolution: source = self.load_frames_using_imageio_from_imgdir(RGB_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) basecolor= self.load_frames_using_imageio_from_imgdir(basecolor_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) depth= self.load_frames_using_imageio_from_imgdir(depth_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, divided_max=True, env_reshape=True) metallic= self.load_frames_using_imageio_from_imgdir(metallic_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) normal= self.load_frames_using_imageio_from_imgdir(normal_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) roughness = self.load_frames_using_imageio_from_imgdir(roughness_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) else: source, std_shape = self.load_frames_using_imageio_from_imgdir(RGB_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, return_shape=True) basecolor= self.load_frames_using_imageio_from_imgdir(basecolor_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, std_shape=std_shape) depth= self.load_frames_using_imageio_from_imgdir(depth_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, divided_max=True, std_shape=std_shape) metallic= self.load_frames_using_imageio_from_imgdir(metallic_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, std_shape=std_shape) normal= self.load_frames_using_imageio_from_imgdir(normal_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, std_shape=std_shape) roughness = self.load_frames_using_imageio_from_imgdir(roughness_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, std_shape=std_shape) if metallic is None: metallic = source * 0.0 if roughness is None: roughness = source * 0.0 ldr = self.load_frames_using_imageio(ldr_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) hdr_log = self.load_frames_using_imageio(hdr_log_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) env_dir = self.load_frames_using_imageio(env_dir_path, start_frame_id, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) if fixed_env: ldr = ldr[:,0:1].repeat(1, self.num_frames, 1, 1) hdr_log = hdr_log[:,0:1].repeat(1, self.num_frames, 1, 1) if self.drop_mr: metallic_weight = torch.tensor(0.0) roughness_weight = torch.tensor(0.0) else: metallic_weight = torch.tensor(1.0) roughness_weight = torch.tensor(1.0) env_self_weight = torch.tensor(1.0) env_cross_weight = torch.tensor(1.0) # print(f"haven't use prompt for {RGB_path}") prompt = "" if self.use_ref_image: image_paths = glob.glob(os.path.join(basecolor_path, '*.png')) + glob.glob(os.path.join(basecolor_path, '*.jpg')) + glob.glob(os.path.join(basecolor_path, '*.exr')) max_idx = (len(image_paths)-1) - (self.num_frames-1) * self.frame_interval if max_idx == 0: start_frame_id_ref = 0 else: start_frame_id_ref = start_frame_id ref_RGB_path = RGB_path if self.full_resolution: ref_source = self.load_frames_using_imageio_from_imgdir(ref_RGB_path, start_frame_id_ref, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) else: ref_source = self.load_frames_using_imageio_from_imgdir(ref_RGB_path, start_frame_id_ref, self.frame_interval, self.num_frames, self.frame_process, std_shape=std_shape) if self.args.use_fixed_frame_and_w_rotate_light or self.args.use_fixed_frame_and_h_rotate_light or self.num_frames == 1: ref_idx = 0 else: ref_idx = ref_source.shape[1] // 2 # Use the mid frame as the reference image ref_image = ref_source[:, ref_idx:ref_idx+1, :, :] if self.args.ref_image_path_with_idddx is not None: ref_image_path_with_idx = self.args.ref_image_path_with_idddx.replace("idddx", f"{index}") if self.full_resolution: ref_image = self.load_frames_using_imageio_from_imgpath(ref_image_path_with_idx, self.max_num_frames, start_frame_id_ref, self.frame_interval, self.num_frames, self.frame_process, env_reshape=True) else: ref_image = self.load_frames_using_imageio_from_imgpath(ref_image_path_with_idx, self.max_num_frames, start_frame_id_ref, self.frame_interval, self.num_frames, self.frame_process, std_shape=None, align='width') ref_image = ref_image[:, -1:, :, :] data = { "source_video": source, "basecolor": basecolor, "depth": depth, "metallic":metallic, "normal": normal, "roughness":roughness, 'metallic_weight': metallic_weight, 'roughness_weight': roughness_weight, 'env_self_weight': env_self_weight, 'env_cross_weight': env_cross_weight, "ldr": ldr, "hdr_log": hdr_log, "env_dir": env_dir, "ref_image": ref_image, "ref_video": ref_source, "path": str(dir_path), "prompt": prompt, } else: data = { "source_video": source, "basecolor": basecolor, "depth": depth, "metallic":metallic, "normal": normal, "roughness":roughness, 'metallic_weight': metallic_weight, 'roughness_weight': roughness_weight, 'env_self_weight': env_self_weight, 'env_cross_weight': env_cross_weight, "ldr": ldr, "hdr_log": hdr_log, "env_dir": env_dir, "path": str(dir_path), "prompt": prompt, } return data def __len__(self): return len(self.path) def parse_args(): parser = argparse.ArgumentParser(description="diffusion renderer Inference") parser.add_argument( "--dataset_path", type=str, default="./example_test_data", help="The path of the Dataset.", ) parser.add_argument( "--env_map_path", type=str, default=None, help="env_map_path.", ) parser.add_argument( "--use_ref_image", default=False, action="store_true", help="Whether to use reference image.", ) parser.add_argument( "--use_muti_ref_image", default=False, action="store_true", help="Whether to use multiple reference images.", ) parser.add_argument( "--ref_image_path_with_idddx", type=str, default=None, help="Path to the reference image with index.", ) parser.add_argument( "--full_resolution", default=False, action="store_true", help="Whether to use full_resolution.", ) parser.add_argument( "--padding_resolution", default=False, action="store_true", help="Whether to use padding_resolution.", ) parser.add_argument( "--dataset_type", type=str, default="relit-live", help="Dataset format to load. Use 'relit-live' for the default Relit-LiVE directory layout.", ) parser.add_argument( "--drop_mr", default=False, action="store_true", help="Ignore metallic and roughness inputs by setting their conditioning weights to zero.", ) parser.add_argument( "--use_rotate_light", default=False, action="store_true", help="Enable light rotation mode during inference.", ) parser.add_argument( "--use_fixed_frame_and_w_rotate_light", default=False, action="store_true", help="Repeat the first frame and rotate the environment map along the width axis across frames.", ) parser.add_argument( "--use_fixed_frame_and_h_rotate_light", default=False, action="store_true", help="Repeat the first frame and rotate the environment map along the height axis across frames.", ) parser.add_argument( "--h_rotate_light", type=int, default=0, help="Rotate the environment map vertically by this many degrees for every frame.", ) parser.add_argument( "--w_rotate_light", type=int, default=0, help="Rotate the environment map horizontally by this many pixels for every frame.", ) parser.add_argument( "--num_frames", type=int, default=81, help="Number of frames.", ) parser.add_argument( "--num_inference_steps", type=int, default=50, help="Number of denoising steps used during inference.", ) parser.add_argument( "--frame_interval", type=int, default=1, help="Sampling interval between frames read from the input sequence.", ) parser.add_argument( "--height", type=int, default=480, help="Image height.", ) parser.add_argument( "--width", type=int, default=832, help="Image width.", ) parser.add_argument( "--ckpt_path", type=str, default=None, help="Path to the fine-tuned checkpoint to load into the pipeline.", ) parser.add_argument( "--output_dir", type=str, default="./results", help="Path to save the results.", ) parser.add_argument( "--output_path", type=str, default=None, help="Optional explicit output file path. Supports .mp4 and .png.", ) parser.add_argument( "--dataloader_num_workers", type=int, default=1, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", ) parser.add_argument( "--cfg_scale", type=float, default=5.0, help="Classifier-free guidance scale.", ) parser.add_argument( "--wo_ref_weight", type=float, default=0.0, help="Weight applied to the branch without reference-image conditioning.", ) parser.add_argument( "--quality", type=int, default=5, help="Video encoding quality passed to imageio when saving mp4 outputs.", ) args = parser.parse_args() return args if __name__ == '__main__': args = parse_args() # Load Wan2.1 pre-trained models model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu") model_manager.load_models([ "models/Wan-AI/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors", "models/Wan-AI/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", "models/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth", ]) pipe = WanVideoRelitlivePipeline.from_model_manager(model_manager, device="cuda") # Initialize additional modules introduced in relit live dim=pipe.dit.blocks[0].self_attn.q.weight.shape[0] for block in pipe.dit.blocks: block.env_encoder = nn.Sequential( nn.Conv3d( 48, dim, kernel_size=(1,2,2), stride=(1,2,2)), nn.SiLU() ) block.env_encoder[0].weight.data.zero_() block.env_encoder[0].bias.data.zero_() block.projector = nn.Linear(dim, dim) block.projector.weight = nn.Parameter(torch.eye(dim)) block.projector.bias = nn.Parameter(torch.zeros(dim)) # if args.use_ref_image: dim = pipe.dit.patch_embedding.out_channels pipe.dit.ref_conv = nn.Conv2d(16, dim, kernel_size=(2,2), stride=(2,2)) # Load checkpoint state_dict = torch.load(args.ckpt_path, map_location="cpu") pipe.dit.load_state_dict(state_dict, strict=False) # True pipe.to("cuda") pipe.to(dtype=torch.bfloat16) output_dir = os.path.join(args.output_dir) if not os.path.exists(output_dir): os.makedirs(output_dir) dataset = PBRVideo_img_Dataset( base_path=args.dataset_path, max_num_frames=args.num_frames, frame_interval=args.frame_interval, num_frames=args.num_frames, height=args.height, width=args.width, env_map_path=args.env_map_path, dataset_type=args.dataset_type, use_ref_image=args.use_ref_image, full_resolution=args.full_resolution, padding_resolution=args.padding_resolution, drop_mr=args.drop_mr, args=args, ) dataloader = torch.utils.data.DataLoader( dataset, shuffle=False, batch_size=1, num_workers=args.dataloader_num_workers ) # Inference ldr_img, log_img, ref_img = None, None, None for batch_idx, batch in enumerate(dataloader): input_rgb = batch["source_video"] # b c t h w basecolor = batch["basecolor"] depth = batch["depth"] metallic = batch["metallic"] normal = batch["normal"] roughness = batch["roughness"] ldr = batch["ldr"] hdr_log = batch["hdr_log"] env_dir = batch["env_dir"] dir_path = batch["path"][0] prompt = batch["prompt"] print(f'processing {dir_path}, prompt: {prompt}.') if args.h_rotate_light != 0: for frame_idx in range(args.num_frames): ldr_now = np.array(ldr[0,:,frame_idx,...].permute(1, 2, 0)) # c h w -> h w c hdr_log_now = np.array(hdr_log[0,:,frame_idx,...].permute(1, 2, 0)) ldr_now = rotate_panorama_around_horizontal_axis(ldr_now, args.h_rotate_light) hdr_log_now = rotate_panorama_around_horizontal_axis(hdr_log_now, args.h_rotate_light) ldr[:,:,frame_idx:frame_idx+1,...] = torch.tensor(ldr_now).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) hdr_log[:,:,frame_idx:frame_idx+1,...] = torch.tensor(hdr_log_now).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) batch["ldr"] = ldr batch["hdr_log"] = hdr_log elif args.w_rotate_light != 0: for frame_idx in range(args.num_frames): ldr_now = np.array(ldr[0,:,frame_idx,...].permute(1, 2, 0)) # c h w -> h w c hdr_log_now = np.array(hdr_log[0,:,frame_idx,...].permute(1, 2, 0)) ldr_rotate = np.concatenate((ldr_now[:,args.w_rotate_light:], ldr_now[:,:args.w_rotate_light]), axis=1) hdr_log_rotate = np.concatenate((hdr_log_now[:,args.w_rotate_light:], hdr_log_now[:,:args.w_rotate_light]), axis=1) ldr[:,:,frame_idx:frame_idx+1,...] = torch.tensor(ldr_rotate).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) hdr_log[:,:,frame_idx:frame_idx+1,...] = torch.tensor(hdr_log_rotate).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) batch["ldr"] = ldr batch["hdr_log"] = hdr_log if (args.use_fixed_frame_and_w_rotate_light or args.use_fixed_frame_and_h_rotate_light or args.use_rotate_light) and args.num_frames == 1: continue if (args.use_fixed_frame_and_w_rotate_light or args.use_fixed_frame_and_h_rotate_light or args.use_rotate_light) and args.num_frames != 1: # Repeat the first frame for num_frames times if not args.use_rotate_light: input_rgb = input_rgb[:, :, :1, :, :].repeat(1, 1, args.num_frames, 1, 1) batch["source_video"] = input_rgb basecolor = basecolor[:, :, :1, :, :].repeat(1, 1, args.num_frames, 1, 1) batch["basecolor"] = basecolor depth = depth[:, :, :1, :, :].repeat(1, 1, args.num_frames, 1, 1) batch["depth"] = depth metallic = metallic[:, :, :1, :, :].repeat(1, 1, args.num_frames, 1, 1) batch["metallic"] = metallic normal = normal[:, :, :1, :, :].repeat(1, 1, args.num_frames, 1, 1) batch["normal"] = normal roughness = roughness[:, :, :1, :, :].repeat(1, 1, args.num_frames, 1, 1) batch["roughness"] = roughness if args.use_fixed_frame_and_h_rotate_light: ldr_first = np.array(ldr[0,:,0,...].permute(1, 2, 0)) # c h w -> h w c hdr_log_first = np.array(hdr_log[0,:,0,...].permute(1, 2, 0)) x_rotate_list = [int(360 / (args.num_frames-1) * (i+1)) for i in range(args.num_frames-1)] for frame_idx, x_rotate in enumerate(x_rotate_list): ldr_now = rotate_panorama_around_horizontal_axis(ldr_first, x_rotate) hdr_log_now = rotate_panorama_around_horizontal_axis(hdr_log_first, x_rotate) ldr[:,:,frame_idx+1:frame_idx+2,...] = torch.tensor(ldr_now).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) hdr_log[:,:,frame_idx+1:frame_idx+2,...] = torch.tensor(hdr_log_now).permute(2, 0, 1).unsqueeze(0).unsqueeze(2) batch["ldr"] = ldr batch["hdr_log"] = hdr_log else: ldr_first = ldr[:,:,:1,...] hdr_log_first = hdr_log[:,:,:1,...] _, c, _, h, w = ldr_first.shape y_rotate_list = [int(w / (args.num_frames-1) * (i+1)) for i in range(args.num_frames-1)] for frame_idx, y_rotate in enumerate(y_rotate_list): try: ldr_now = torch.zeros_like(ldr_first) ldr_now[:, :, :, :, -y_rotate:] = ldr_first[:, :, :, :, :y_rotate] ldr_now[:, :, :, :, :-y_rotate] = ldr_first[:, :, :, :, y_rotate:] hdr_log_now = torch.zeros_like(hdr_log_first) hdr_log_now[:, :, :, :, -y_rotate:] = hdr_log_first[:, :, :, :, :y_rotate] hdr_log_now[:, :, :, :, :-y_rotate] = hdr_log_first[:, :, :, :, y_rotate:] ldr[:,:,frame_idx+1:frame_idx+2,...] = ldr_now hdr_log[:,:,frame_idx+1:frame_idx+2,...] = hdr_log_now except: pdb.set_trace() batch["ldr"] = ldr batch["hdr_log"] = hdr_log seq_name = os.path.basename(dir_path) if args.use_rotate_light: seq_name += "_w_rotate_light" elif args.use_fixed_frame_and_w_rotate_light: seq_name += "_fixed_frame_and_w_rotate_light" elif args.use_fixed_frame_and_h_rotate_light: seq_name += "_fixed_frame_and_h_rotate_light" model_name = os.path.basename(args.ckpt_path).replace("ckpt", f'{seq_name}_{args.height}_{args.width}') if not args.use_ref_image: model_name += f'_no_ref_image' elif args.use_muti_ref_image: model_name += f'_ref_muti_image' else: model_name += f'_ref_mid_image' if args.env_map_path is not None: env_str = os.path.basename(args.env_map_path) model_name += f'_envdir_{env_str}' if args.drop_mr: model_name += f'_drop_mr' if args.padding_resolution: model_name += f'_padding' if args.full_resolution: model_name += f'_full' if args.num_inference_steps != 50: model_name += f'_steps{args.num_inference_steps}' if args.wo_ref_weight != 0.0: args.wo_ref_weight = round(args.wo_ref_weight, 2) model_name += f'_worw{args.wo_ref_weight}' if args.h_rotate_light != 0: model_name += f'_h-rotation-{args.h_rotate_light}' if args.w_rotate_light != 0: model_name += f'_w-rotation-{args.w_rotate_light}' if args.ref_image_path_with_idddx is not None: model_name += f'_ref_hr_{batch_idx}' model_name += f'_frames{args.num_frames}' if args.num_frames in [1]: save_path = os.path.join(output_dir, f"{model_name}_cfg{args.cfg_scale}_render.png") else: save_path = os.path.join(output_dir, f"{model_name}_cfg{args.cfg_scale}_video.mp4") video, envs = pipe( prompt=prompt, negative_prompt="", batch=batch, height=args.height, width=args.width, num_frames=args.num_frames, cfg_scale=args.cfg_scale, num_inference_steps=args.num_inference_steps, seed=0, tiled=True, wo_ref_weight=args.wo_ref_weight, use_muti_ref_image=args.use_muti_ref_image ) input_rgb = pipe.tensor2video(input_rgb[0]) basecolor = pipe.tensor2video(basecolor[0]) depth = pipe.tensor2video(depth[0]) metallic = pipe.tensor2video(metallic[0]) normal = pipe.tensor2video(normal[0]) roughness = pipe.tensor2video(roughness[0]) ldr = pipe.tensor2video(ldr[0]) hdr_log = pipe.tensor2video(hdr_log[0]) env_dir = pipe.tensor2video(env_dir[0]) if args.use_ref_image: ref_image = batch["ref_image"].repeat(1, 1, args.num_frames, 1, 1) # b c t h w ref_image = pipe.tensor2video(ref_image[0]) stitched_results = stitch_frames([basecolor, metallic, roughness, depth, normal, ldr, ref_image, video, input_rgb, env_dir, envs, env_dir], rows=4, cols=3) else: stitched_results = stitch_frames([basecolor, metallic, roughness, depth, normal, ldr, env_dir, video, input_rgb, env_dir, envs, env_dir], rows=4, cols=3) print(f'Finish the inference of {model_name}.') if args.num_frames in [1]: output = video[0] output_all = stitched_results[0] output.save(os.path.join(output_dir, f"{model_name}_cfg{args.cfg_scale}_render.png")) output_all.save(os.path.join(output_dir, f"{model_name}_cfg{args.cfg_scale}.png")) else: save_video(video, os.path.join(output_dir, f"{model_name}_cfg{args.cfg_scale}_video.mp4"), fps=30, quality=args.quality) save_video(stitched_results, os.path.join(output_dir, f"{model_name}_cfg{args.cfg_scale}.mp4"), fps=30, quality=5) if args.output_path is not None: if '.mp4' in args.output_path: save_video(video, args.output_path, fps=30, quality=args.quality) elif '.png' in args.output_path: output = video[0] output.save(args.output_path) else: print(f'Error output_path: {args.output_path}')