Spaces:
Running
on
L40S
Running
on
L40S
Upload 5 files
Browse files- camera_utils.py +33 -0
- config.py +16 -0
- model_loader.py +330 -0
- test_data.py +128 -0
- 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
|