jbilcke-hf HF Staff commited on
Commit
099dc67
·
verified ·
1 Parent(s): 2932acc

Upload 5 files

Browse files
Files changed (5) hide show
  1. camera_utils.py +33 -0
  2. config.py +16 -0
  3. model_loader.py +330 -0
  4. test_data.py +128 -0
  5. video_processor.py +118 -0
camera_utils.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ class Camera(object):
4
+ def __init__(self, c2w):
5
+ c2w_mat = np.array(c2w).reshape(4, 4)
6
+ self.c2w_mat = c2w_mat
7
+ self.w2c_mat = np.linalg.inv(c2w_mat)
8
+
9
+ def parse_matrix(matrix_str):
10
+ """Parse camera matrix string from JSON format"""
11
+ rows = matrix_str.strip().split('] [')
12
+ matrix = []
13
+ for row in rows:
14
+ row = row.replace('[', '').replace(']', '')
15
+ matrix.append(list(map(float, row.split())))
16
+ return np.array(matrix)
17
+
18
+ def get_relative_pose(cam_params):
19
+ """Calculate relative camera poses"""
20
+ abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
21
+ abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]
22
+
23
+ cam_to_origin = 0
24
+ target_cam_c2w = np.array([
25
+ [1, 0, 0, 0],
26
+ [0, 1, 0, -cam_to_origin],
27
+ [0, 0, 1, 0],
28
+ [0, 0, 0, 1]
29
+ ])
30
+ abs2rel = target_cam_c2w @ abs_w2cs[0]
31
+ ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
32
+ ret_poses = np.array(ret_poses, dtype=np.float32)
33
+ return ret_poses
config.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Camera transformation types
2
+ CAMERA_TRANSFORMATIONS = {
3
+ "1": "Pan Right",
4
+ "2": "Pan Left",
5
+ "3": "Tilt Up",
6
+ "4": "Tilt Down",
7
+ "5": "Zoom In",
8
+ "6": "Zoom Out",
9
+ "7": "Translate Up (with rotation)",
10
+ "8": "Translate Down (with rotation)",
11
+ "9": "Arc Left (with rotation)",
12
+ "10": "Arc Right (with rotation)"
13
+ }
14
+
15
+ # Define test data directory
16
+ TEST_DATA_DIR = "example_test_data"
model_loader.py ADDED
@@ -0,0 +1,330 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import torch.nn as nn
4
+ import logging
5
+ from pathlib import Path
6
+ from huggingface_hub import hf_hub_download
7
+ from diffsynth import ModelManager, WanVideoReCamMasterPipeline
8
+
9
+ logger = logging.getLogger(__name__)
10
+
11
+ # Get model storage path from environment variable or use default
12
+ MODELS_ROOT_DIR = os.environ.get("RECAMMASTER_MODELS_DIR", "/data/models")
13
+ logger.info(f"Using models root directory: {MODELS_ROOT_DIR}")
14
+
15
+ # Define model repositories and files
16
+ WAN21_REPO_ID = "Wan-AI/Wan2.1-T2V-1.3B"
17
+ WAN21_LOCAL_DIR = f"{MODELS_ROOT_DIR}/Wan-AI/Wan2.1-T2V-1.3B"
18
+ WAN21_FILES = [
19
+ "diffusion_pytorch_model.safetensors",
20
+ "models_t5_umt5-xxl-enc-bf16.pth",
21
+ "Wan2.1_VAE.pth"
22
+ ]
23
+
24
+ # Define tokenizer files to download
25
+ UMT5_XXL_TOKENIZER_FILES = [
26
+ "google/umt5-xxl/special_tokens_map.json",
27
+ "google/umt5-xxl/spiece.model",
28
+ "google/umt5-xxl/tokenizer.json",
29
+ "google/umt5-xxl/tokenizer_config.json"
30
+ ]
31
+
32
+ RECAMMASTER_REPO_ID = "KwaiVGI/ReCamMaster-Wan2.1"
33
+ RECAMMASTER_CHECKPOINT_FILE = "step20000.ckpt"
34
+ RECAMMASTER_LOCAL_DIR = f"{MODELS_ROOT_DIR}/ReCamMaster/checkpoints"
35
+
36
+ class ModelLoader:
37
+ def __init__(self):
38
+ self.model_manager = None
39
+ self.pipe = None
40
+ self.is_loaded = False
41
+
42
+ def download_umt5_xxl_tokenizer(self, progress_callback=None):
43
+ """Download UMT5-XXL tokenizer files from HuggingFace"""
44
+
45
+ total_files = len(UMT5_XXL_TOKENIZER_FILES)
46
+ downloaded_paths = []
47
+
48
+ for i, file_path in enumerate(UMT5_XXL_TOKENIZER_FILES):
49
+ local_dir = f"{WAN21_LOCAL_DIR}/{os.path.dirname(file_path)}"
50
+ filename = os.path.basename(file_path)
51
+ full_local_path = f"{WAN21_LOCAL_DIR}/{file_path}"
52
+
53
+ # Update progress
54
+ if progress_callback:
55
+ progress_callback(i/total_files, desc=f"Checking tokenizer file {i+1}/{total_files}: {filename}")
56
+
57
+ # Check if already exists
58
+ if os.path.exists(full_local_path):
59
+ logger.info(f"✓ Tokenizer file {filename} already exists at {full_local_path}")
60
+ downloaded_paths.append(full_local_path)
61
+ continue
62
+
63
+ # Create directory if it doesn't exist
64
+ os.makedirs(local_dir, exist_ok=True)
65
+
66
+ # Download the file
67
+ logger.info(f"Downloading tokenizer file {filename} from {WAN21_REPO_ID}/{file_path}...")
68
+
69
+ if progress_callback:
70
+ progress_callback(i/total_files, desc=f"Downloading tokenizer file {i+1}/{total_files}: {filename}")
71
+
72
+ try:
73
+ # Download using huggingface_hub
74
+ downloaded_path = hf_hub_download(
75
+ repo_id=WAN21_REPO_ID,
76
+ filename=file_path,
77
+ local_dir=WAN21_LOCAL_DIR,
78
+ local_dir_use_symlinks=False
79
+ )
80
+ logger.info(f"✓ Successfully downloaded tokenizer file {filename} to {downloaded_path}!")
81
+ downloaded_paths.append(downloaded_path)
82
+ except Exception as e:
83
+ logger.error(f"✗ Error downloading tokenizer file {filename}: {e}")
84
+ raise
85
+
86
+ if progress_callback:
87
+ progress_callback(1.0, desc=f"All tokenizer files downloaded successfully!")
88
+
89
+ return downloaded_paths
90
+
91
+ def download_wan21_models(self, progress_callback=None):
92
+ """Download Wan2.1 model files from HuggingFace"""
93
+
94
+ total_files = len(WAN21_FILES)
95
+ downloaded_paths = []
96
+
97
+ # Create directory if it doesn't exist
98
+ Path(WAN21_LOCAL_DIR).mkdir(parents=True, exist_ok=True)
99
+
100
+ for i, filename in enumerate(WAN21_FILES):
101
+ local_path = Path(WAN21_LOCAL_DIR) / filename
102
+
103
+ # Update progress
104
+ if progress_callback:
105
+ progress_callback(i/total_files, desc=f"Checking Wan2.1 file {i+1}/{total_files}: {filename}")
106
+
107
+ # Check if already exists
108
+ if local_path.exists():
109
+ logger.info(f"✓ {filename} already exists at {local_path}")
110
+ downloaded_paths.append(str(local_path))
111
+ continue
112
+
113
+ # Download the file
114
+ logger.info(f"Downloading {filename} from {WAN21_REPO_ID}...")
115
+
116
+ if progress_callback:
117
+ progress_callback(i/total_files, desc=f"Downloading Wan2.1 file {i+1}/{total_files}: {filename}")
118
+
119
+ try:
120
+ # Download using huggingface_hub
121
+ downloaded_path = hf_hub_download(
122
+ repo_id=WAN21_REPO_ID,
123
+ filename=filename,
124
+ local_dir=WAN21_LOCAL_DIR,
125
+ local_dir_use_symlinks=False
126
+ )
127
+ logger.info(f"✓ Successfully downloaded {filename} to {downloaded_path}!")
128
+ downloaded_paths.append(downloaded_path)
129
+ except Exception as e:
130
+ logger.error(f"✗ Error downloading {filename}: {e}")
131
+ raise
132
+
133
+ if progress_callback:
134
+ progress_callback(1.0, desc=f"All Wan2.1 models downloaded successfully!")
135
+
136
+ return downloaded_paths
137
+
138
+ def download_recammaster_checkpoint(self, progress_callback=None):
139
+ """Download ReCamMaster checkpoint from HuggingFace using huggingface_hub"""
140
+ checkpoint_path = Path(RECAMMASTER_LOCAL_DIR) / RECAMMASTER_CHECKPOINT_FILE
141
+
142
+ # Check if already exists
143
+ if checkpoint_path.exists():
144
+ logger.info(f"✓ ReCamMaster checkpoint already exists at {checkpoint_path}")
145
+ return checkpoint_path
146
+
147
+ # Create directory if it doesn't exist
148
+ Path(RECAMMASTER_LOCAL_DIR).mkdir(parents=True, exist_ok=True)
149
+
150
+ # Download the checkpoint
151
+ logger.info("Downloading ReCamMaster checkpoint from HuggingFace...")
152
+ logger.info(f"Repository: {RECAMMASTER_REPO_ID}")
153
+ logger.info(f"File: {RECAMMASTER_CHECKPOINT_FILE}")
154
+ logger.info(f"Destination: {checkpoint_path}")
155
+
156
+ if progress_callback:
157
+ progress_callback(0.0, desc=f"Downloading ReCamMaster checkpoint...")
158
+
159
+ try:
160
+ # Download using huggingface_hub
161
+ downloaded_path = hf_hub_download(
162
+ repo_id=RECAMMASTER_REPO_ID,
163
+ filename=RECAMMASTER_CHECKPOINT_FILE,
164
+ local_dir=RECAMMASTER_LOCAL_DIR,
165
+ local_dir_use_symlinks=False
166
+ )
167
+ logger.info(f"✓ Successfully downloaded ReCamMaster checkpoint to {downloaded_path}!")
168
+
169
+ if progress_callback:
170
+ progress_callback(1.0, desc=f"ReCamMaster checkpoint downloaded successfully!")
171
+
172
+ return downloaded_path
173
+ except Exception as e:
174
+ logger.error(f"✗ Error downloading checkpoint: {e}")
175
+ raise
176
+
177
+ def create_symlink_for_tokenizer(self):
178
+ """Create symlink for google/umt5-xxl to handle potential path issues"""
179
+ try:
180
+ google_dir = f"{MODELS_ROOT_DIR}/google"
181
+ if not os.path.exists(google_dir):
182
+ os.makedirs(google_dir, exist_ok=True)
183
+
184
+ umt5_xxl_symlink = f"{google_dir}/umt5-xxl"
185
+ umt5_xxl_source = f"{WAN21_LOCAL_DIR}/google/umt5-xxl"
186
+
187
+ # Create a symlink if it doesn't exist
188
+ if not os.path.exists(umt5_xxl_symlink) and os.path.exists(umt5_xxl_source):
189
+ if os.name == 'nt': # Windows
190
+ import ctypes
191
+ kdll = ctypes.windll.LoadLibrary("kernel32.dll")
192
+ kdll.CreateSymbolicLinkA(umt5_xxl_symlink.encode(), umt5_xxl_source.encode(), 1)
193
+ else: # Unix/Linux
194
+ os.symlink(umt5_xxl_source, umt5_xxl_symlink)
195
+ logger.info(f"Created symlink from {umt5_xxl_source} to {umt5_xxl_symlink}")
196
+ except Exception as e:
197
+ logger.warning(f"Could not create symlink for google/umt5-xxl: {str(e)}")
198
+ # This is a warning, not an error, as we'll try to proceed anyway
199
+
200
+ def load_models(self, progress_callback=None):
201
+ """Load the ReCamMaster models"""
202
+
203
+ if self.is_loaded:
204
+ return "Models already loaded!"
205
+
206
+ try:
207
+ logger.info("Starting model loading...")
208
+
209
+ # Import test data creator
210
+ from test_data import create_test_data_structure
211
+
212
+ # First create the test data structure
213
+ if progress_callback:
214
+ progress_callback(0.05, desc="Setting up test data structure...")
215
+
216
+ try:
217
+ create_test_data_structure(progress_callback)
218
+ except Exception as e:
219
+ error_msg = f"Error creating test data structure: {str(e)}"
220
+ logger.error(error_msg)
221
+ return error_msg
222
+
223
+ # Second, ensure the checkpoint is downloaded
224
+ if progress_callback:
225
+ progress_callback(0.1, desc="Checking for ReCamMaster checkpoint...")
226
+
227
+ try:
228
+ ckpt_path = self.download_recammaster_checkpoint(progress_callback)
229
+ logger.info(f"Using checkpoint at {ckpt_path}")
230
+ except Exception as e:
231
+ error_msg = f"Error downloading ReCamMaster checkpoint: {str(e)}"
232
+ logger.error(error_msg)
233
+ return error_msg
234
+
235
+ # Third, download Wan2.1 models if needed
236
+ if progress_callback:
237
+ progress_callback(0.2, desc="Checking for Wan2.1 models...")
238
+
239
+ try:
240
+ wan21_paths = self.download_wan21_models(progress_callback)
241
+ logger.info(f"Using Wan2.1 models: {wan21_paths}")
242
+ except Exception as e:
243
+ error_msg = f"Error downloading Wan2.1 models: {str(e)}"
244
+ logger.error(error_msg)
245
+ return error_msg
246
+
247
+ # Fourth, download UMT5-XXL tokenizer files
248
+ if progress_callback:
249
+ progress_callback(0.3, desc="Checking for UMT5-XXL tokenizer files...")
250
+
251
+ try:
252
+ tokenizer_paths = self.download_umt5_xxl_tokenizer(progress_callback)
253
+ logger.info(f"Using UMT5-XXL tokenizer files: {tokenizer_paths}")
254
+ except Exception as e:
255
+ error_msg = f"Error downloading UMT5-XXL tokenizer files: {str(e)}"
256
+ logger.error(error_msg)
257
+ return error_msg
258
+
259
+ # Now, load the models
260
+ if progress_callback:
261
+ progress_callback(0.4, desc="Loading model manager...")
262
+
263
+ # Create symlink for tokenizer
264
+ self.create_symlink_for_tokenizer()
265
+
266
+ # Load Wan2.1 pre-trained models
267
+ self.model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
268
+
269
+ if progress_callback:
270
+ progress_callback(0.5, desc="Loading Wan2.1 models...")
271
+
272
+ # Build full paths for the model files
273
+ model_files = [f"{WAN21_LOCAL_DIR}/{filename}" for filename in WAN21_FILES]
274
+
275
+ for model_file in model_files:
276
+ logger.info(f"Loading model from: {model_file}")
277
+ if not os.path.exists(model_file):
278
+ error_msg = f"Error: Model file not found: {model_file}"
279
+ logger.error(error_msg)
280
+ return error_msg
281
+
282
+ # Set environment variable for transformers to find the tokenizer
283
+ os.environ["TRANSFORMERS_CACHE"] = MODELS_ROOT_DIR
284
+ os.environ["TOKENIZERS_PARALLELISM"] = "false" # Disable tokenizers parallelism warning
285
+
286
+ self.model_manager.load_models(model_files)
287
+
288
+ if progress_callback:
289
+ progress_callback(0.7, desc="Creating pipeline...")
290
+
291
+ self.pipe = WanVideoReCamMasterPipeline.from_model_manager(self.model_manager, device="cuda")
292
+
293
+ if progress_callback:
294
+ progress_callback(0.8, desc="Initializing ReCamMaster modules...")
295
+
296
+ # Initialize additional modules introduced in ReCamMaster
297
+ dim = self.pipe.dit.blocks[0].self_attn.q.weight.shape[0]
298
+ for block in self.pipe.dit.blocks:
299
+ block.cam_encoder = nn.Linear(12, dim)
300
+ block.projector = nn.Linear(dim, dim)
301
+ block.cam_encoder.weight.data.zero_()
302
+ block.cam_encoder.bias.data.zero_()
303
+ block.projector.weight = nn.Parameter(torch.eye(dim))
304
+ block.projector.bias = nn.Parameter(torch.zeros(dim))
305
+
306
+ if progress_callback:
307
+ progress_callback(0.9, desc="Loading ReCamMaster checkpoint...")
308
+
309
+ # Load ReCamMaster checkpoint
310
+ if not os.path.exists(ckpt_path):
311
+ error_msg = f"Error: ReCamMaster checkpoint not found at {ckpt_path} even after download attempt."
312
+ logger.error(error_msg)
313
+ return error_msg
314
+
315
+ state_dict = torch.load(ckpt_path, map_location="cpu")
316
+ self.pipe.dit.load_state_dict(state_dict, strict=True)
317
+ self.pipe.to("cuda")
318
+ self.pipe.to(dtype=torch.bfloat16)
319
+
320
+ self.is_loaded = True
321
+
322
+ if progress_callback:
323
+ progress_callback(1.0, desc="Models loaded successfully!")
324
+
325
+ logger.info("Models loaded successfully!")
326
+ return "Models loaded successfully!"
327
+
328
+ except Exception as e:
329
+ logger.error(f"Error loading models: {str(e)}")
330
+ return f"Error loading models: {str(e)}"
test_data.py ADDED
@@ -0,0 +1,128 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import numpy as np
3
+ import logging
4
+ from pathlib import Path
5
+ from config import TEST_DATA_DIR
6
+
7
+ logger = logging.getLogger(__name__)
8
+
9
+ def create_test_data_structure(progress_callback=None):
10
+ """Create sample camera extrinsics data for testing"""
11
+
12
+ if progress_callback:
13
+ progress_callback(0.0, desc="Creating test data structure...")
14
+
15
+ # Create directories
16
+ data_dir = Path(f"{TEST_DATA_DIR}/cameras")
17
+ videos_dir = Path(f"{TEST_DATA_DIR}/videos")
18
+ data_dir.mkdir(parents=True, exist_ok=True)
19
+ videos_dir.mkdir(parents=True, exist_ok=True)
20
+
21
+ camera_file = data_dir / "camera_extrinsics.json"
22
+
23
+ # Skip if file already exists
24
+ if camera_file.exists():
25
+ logger.info(f"✓ Camera extrinsics already exist at {camera_file}")
26
+
27
+ if progress_callback:
28
+ progress_callback(1.0, desc="Test data structure already exists")
29
+
30
+ return
31
+
32
+ if progress_callback:
33
+ progress_callback(0.3, desc="Generating camera extrinsics data...")
34
+
35
+ # Generate sample camera data
36
+ camera_data = {}
37
+
38
+ # Create 81 frames with 10 camera trajectories each
39
+ for frame_idx in range(81):
40
+ frame_key = f"frame{frame_idx}"
41
+ camera_data[frame_key] = {}
42
+
43
+ for cam_idx in range(1, 11): # Camera types 1-10
44
+ # Create a sample camera matrix (this is just an example - replace with actual logic if needed)
45
+ # In reality, these would be calculated based on specific camera movement patterns
46
+
47
+ # Create a base identity matrix
48
+ base_matrix = np.eye(4)
49
+
50
+ # Add some variation based on frame and camera type
51
+ # This is a simplistic example - real camera movements would be more complex
52
+ if cam_idx == 1: # Pan Right
53
+ base_matrix[0, 3] = 0.01 * frame_idx # Move right over time
54
+ elif cam_idx == 2: # Pan Left
55
+ base_matrix[0, 3] = -0.01 * frame_idx # Move left over time
56
+ elif cam_idx == 3: # Tilt Up
57
+ # Rotate around X-axis
58
+ angle = 0.005 * frame_idx
59
+ base_matrix[1, 1] = np.cos(angle)
60
+ base_matrix[1, 2] = -np.sin(angle)
61
+ base_matrix[2, 1] = np.sin(angle)
62
+ base_matrix[2, 2] = np.cos(angle)
63
+ elif cam_idx == 4: # Tilt Down
64
+ # Rotate around X-axis (opposite direction)
65
+ angle = -0.005 * frame_idx
66
+ base_matrix[1, 1] = np.cos(angle)
67
+ base_matrix[1, 2] = -np.sin(angle)
68
+ base_matrix[2, 1] = np.sin(angle)
69
+ base_matrix[2, 2] = np.cos(angle)
70
+ elif cam_idx == 5: # Zoom In
71
+ base_matrix[2, 3] = -0.01 * frame_idx # Move forward over time
72
+ elif cam_idx == 6: # Zoom Out
73
+ base_matrix[2, 3] = 0.01 * frame_idx # Move backward over time
74
+ elif cam_idx == 7: # Translate Up (with rotation)
75
+ base_matrix[1, 3] = 0.01 * frame_idx # Move up over time
76
+ angle = 0.003 * frame_idx
77
+ base_matrix[0, 0] = np.cos(angle)
78
+ base_matrix[0, 2] = np.sin(angle)
79
+ base_matrix[2, 0] = -np.sin(angle)
80
+ base_matrix[2, 2] = np.cos(angle)
81
+ elif cam_idx == 8: # Translate Down (with rotation)
82
+ base_matrix[1, 3] = -0.01 * frame_idx # Move down over time
83
+ angle = -0.003 * frame_idx
84
+ base_matrix[0, 0] = np.cos(angle)
85
+ base_matrix[0, 2] = np.sin(angle)
86
+ base_matrix[2, 0] = -np.sin(angle)
87
+ base_matrix[2, 2] = np.cos(angle)
88
+ elif cam_idx == 9: # Arc Left (with rotation)
89
+ angle = 0.005 * frame_idx
90
+ radius = 2.0
91
+ base_matrix[0, 3] = -radius * np.sin(angle)
92
+ base_matrix[2, 3] = -radius * np.cos(angle) + radius
93
+ # Rotate to look at center
94
+ look_angle = angle + np.pi
95
+ base_matrix[0, 0] = np.cos(look_angle)
96
+ base_matrix[0, 2] = np.sin(look_angle)
97
+ base_matrix[2, 0] = -np.sin(look_angle)
98
+ base_matrix[2, 2] = np.cos(look_angle)
99
+ elif cam_idx == 10: # Arc Right (with rotation)
100
+ angle = -0.005 * frame_idx
101
+ radius = 2.0
102
+ base_matrix[0, 3] = -radius * np.sin(angle)
103
+ base_matrix[2, 3] = -radius * np.cos(angle) + radius
104
+ # Rotate to look at center
105
+ look_angle = angle + np.pi
106
+ base_matrix[0, 0] = np.cos(look_angle)
107
+ base_matrix[0, 2] = np.sin(look_angle)
108
+ base_matrix[2, 0] = -np.sin(look_angle)
109
+ base_matrix[2, 2] = np.cos(look_angle)
110
+
111
+ # Format the matrix as a string (as expected by the app)
112
+ matrix_str = ' '.join([' '.join([str(base_matrix[i, j]) for j in range(4)]) for i in range(4)])
113
+ matrix_str = '[ ' + matrix_str.replace(' ', ' ] [ ', 3) + ' ]'
114
+
115
+ camera_data[frame_key][f"cam{cam_idx:02d}"] = matrix_str
116
+
117
+ if progress_callback:
118
+ progress_callback(0.7, desc="Saving camera extrinsics data...")
119
+
120
+ # Save camera extrinsics to JSON file
121
+ with open(camera_file, 'w') as f:
122
+ json.dump(camera_data, f, indent=2)
123
+
124
+ logger.info(f"Created sample camera extrinsics at {camera_file}")
125
+ logger.info(f"Created directory for example videos at {videos_dir}")
126
+
127
+ if progress_callback:
128
+ progress_callback(1.0, desc="Test data structure created successfully!")
video_processor.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import json
4
+ import imageio
5
+ from PIL import Image
6
+ from torchvision.transforms import v2
7
+ from einops import rearrange
8
+ import torchvision
9
+ import logging
10
+ from config import TEST_DATA_DIR
11
+ from camera_utils import Camera, parse_matrix, get_relative_pose
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+ class VideoProcessor:
16
+ def __init__(self, pipe):
17
+ self.pipe = pipe
18
+ self.height = 480
19
+ self.width = 832
20
+
21
+ # Create frame processor
22
+ self.frame_process = v2.Compose([
23
+ v2.CenterCrop(size=(self.height, self.width)),
24
+ v2.Resize(size=(self.height, self.width), antialias=True),
25
+ v2.ToTensor(),
26
+ v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
27
+ ])
28
+
29
+ def crop_and_resize(self, image):
30
+ """Crop and resize image to match target dimensions"""
31
+ width_img, height_img = image.size
32
+ scale = max(self.width / width_img, self.height / height_img)
33
+ image = torchvision.transforms.functional.resize(
34
+ image,
35
+ (round(height_img*scale), round(width_img*scale)),
36
+ interpolation=torchvision.transforms.InterpolationMode.BILINEAR
37
+ )
38
+ return image
39
+
40
+ def load_video_frames(self, video_path):
41
+ """Load and process video frames"""
42
+ reader = imageio.get_reader(video_path)
43
+ frames = []
44
+
45
+ for i in range(81): # ReCamMaster needs exactly 81 frames
46
+ try:
47
+ frame = reader.get_data(i)
48
+ frame = Image.fromarray(frame)
49
+ frame = self.crop_and_resize(frame)
50
+ frame = self.frame_process(frame)
51
+ frames.append(frame)
52
+ except:
53
+ # If we run out of frames, repeat the last one
54
+ if frames:
55
+ frames.append(frames[-1])
56
+ else:
57
+ raise ValueError("Video is too short!")
58
+
59
+ reader.close()
60
+
61
+ frames = torch.stack(frames, dim=0)
62
+ frames = rearrange(frames, "T C H W -> C T H W")
63
+ video_tensor = frames.unsqueeze(0) # Add batch dimension
64
+
65
+ return video_tensor
66
+
67
+ def load_camera_trajectory(self, cam_type):
68
+ """Load camera trajectory for the selected type"""
69
+ tgt_camera_path = f"./{TEST_DATA_DIR}/cameras/camera_extrinsics.json"
70
+ with open(tgt_camera_path, 'r') as file:
71
+ cam_data = json.load(file)
72
+
73
+ # Get camera trajectory for selected type
74
+ cam_idx = list(range(81))[::4] # Sample every 4 frames
75
+ traj = [parse_matrix(cam_data[f"frame{idx}"][f"cam{int(cam_type):02d}"]) for idx in cam_idx]
76
+ traj = np.stack(traj).transpose(0, 2, 1)
77
+
78
+ c2ws = []
79
+ for c2w in traj:
80
+ c2w = c2w[:, [1, 2, 0, 3]]
81
+ c2w[:3, 1] *= -1.
82
+ c2w[:3, 3] /= 100
83
+ c2ws.append(c2w)
84
+
85
+ tgt_cam_params = [Camera(cam_param) for cam_param in c2ws]
86
+ relative_poses = []
87
+ for i in range(len(tgt_cam_params)):
88
+ relative_pose = get_relative_pose([tgt_cam_params[0], tgt_cam_params[i]])
89
+ relative_poses.append(torch.as_tensor(relative_pose)[:,:3,:][1])
90
+
91
+ pose_embedding = torch.stack(relative_poses, dim=0) # 21x3x4
92
+ pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')
93
+ camera_tensor = pose_embedding.to(torch.bfloat16).unsqueeze(0) # Add batch dimension
94
+
95
+ return camera_tensor
96
+
97
+ def process_video(self, video_path, text_prompt, cam_type):
98
+ """Process video through ReCamMaster model"""
99
+
100
+ # Load video frames
101
+ video_tensor = self.load_video_frames(video_path)
102
+
103
+ # Load camera trajectory
104
+ camera_tensor = self.load_camera_trajectory(cam_type)
105
+
106
+ # Generate video with ReCamMaster
107
+ video = self.pipe(
108
+ prompt=[text_prompt],
109
+ negative_prompt=["worst quality, low quality, blurry, jittery, distorted"],
110
+ source_video=video_tensor,
111
+ target_camera=camera_tensor,
112
+ cfg_scale=5.0,
113
+ num_inference_steps=50,
114
+ seed=0,
115
+ tiled=True
116
+ )
117
+
118
+ return video