Spaces:
Paused
Paused
File size: 3,031 Bytes
05fcd0f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import torch
import numpy as np
import traceback
from diffusers_helper.utils import save_bcthw_as_mp4
@torch.no_grad()
def combine_videos_sequentially_from_tensors(processed_input_frames_np,
generated_frames_pt,
output_path,
target_fps,
crf_value):
"""
Combines processed input frames (NumPy) with generated frames (PyTorch Tensor) sequentially
and saves the result as an MP4 video using save_bcthw_as_mp4.
Args:
processed_input_frames_np: NumPy array of processed input frames (T_in, H, W_in, C), uint8.
generated_frames_pt: PyTorch tensor of generated frames (B_gen, C_gen, T_gen, H, W_gen), float32 [-1,1].
(This will be history_pixels from worker.py)
output_path: Path to save the combined video.
target_fps: FPS for the output combined video.
crf_value: CRF value for video encoding.
Returns:
Path to the combined video, or None if an error occurs.
"""
try:
# 1. Convert processed_input_frames_np to PyTorch tensor BCTHW, float32, [-1,1]
# processed_input_frames_np shape: (T_in, H, W_in, C)
input_frames_pt = torch.from_numpy(processed_input_frames_np).float() / 127.5 - 1.0 # (T,H,W,C)
input_frames_pt = input_frames_pt.permute(3, 0, 1, 2) # (C,T,H,W)
input_frames_pt = input_frames_pt.unsqueeze(0) # (1,C,T,H,W) -> BCTHW
# Ensure generated_frames_pt is on the same device and dtype for concatenation
input_frames_pt = input_frames_pt.to(device=generated_frames_pt.device, dtype=generated_frames_pt.dtype)
# 2. Dimension Check (Heights and Widths should match)
# They should match, since the input frames should have been processed to match the generation resolution.
# But sanity check to ensure no mismatch occurs when the code is refactored.
if input_frames_pt.shape[3:] != generated_frames_pt.shape[3:]: # Compare (H,W)
print(f"Warning: Dimension mismatch for sequential combination! Input: {input_frames_pt.shape[3:]}, Generated: {generated_frames_pt.shape[3:]}.")
print("Attempting to proceed, but this might lead to errors or unexpected video output.")
# Potentially add resizing logic here if necessary, but for now, assume they match
# 3. Concatenate Tensors along the time dimension (dim=2 for BCTHW)
combined_video_pt = torch.cat([input_frames_pt, generated_frames_pt], dim=2)
# 4. Save
save_bcthw_as_mp4(combined_video_pt, output_path, fps=target_fps, crf=crf_value)
print(f"Sequentially combined video (from tensors) saved to {output_path}")
return output_path
except Exception as e:
print(f"Error in combine_videos_sequentially_from_tensors: {str(e)}")
import traceback
traceback.print_exc()
return None |