Coverage for media_utils.py: 86%
693 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1import os
2import torch
3import tempfile
4import math
5import logging
6import wave
7import asyncio
8import subprocess
9import json
10import re
11import aiofiles
12import aiofiles.os
13import textwrap
15import imageio
16# from imageio import formats
18import imageio_ffmpeg as ffmpeg
20from contextlib import asynccontextmanager
22import numpy as np
24from typing import List
25from typing import Sequence
26from typing import TypedDict
27from typing import Optional
28from typing import Union
29from typing import Dict
30from typing import Tuple
31from typing import Any
32from typing import AsyncIterator
34from aiofiles.threadpool.binary import AsyncBufferedIOBase
36from io import BytesIO
37from PIL import Image
38from PIL import ImageDraw
39from PIL import ImageFont
41from file_utils import read_file_bytes
42from file_utils import base64_to_binary
43from file_utils import save_base64_as_binary
44from file_utils import binary_to_base64
46PIX_FMT_RGB24 = "rgb24"
48# DEFAULT_VIDEO_CODEC = "libx264rgb"
49# DEFAULT_PIX_FMT = "rgb24"
50DEFAULT_VIDEO_CODEC = "libx264"
51DEFAULT_AUDIO_CODEC = "aac"
52DEFAULT_PIX_FMT = "yuv420p"
55def tensor_to_base64(tensor: torch.Tensor) -> str:
56 """Converts a PyTorch tensor to a base64-encoded string."""
57 if not isinstance(tensor, torch.Tensor):
58 raise TypeError(f"Expected torch.Tensor for tensor, got {type(tensor)}")
59 buffer = BytesIO()
60 torch.save(tensor, buffer)
61 buffer.seek(0)
62 tensor_bytes = buffer.read()
63 base64_str = binary_to_base64(tensor_bytes)
64 return base64_str
67def bytes_to_tensor(binary_data: bytes) -> torch.Tensor:
68 """Converts binary data to a PyTorch tensor."""
69 if not isinstance(binary_data, bytes):
70 raise TypeError(f"Expected bytes for binary_data, got {type(binary_data)}")
71 buffer = BytesIO(binary_data)
72 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
73 tensor = torch.load(buffer, map_location=device, weights_only=True)
74 return tensor
77def base64_to_tensor(base64_str: str) -> torch.Tensor:
78 """Converts a base64-encoded string to a PyTorch tensor."""
79 if not isinstance(base64_str, str):
80 raise TypeError(f"Expected str for base64_str, got {type(base64_str)}")
81 tensor_bytes = base64_to_binary(base64_str)
82 tensor = bytes_to_tensor(tensor_bytes)
83 return tensor
86def base64_to_video_frames(
87 video_base64: str,
88 video_format: str = "mp4",
89) -> List[Image.Image]:
90 """Converts a base64-encoded string to a list of PIL Image frames."""
91 if not isinstance(video_base64, str):
92 raise TypeError(f"Expected str for video_base64, got {type(video_base64)}")
93 video_bytes = base64_to_binary(video_base64)
94 video_buffer = BytesIO(video_bytes)
95 # fmt: imageio.Format = formats[video_format]
96 video_frames = imageio.get_reader(
97 video_buffer,
98 format=video_format) # type: ignore[arg-type]
99 frames = [
100 Image.fromarray(frame_np).convert("RGB")
101 for frame_np in video_frames # type: ignore[attr-defined]
102 ]
103 return frames
106def video_frames_to_base64(
107 video_frames: List[Image.Image],
108 fps: float = 30.0,
109 format: str = "mp4"
110) -> str:
111 """Converts a list of PIL Image frames to a base64-encoded video string."""
112 if not isinstance(video_frames, list):
113 raise TypeError(f"Expected list for video_frames, got {type(video_frames)}")
114 if len(video_frames) <= 0:
115 raise ValueError("video_frames cannot be empty")
116 video_buffer = BytesIO()
117 # fmt: imageio.Format = formats[format]
118 with imageio.get_writer(
119 video_buffer,
120 format=format, # type: ignore[arg-type]
121 fps=fps
122 ) as writer:
123 for frame in video_frames:
124 frame_np = np.array(frame)
125 writer.append_data(frame_np) # type: ignore[attr-defined]
126 video_bytes = video_buffer.getvalue()
127 video_base64 = binary_to_base64(video_bytes)
128 return video_base64
131def fix_json_like_string(s: str) -> str:
132 """Add double quotes around keys (only if not already quoted)."""
133 s = re.sub(r'([{,]\s*)(\w+)(\s*:\s*)', r'\1"\2"\3', s)
134 return s
137def chunk_audio_base64(
138 audio_base64: str,
139 start_seconds: float = 0.0,
140 end_seconds: float = 2.0,
141) -> str:
142 """Chunk audio base64 between start_seconds and end_seconds."""
143 if not isinstance(audio_base64, str):
144 raise TypeError(f"Expected str for audio_base64, got {type(audio_base64)}")
146 import soundfile
148 audio_bytes = base64_to_binary(audio_base64)
149 audio_buffer = BytesIO(audio_bytes)
150 audio_data, sample_rate = soundfile.read(audio_buffer)
151 start_sample = int(start_seconds * sample_rate)
152 end_sample = int(end_seconds * sample_rate)
153 if not (0 <= start_sample < end_sample <= len(audio_data)):
154 raise ValueError(
155 f"Invalid trim range {start_sample} to {end_sample} sample for audio with {len(audio_data)} samples.")
156 trimmed_audio_data = audio_data[start_sample:end_sample]
157 trimmed_buffer = BytesIO()
158 soundfile.write(
159 trimmed_buffer,
160 trimmed_audio_data,
161 sample_rate,
162 format="WAV")
163 trimmed_buffer.seek(0)
164 return binary_to_base64(trimmed_buffer.getvalue())
167def chunk_video_binary(
168 video_binary: bytes,
169 start_seconds: float = 0.0,
170 end_seconds: float = 2.0,
171 video_codec: str = DEFAULT_VIDEO_CODEC,
172 audio_codec: str = DEFAULT_AUDIO_CODEC,
173 pix_fmt: str = DEFAULT_PIX_FMT,
174 width: Optional[int] = None,
175 height: Optional[int] = None,
176) -> bytes:
177 """Chunk video binary between start_seconds and end_seconds."""
178 if not video_binary:
179 raise ValueError("Video input is empty.")
180 if not isinstance(video_binary, bytes):
181 raise TypeError(f"Expected bytes for video binary, got {type(video_binary)}")
182 if start_seconds < 0:
183 raise ValueError("start_seconds must be non-negative")
184 if end_seconds <= start_seconds:
185 raise ValueError("end_seconds must be greater than start_seconds")
187 with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as input_file, \
188 tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as output_file:
189 try:
190 # Write input bytes
191 input_file.write(video_binary)
192 input_file.flush()
194 # Run ffmpeg to trim and re-encode
195 cmd = [
196 "ffmpeg",
197 "-hide_banner",
198 "-loglevel", "error",
199 "-y", # Overwrite created output file
200 "-ss", str(start_seconds),
201 "-to", str(end_seconds),
202 "-i", input_file.name,
203 ]
205 if width and height:
206 cmd += ["-vf", f"scale={width}:{height}"]
208 cmd += [
209 "-c:v", video_codec,
210 "-c:a", audio_codec,
211 "-pix_fmt", pix_fmt,
212 output_file.name
213 ]
214 subprocess.run(cmd, check=True)
216 # Return trimmed video as bytes
217 output_file.seek(0)
218 return output_file.read()
219 finally:
220 # Clean up temporary files
221 for f in (input_file.name, output_file.name):
222 try:
223 os.remove(f)
224 except OSError:
225 pass
228def get_audio_duration(
229 audio_content: Union[str, bytes]
230) -> float:
231 """Get audio duration in seconds from binary content or base64 string."""
232 if isinstance(audio_content, bytes):
233 audio_bytes = audio_content
234 elif isinstance(audio_content, str):
235 audio_base64 = audio_content
236 audio_bytes = base64_to_binary(audio_base64)
237 else:
238 raise TypeError(f"Expected str|bytes for audio_content, got {type(audio_content)}")
239 if not isinstance(audio_bytes, bytes):
240 raise TypeError(f"Expected str|bytes for audio_content, got {type(audio_content)}")
242 import soundfile
244 audio_buffer = BytesIO(audio_bytes)
245 audio_data, sample_rate = soundfile.read(audio_buffer)
246 duration_seconds = len(audio_data) / sample_rate
247 return duration_seconds
250class AudioFileInfo(TypedDict):
251 num_frames: int
252 samplerate: int
253 channels: int
254 duration_seconds: float
257def get_audio_file_info(
258 audio_path: str
259) -> AudioFileInfo:
260 """Get audio file info from file path."""
261 if not isinstance(audio_path, str):
262 raise TypeError(f"Expected str for audio_path, got {type(audio_path)}")
263 if not os.path.exists(audio_path):
264 raise FileNotFoundError(f"Audio file does not exist: {audio_path}")
266 import soundfile
268 info = soundfile.info(audio_path)
269 return {
270 "num_frames": info.frames,
271 "samplerate": info.samplerate,
272 "channels": info.channels,
273 "duration_seconds": info.frames / info.samplerate
274 }
277class OverallInfo(TypedDict):
278 duration_seconds: float
279 bitrate: Optional[int]
280 num_bytes: int
283class VideoInfo(TypedDict, total=False):
284 codec: Optional[str]
285 pix_fmt: Optional[str]
286 bitrate: Optional[int]
287 num_frames: Optional[int]
288 fps: Optional[float]
289 duration_seconds: Optional[float]
290 width: Optional[int]
291 height: Optional[int]
292 aspect_ratio: Optional[float]
295class AudioInfo(TypedDict, total=False):
296 codec: Optional[str]
297 bitrate: Optional[int]
298 sample_rate: Optional[int]
299 channels: Optional[int]
300 duration_seconds: Optional[float]
303class VideoFileInfo(TypedDict, total=False):
304 overall: OverallInfo
305 video: VideoInfo
306 audio: AudioInfo
309def get_video_file_info(
310 video_content: Union[bytes, str]
311) -> VideoFileInfo:
312 """Get video file info from binary content or file path."""
313 if not video_content:
314 raise ValueError("Video input is empty.")
315 if isinstance(video_content, bytes):
316 with tempfile.NamedTemporaryFile(suffix=".mp4") as temp_video:
317 video_binary = video_content
318 temp_video.write(video_binary)
319 temp_video.flush()
320 video_path = temp_video.name
321 return get_video_file_info_path(video_path)
322 if isinstance(video_content, str):
323 video_path = video_content
324 return get_video_file_info_path(video_path)
325 raise TypeError(f"Expected bytes|str for video content, got {type(video_content)}")
328def get_video_file_info_path(
329 video_path: str
330) -> VideoFileInfo:
331 """Get video file info from file path."""
332 if not isinstance(video_path, str):
333 raise TypeError(f"Expected str for video_path, got {type(video_path)}")
334 if not os.path.exists(video_path):
335 raise FileNotFoundError(f"Video file does not exist: {video_path}")
336 cmd = [
337 "ffprobe",
338 "-v", "error",
339 "-count_frames",
340 "-show_entries",
341 (
342 "stream=codec_name,codec_type,avg_frame_rate,nb_frames,nb_read_frames,"
343 "duration,width,height,bit_rate,sample_rate,channels,pix_fmt"
344 ),
345 "-show_entries",
346 "format=duration,bit_rate,size",
347 "-of", "json",
348 video_path,
349 ]
350 result = subprocess.run(
351 cmd,
352 stdout=subprocess.PIPE,
353 stderr=subprocess.PIPE,
354 text=True)
355 info: Dict[str, Any] = json.loads(result.stdout)
357 if "format" not in info:
358 ffprobe_error = result.stderr.strip()
359 logging.error(f"No format found in ffprobe error: {ffprobe_error}")
360 raise ValueError("The video binary is corrupted or has an unsupported format")
362 overall_info: OverallInfo = {
363 "duration_seconds": float(info["format"].get("duration", 0)),
364 "bitrate": int(info["format"].get("bit_rate", 0)) if "bit_rate" in info["format"] else None,
365 "num_bytes": int(info["format"].get("size", 0))
366 }
367 file_info: VideoFileInfo = {
368 "overall": overall_info
369 }
371 if "streams" not in info:
372 logging.error(f"No streams found in ffprobe output: {result.stdout}")
374 # Video info
375 video_stream = next((s for s in info["streams"] if s["codec_type"] == "video"), None)
376 if video_stream:
377 fps_val: Optional[float] = None
378 fps = video_stream.get("avg_frame_rate", "0/0")
379 if fps != "0/0":
380 num, den = map(int, fps.split("/"))
381 fps_val = num / den if den else None
382 num_frames: Optional[int] = None
383 if "nb_frames" in video_stream and video_stream["nb_frames"].isdigit():
384 num_frames = int(video_stream["nb_frames"])
385 elif "nb_read_frames" in video_stream and video_stream["nb_read_frames"].isdigit():
386 num_frames = int(video_stream["nb_read_frames"])
388 video_info: VideoInfo = {
389 "codec": video_stream.get("codec_name"),
390 "pix_fmt": video_stream.get("pix_fmt"),
391 "bitrate": int(video_stream["bit_rate"]) if "bit_rate" in video_stream else None,
392 "num_frames": num_frames,
393 "fps": fps_val,
394 # "duration_seconds": float(video_stream.get("duration", 0)),
395 }
396 # Formats like MKV may not have duration in stream
397 if "duration" in video_stream and float(video_stream["duration"]) > 0:
398 video_info["duration_seconds"] = float(video_stream["duration"])
400 width = int(video_stream.get("width", 0))
401 height = int(video_stream.get("height", 0))
402 if width > 0 and height > 0:
403 video_info["width"] = width
404 video_info["height"] = height
405 video_info["aspect_ratio"] = width / height
407 file_info["video"] = video_info
409 # Audio info
410 audio_stream = next((s for s in info["streams"] if s["codec_type"] == "audio"), None)
411 if audio_stream:
412 audio_info: AudioInfo = {
413 "codec": audio_stream.get("codec_name"),
414 "bitrate": int(audio_stream["bit_rate"]) if "bit_rate" in audio_stream else None,
415 "sample_rate": int(audio_stream["sample_rate"]) if "sample_rate" in audio_stream else None,
416 "channels": int(audio_stream["channels"]) if "channels" in audio_stream else None,
417 }
418 # Formats like MKV may not have duration in stream
419 if "duration" in audio_stream and float(audio_stream["duration"]) > 0:
420 audio_info["duration_seconds"] = float(audio_stream["duration"])
421 file_info["audio"] = audio_info
423 return file_info
426class ImageFileInfo(TypedDict):
427 width: int
428 height: int
429 aspect_ratio: float
432def get_image_file_info(
433 image_path: str
434) -> ImageFileInfo:
435 if not isinstance(image_path, str):
436 raise TypeError(f"Expected str for image_path, got {type(image_path)}")
437 if not os.path.exists(image_path):
438 raise FileNotFoundError(f"Image file does not exist: {image_path}")
439 img = Image.open(image_path)
440 return {
441 "width": img.width,
442 "height": img.height,
443 "aspect_ratio": img.width / img.height,
444 }
447class TextFileInfo(TypedDict):
448 num_chars: int
449 num_words: int
450 num_lines: int
453def get_text_file_info(
454 text_path: str
455) -> TextFileInfo:
456 if not isinstance(text_path, str):
457 raise TypeError(f"Expected str for text_path, got {type(text_path)}")
458 if not os.path.exists(text_path):
459 raise FileNotFoundError(f"Text file does not exist: {text_path}")
460 with open(text_path, 'r', encoding="utf-8", errors="ignore") as f:
461 content = f.read()
462 lines = content.split('\n')
463 words = content.split()
464 return {
465 "num_chars": len(content),
466 "num_words": len(words),
467 "num_lines": len(lines),
468 }
471class TensorFileInfo(TypedDict):
472 dtype: str
473 shape: str
474 device: str
475 mean: str
476 min: str
477 max: str
478 numel: int
481def get_tensor_file_info(
482 tensor_path: str
483) -> TensorFileInfo:
484 if not isinstance(tensor_path, str):
485 raise TypeError(f"Expected str for tensor_path, got {type(tensor_path)}")
486 if not os.path.exists(tensor_path):
487 raise FileNotFoundError(f"Tensor file does not exist: {tensor_path}")
488 tensor = torch.load(tensor_path, weights_only=True)
489 return {
490 "dtype": str(tensor.dtype),
491 "shape": str(tensor.shape),
492 "device": str(tensor.device),
493 "mean": str(tensor.mean().item()),
494 "min": str(tensor.min().item()),
495 "max": str(tensor.max().item()),
496 "numel": tensor.numel(),
497 }
500async def base64_to_audio_file(
501 audio_base64: str,
502 audio_path: Optional[str] = None
503) -> str:
504 if not isinstance(audio_base64, str):
505 raise TypeError(f"Expected str for audio_base64, got {type(audio_base64)}")
506 if audio_path is None:
507 audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
508 return await save_base64_as_binary(
509 audio_path,
510 audio_base64)
513def empty_audio_file(
514 duration_seconds: float = 0.2,
515 sample_rate: int = 44100, # KHz
516 num_channels: int = 1,
517 sample_width: int = 2
518) -> str:
519 n_samples = int(sample_rate * duration_seconds)
520 silence = np.zeros(n_samples, dtype=np.int16) # 16-bit PCM
521 audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
522 with wave.open(audio_path, "w") as wav_file:
523 wav_file.setnchannels(num_channels)
524 wav_file.setsampwidth(sample_width)
525 wav_file.setframerate(sample_rate)
526 wav_file.writeframes(silence.tobytes())
527 return audio_path
530def fit_audio_to_duration(
531 input_path: str,
532 target_duration: float,
533 output_path: Optional[str] = None
534) -> str:
535 if not isinstance(input_path, str):
536 raise TypeError(f"Expected str for input_path, got {type(input_path)}")
537 if not os.path.exists(input_path):
538 raise FileNotFoundError(f"Input WAV file does not exist: {input_path}")
539 if not input_path.lower().endswith(".wav"):
540 raise ValueError(f"Input file must be a WAV file: {input_path}")
541 if target_duration <= 0:
542 raise ValueError("target_duration must be positive")
544 with wave.open(input_path, "rb") as wf:
545 params = wf.getparams()
546 framerate = wf.getframerate()
547 n_channels = wf.getnchannels()
548 sampwidth = wf.getsampwidth()
549 nframes = wf.getnframes()
550 frames = wf.readframes(nframes)
552 dtype = {1: np.int8, 2: np.int16, 4: np.int32}[sampwidth]
553 audio = np.frombuffer(frames, dtype=dtype)
555 if not output_path:
556 output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
558 target_samples = int(target_duration * framerate * n_channels)
560 if len(audio) > target_samples:
561 audio = audio[:target_samples] # Truncate
562 elif len(audio) < target_samples:
563 silence = np.zeros(target_samples - len(audio), dtype=dtype)
564 audio = np.concatenate([audio, silence]) # Pad with silence
566 with wave.open(output_path, "wb") as wf:
567 wf.setparams(params._replace(nframes=target_samples))
568 wf.writeframes(audio.tobytes())
569 return output_path
572@asynccontextmanager
573async def async_tempfile(
574 suffix: str = "",
575 delete: bool = True
576) -> AsyncIterator[AsyncBufferedIOBase]:
577 """Asynchronous context manager for a temporary file."""
578 fd, path = tempfile.mkstemp(suffix=suffix)
579 os.close(fd) # close fd so aiofiles can open
580 try:
581 file = await aiofiles.open(path, mode="w+b")
582 try:
583 yield file
584 finally:
585 await file.close()
586 finally:
587 if delete and await aiofiles.os.path.exists(path):
588 await aiofiles.os.remove(path)
591async def get_video_frame(
592 video_binary: bytes,
593 frame_index: int = 0 # Get first frame by default
594) -> Image.Image:
595 """Extracts a specific frame from a video binary using imageio-ffmpeg."""
596 if not isinstance(video_binary, bytes):
597 raise TypeError(f"Expected bytes for video binary, got {type(video_binary)}")
598 if len(video_binary) == 0:
599 raise ValueError("Video binary is empty.")
600 if frame_index < 0:
601 raise ValueError("Frame index must be non-negative.")
603 video_frames = await get_video_frames(video_binary)
604 if frame_index >= len(video_frames):
605 raise ValueError(f"Frame index {frame_index} exceeds number of frames {len(video_frames)}.")
606 return video_frames[frame_index]
609def get_video_frames_at_fps(
610 video_frames: List[Image.Image],
611 src_fps: float = 30.0, # Hunyuan FramePack
612 dst_fps: float = 23.0, # Fantasy Talking
613) -> List[Image.Image]:
614 """
615 Adjusts the frame rate of a list of video frames by either up-sampling or down-sampling.
616 Args:
617 video_frames (List[Image.Image]): List of input frames.
618 src_fps (int): Original frames per second.
619 dst_fps (int): Desired frames per second.
620 Returns:
621 List[Image.Image]: New list of frames adjusted to the target FPS.
622 """
623 if not video_frames:
624 return []
625 if src_fps <= 0 or dst_fps <= 0:
626 raise ValueError("FPS must be a positive integer.")
627 if src_fps == dst_fps:
628 return video_frames # No FPS change needed
630 num_src_frames = len(video_frames)
631 duration_secs = 1.0 * num_src_frames / src_fps
632 # target_frame_count = int(math.ceil(duration_secs * dst_fps))
633 target_frame_count = int(round(duration_secs * dst_fps))
635 new_frames = []
636 for ix in range(target_frame_count):
637 t = 1.0 * ix / dst_fps # target time in seconds
638 src_index = min(int(t * src_fps), num_src_frames - 1)
639 video_frame = video_frames[src_index]
640 new_frames.append(video_frame)
642 return new_frames
645def get_ffmpeg_version() -> str:
646 # Use dpkg -s ffmpeg to get version on Ubuntu
647 # Example: Version: 7:4.2.7-0ubuntu0.1
648 if os.path.exists("/usr/bin/dpkg"):
649 result = subprocess.run(
650 ["dpkg", "-s", "ffmpeg"],
651 stdout=subprocess.PIPE,
652 stderr=subprocess.PIPE,
653 text=True)
654 for line in result.stdout.split('\n'):
655 if line.startswith("Version:"):
656 return line.split("Version:")[1].strip()
658 # Otherwise use ffmpeg -version
659 result = subprocess.run(
660 ["ffmpeg", "-version"],
661 stdout=subprocess.PIPE,
662 stderr=subprocess.PIPE,
663 text=True)
664 first_line = result.stdout.split('\n')[0]
665 # Extract from line: ffmpeg version 6.1.1-3ubuntu5 Copyright (c) 2000-2023 the FFmpeg developers
666 if not first_line.startswith("ffmpeg version"):
667 raise ValueError(f"Unexpected ffmpeg version output: {first_line}")
668 version = first_line.split(" ")[2].strip()
669 return version
672def change_video_fps(
673 video_binary: bytes,
674 new_fps: float,
675 video_codec: str = DEFAULT_VIDEO_CODEC,
676 audio_codec: str = DEFAULT_AUDIO_CODEC,
677 pix_fmt: str = DEFAULT_PIX_FMT,
678) -> bytes:
679 if not isinstance(video_binary, bytes):
680 raise TypeError(f"Expected bytes for video binary, got {type(video_binary)}")
681 if new_fps <= 0:
682 raise ValueError("FPS must be a positive integer.")
684 with tempfile.NamedTemporaryFile(suffix=".mp4", delete=True) as input_file, \
685 tempfile.NamedTemporaryFile(suffix=".mp4", delete=True) as output_file:
686 input_file.write(video_binary)
687 input_file.flush()
689 subprocess.run([
690 "ffmpeg",
691 "-hide_banner",
692 "-loglevel", "error",
693 "-y", # Overwrite created output file
694 "-i", input_file.name,
695 "-filter:v", f"fps={new_fps}",
696 "-c:v", video_codec,
697 "-c:a", audio_codec,
698 "-pix_fmt", pix_fmt,
699 output_file.name
700 ], check=True)
702 output_file.seek(0)
703 return output_file.read()
706def get_video_size(
707 video_frames: List[Image.Image]
708) -> int:
709 """
710 Get approximate size in bytes of a list of video frames.
711 This is a rough estimate based on width, height, and RGB channels.
712 """
713 if not video_frames:
714 return 0
715 ret = 0
716 NUM_RGB_CHANNELS = 3
717 for frame in video_frames:
718 if frame is not None:
719 ret += frame.width * frame.height * NUM_RGB_CHANNELS
720 return ret
723async def get_video_frames(
724 video_binary: Optional[bytes],
725 extend_last_frame: bool = False
726) -> List[Image.Image]:
727 """
728 Extracts frames from a video binary using imageio-ffmpeg.
729 This is slow.
730 """
731 if not video_binary:
732 return []
733 if not isinstance(video_binary, bytes):
734 raise TypeError(f"Expected bytes for video binary, got {type(video_binary)}")
735 if len(video_binary) == 0:
736 raise ValueError("Video binary is empty.")
738 NUM_RGB_CHANNELS = 3
740 async with async_tempfile(suffix=".mp4") as temp_video:
741 await temp_video.write(video_binary)
742 await temp_video.flush()
743 video_path = temp_video.name
745 reader = ffmpeg.read_frames(video_path, pix_fmt=PIX_FMT_RGB24)
746 meta = reader.__next__()
747 # 'ffmpeg_version', 'codec', 'pix_fmt', 'fps', 'source_size', 'size', 'rotate', 'duration'
748 width, height = meta["size"]
750 frames: List[Image.Image] = []
751 for frame_bytes in reader:
752 frame_array = np.frombuffer(frame_bytes, dtype=np.uint8)
753 frame_array = frame_array.reshape((height, width, NUM_RGB_CHANNELS))
754 frame_img_pil = Image.fromarray(frame_array)
755 frames.append(frame_img_pil)
757 # Extend last frame if needed
758 # expected_num_frames = int(math.ceil(meta["fps"] * meta["duration"]))
759 expected_num_frames = int(round(meta["fps"] * meta["duration"]))
760 if len(frames) < expected_num_frames:
761 logging.warning(
762 f"Video has {len(frames)} frames but expected {expected_num_frames} "
763 f"({meta['fps']}x{meta['duration']}). "
764 "Exending last frame.")
765 if extend_last_frame:
766 frames.extend([frames[-1]] * (expected_num_frames - len(frames)))
768 return frames
771async def get_video_duration(
772 video_content: Union[bytes, str]
773) -> float:
774 """Get video duration in seconds from binary content or file path."""
775 video_file_info = get_video_file_info(video_content)
776 video_info = video_file_info.get("video", {})
777 result = video_info.get("duration_seconds", -1.0)
778 return result if result is not None else -1.0
781def get_video_fps(
782 video_content: Union[bytes, str]
783) -> float:
784 """Get video frames per second (FPS) from binary content or file path."""
785 video_file_info = get_video_file_info(video_content)
786 video_info = video_file_info.get("video", {})
787 result = video_info.get("fps", -1.0)
788 return result if result is not None else -1.0
791async def get_video_num_frames(
792 video_content: Union[bytes, str]
793) -> int:
794 """Get number of video frames from binary content or file path."""
795 video_file_info = get_video_file_info(video_content)
796 video_info = video_file_info.get("video", {})
797 result = video_info.get("num_frames", -1)
798 return result if result is not None else -1
801def save_audio(
802 audio: torch.Tensor,
803 audio_path: str,
804 sample_rate: int = 24000
805) -> str:
806 if not isinstance(audio, torch.Tensor):
807 raise TypeError(f"Expected torch.Tensor for audio, got {type(audio)}")
808 if not isinstance(audio_path, str):
809 raise TypeError(f"Expected str for audio_path, got {type(audio_path)}")
810 import soundfile
811 soundfile.write(
812 audio_path,
813 audio,
814 sample_rate,
815 format="WAV",
816 subtype="PCM_16")
817 return audio_path
820def save_audio_np(
821 wav: np.ndarray,
822 path: str,
823 sample_rate: Optional[int] = None
824) -> None:
825 import scipy
827 # https://github.com/coqui-ai/TTS/blob/main/TTS/utils/audio/numpy_transforms.py#L430
828 FREQ_32KHz = (32 * 1024) - 1
829 wav_norm = wav * (FREQ_32KHz / max(0.01, np.max(np.abs(wav))))
831 wav_norm = wav_norm.astype(np.int16) # PCM16
832 scipy.io.wavfile.write(path, sample_rate, wav_norm)
835async def save_video_frames(
836 video_frames: Union[Sequence[Union[Image.Image, np.ndarray, torch.Tensor]], np.ndarray],
837 out_video_path: Optional[str] = None,
838 fps: float = 30.0,
839 video_codec: str = DEFAULT_VIDEO_CODEC,
840 pix_fmt: str = DEFAULT_PIX_FMT,
841 output_format: str = "mp4", # "mkv"
842) -> str:
843 """Save video frames and optional audio to a file asynchronously using ffmpeg."""
844 if not video_frames:
845 raise ValueError("No video frames provided.")
846 if not isinstance(video_frames, (list, np.ndarray)):
847 raise TypeError(f"Expected list or np.ndarray for video_frames, got {type(video_frames)}")
849 frame0 = video_frames[0]
850 if isinstance(frame0, np.ndarray):
851 width, height = frame0.shape[1], frame0.shape[0]
852 else:
853 width, height = frame0.size
855 if not out_video_path:
856 out_video_path = tempfile.NamedTemporaryFile(suffix=f".{output_format}", delete=False).name
858 num_frames = len(video_frames)
859 duration = float(num_frames) / float(fps)
861 cmd = [
862 "ffmpeg",
863 "-hide_banner",
864 "-loglevel", "error",
865 "-y", # Overwrite created output file
866 "-f", "rawvideo",
867 "-vcodec", "rawvideo",
868 "-pix_fmt", "rgb24",
869 "-s", f"{width}x{height}",
870 "-r", str(fps),
871 "-i", "-", # input frames from stdin
872 "-frames:v", str(num_frames),
873 "-t", str(duration),
874 "-c:v", video_codec,
875 "-pix_fmt", pix_fmt,
876 "-preset", "fast", # "slow",
877 out_video_path
878 ]
879 cmd = [c for c in cmd if c]
881 proc = await asyncio.create_subprocess_exec(
882 *cmd,
883 stdin=asyncio.subprocess.PIPE,
884 stdout=asyncio.subprocess.DEVNULL,
885 stderr=asyncio.subprocess.DEVNULL)
886 if proc.stdin is None:
887 raise RuntimeError("Failed to open ffmpeg stdin for writing frames.")
889 for frame in video_frames[:num_frames]:
890 if isinstance(frame, Image.Image):
891 frame_array = np.array(frame.convert("RGB"))
892 elif isinstance(frame, np.ndarray):
893 frame_array = frame
894 elif isinstance(frame, torch.Tensor):
895 frame_array = frame.cpu().numpy()
896 else:
897 frame_array = frame
898 frame_bytes = frame_array.tobytes()
899 proc.stdin.write(frame_bytes)
900 proc.stdin.close()
901 await proc.wait()
903 return out_video_path
906async def extract_audio_from_video(
907 video_path: str,
908 out_audio_path: Optional[str] = None,
909) -> str:
910 """
911 Extract audio from video.
912 It is forced to WAV format for compatibility.
913 44100 Hz, 16-bit PCM, stereo.
914 """
915 if not os.path.exists(video_path):
916 raise FileNotFoundError(f"Video file does not exist: {video_path}")
918 if not out_audio_path:
919 out_audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
921 cmd = [
922 "ffmpeg",
923 "-hide_banner",
924 "-loglevel", "error",
925 "-y", # Overwrite created output file
926 "-i", video_path,
927 "-vn", # No video
928 "-acodec", "pcm_s16le", # Convert audio to WAV
929 "-ar", "44100",
930 "-ac", "2",
931 out_audio_path
932 ]
933 cmd = [c for c in cmd if c]
935 proc = await asyncio.create_subprocess_exec(
936 *cmd,
937 stdout=asyncio.subprocess.DEVNULL,
938 stderr=asyncio.subprocess.DEVNULL)
939 await proc.wait()
941 return out_audio_path
944async def save_video_audio(
945 video_content: Union[bytes, str, List[Image.Image], np.ndarray],
946 audio_path: str,
947 fps: float = 30.0,
948 out_video_path: Optional[str] = None,
949 video_codec: str = DEFAULT_VIDEO_CODEC,
950 audio_codec: str = DEFAULT_AUDIO_CODEC,
951 output_format: str = "mp4", # "mkv"
952) -> str:
953 """Merge audio and video."""
954 if isinstance(video_content, (list, np.ndarray)):
955 video_frames = video_content
956 video_path = await save_video_frames(
957 video_frames=video_frames,
958 fps=fps,
959 video_codec=video_codec,
960 out_video_path=None,
961 pix_fmt=DEFAULT_PIX_FMT,
962 output_format=output_format)
963 elif isinstance(video_content, bytes):
964 video_binary = video_content
965 with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
966 temp_video.write(video_binary)
967 temp_video.flush()
968 video_path = temp_video.name
969 elif isinstance(video_content, str):
970 video_path = video_content
971 else:
972 raise TypeError(f"Expected bytes|str|list|np.ndarray for video_content, got {type(video_content)}")
974 if not await aiofiles.os.path.exists(video_path):
975 raise FileNotFoundError(f"Video file does not exist: {video_path}")
976 if not await aiofiles.os.path.exists(audio_path):
977 raise FileNotFoundError(f"Audio file does not exist: {audio_path}")
979 if not out_video_path:
980 out_video_path = tempfile.NamedTemporaryFile(suffix=f".{output_format}", delete=False).name
982 video_info = get_video_file_info_path(video_path)
983 video_duration = video_info["video"]["duration_seconds"]
985 cmd = [
986 "ffmpeg",
987 "-hide_banner",
988 "-loglevel", "error",
989 "-y", # Overwrite created output file
990 "-i", video_path,
991 "-i", audio_path,
992 "-c:v", "copy",
993 "-c:a", audio_codec,
994 ]
995 cmd += [
996 "-af", f"apad,atrim=0:{video_duration}", # Align with the video
997 "-shortest",
998 ]
999 cmd += [
1000 out_video_path
1001 ]
1002 cmd = [c for c in cmd if c]
1004 proc = await asyncio.create_subprocess_exec(
1005 *cmd,
1006 stdout=asyncio.subprocess.DEVNULL,
1007 stderr=asyncio.subprocess.DEVNULL)
1008 await proc.wait()
1010 return out_video_path
1013def save_diffusers_video(
1014 video_frames: List[Image.Image],
1015 out_video_path: Optional[str] = None,
1016 fps: float = 30.0,
1017 quality: float = 5.0,
1018 bitrate: Optional[int] = None,
1019 macro_block_size: Optional[int] = 16,
1020) -> str:
1021 """
1022 Save a list of PIL Image frames as a video using the diffusers library.
1023 https://github.com/huggingface/diffusers/blob/main/src/diffusers/utils/export_utils.py#L141
1024 """
1025 if not video_frames:
1026 raise ValueError("No video frames provided.")
1027 if out_video_path is None:
1028 out_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
1029 with imageio.get_writer(
1030 out_video_path, fps=fps, quality=quality, bitrate=bitrate,
1031 macro_block_size=macro_block_size
1032 ) as writer:
1033 for frame in video_frames:
1034 frame_np = np.array(frame)
1035 writer.append_data(frame_np) # type: ignore[attr-defined]
1036 return out_video_path
1039def save_bcthw_as_mp4(
1040 x: torch.Tensor,
1041 output_filename: str,
1042 fps: float = 30.0,
1043 video_codec: str = DEFAULT_VIDEO_CODEC,
1044) -> str:
1045 import einops
1047 # Version with no torchvision warning from
1048 # https://github.com/lllyasviel/FramePack/blob/main/diffusers_helper/utils.py#L266
1049 b, c, t, h, w = x.shape
1050 per_row = b
1051 for p in [6, 5, 4, 3, 2]:
1052 if b % p == 0:
1053 per_row = p
1054 break
1055 x = (torch.clamp(x.float(), -1., 1.) + 1) * 255 / 2
1056 x = x.detach().cpu().to(torch.uint8)
1057 x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
1058 # torchvision.io.write_video(output_filename, x, fps=fps, video_codec=codec, options={'crf': str(int(crf))})
1059 with imageio.get_writer(output_filename, fps=fps, codec=video_codec) as writer:
1060 for frame in x.numpy():
1061 writer.append_data(frame) # type: ignore[attr-defined]
1062 return output_filename
1065def get_aligned_duration(
1066 duration: float,
1067 fps: float = 30.0,
1068 vae: int = 4
1069) -> float:
1070 """
1071 Aligns the duration to be compatible with diffusion models that use a VAE and FPS.
1072 This ensures that the number of frames is a multiple of the VAE's downsampling factor.
1073 VAE does: 1+4n
1074 """
1075 if duration == 0:
1076 return 0.0
1077 if duration < 0:
1078 raise ValueError("Duration must be positive.")
1079 if fps <= 0:
1080 raise ValueError("FPS must be positive.")
1081 if vae <= 0:
1082 raise ValueError("VAE must be positive.")
1083 num_frames = duration * fps
1084 lat_num_frames = int(math.ceil((num_frames - 1) / vae)) + 1
1085 aligned_num_frames = (lat_num_frames - 1) * vae + 1
1086 aligned_duration = aligned_num_frames / fps
1087 return aligned_duration
1090async def concatenate_videos(
1091 video_inputs: Union[List[bytes], List[str]],
1092 fast_copy: bool = True,
1093 video_codec: str = DEFAULT_VIDEO_CODEC,
1094 audio_codec: str = DEFAULT_AUDIO_CODEC,
1095 pix_fmt: str = DEFAULT_PIX_FMT,
1096) -> bytes:
1097 """
1098 Concatenate multiple videos (given as raw bytes) into a single video.
1099 Returns the final video as bytes.
1100 """
1101 if not video_inputs:
1102 raise ValueError("No videos provided")
1104 with tempfile.TemporaryDirectory() as tmpdir:
1105 input_paths = []
1106 for i, video_input in enumerate(video_inputs):
1107 if isinstance(video_input, str):
1108 input_path = video_input
1109 if not await aiofiles.os.path.exists(input_path):
1110 raise FileNotFoundError(f"Video file does not exist: {input_path}")
1111 elif isinstance(video_input, bytes):
1112 video_bytes = video_input
1113 input_path = f"{tmpdir}/input_{i}.mp4"
1114 async with aiofiles.open(input_path, "wb") as file:
1115 await file.write(video_bytes)
1116 else:
1117 raise TypeError(f"Video input {i} is not bytes or str. Found: {type(video_input)}")
1118 input_paths.append(input_path)
1120 list_file = f"{tmpdir}/videos.txt"
1121 async with aiofiles.open(list_file, "w") as f:
1122 for path in input_paths:
1123 await f.write(f"file '{path}'\n")
1125 output_path = f"{tmpdir}/output.mp4"
1127 # Build ffmpeg command with stream copy (fast)
1128 cmd = [
1129 "ffmpeg",
1130 "-hide_banner",
1131 "-loglevel", "error",
1132 "-y", # Overwrite created output file
1133 "-f", "concat",
1134 "-safe", "0",
1135 "-i", str(list_file),
1136 ]
1137 if fast_copy:
1138 # This is fast but may leave the audio and video out of sync
1139 cmd += [
1140 "-c", "copy",
1141 ]
1142 else:
1143 cmd += [
1144 "-c:v", video_codec,
1145 "-c:a", audio_codec,
1146 "-pix_fmt", pix_fmt,
1147 ]
1148 cmd += [
1149 str(output_path)
1150 ]
1152 proc = await asyncio.create_subprocess_exec(
1153 *cmd,
1154 stdout=asyncio.subprocess.PIPE,
1155 stderr=asyncio.subprocess.PIPE
1156 )
1157 _, stderr = await proc.communicate()
1159 if proc.returncode != 0:
1160 err_msg = stderr.decode()
1161 if "Invalid data found when processing input" in err_msg:
1162 raise ValueError("One of the inputs is corrupted or has an unsupported format.")
1163 raise RuntimeError(f"ffmpeg failed: {err_msg}")
1165 concatenated_binary = await read_file_bytes(output_path)
1167 return concatenated_binary
1170DEFAULT_FONT = "/usr/share/fonts/truetype/noto/NotoColorEmoji.ttf"
1171DEFAULT_FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
1174def get_font_size(width: int, height: int) -> int:
1175 """Get the font size based on the frame dimensions."""
1176 if width >= 1280 or height >= 800:
1177 return 32
1178 if width >= 640 or height >= 400:
1179 return 24
1180 return 16
1183def get_font(
1184 font_path: str = DEFAULT_FONT,
1185 font_size: int = 32
1186) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
1187 """Load a truetype font, or default if not found."""
1188 font: ImageFont.FreeTypeFont | ImageFont.ImageFont
1189 try:
1190 font = ImageFont.truetype(font_path, font_size)
1191 except OSError:
1192 # apt-get install fonts-noto-color-emoji
1193 logging.warning(f"Could not load font at {font_path}. Using default font.")
1194 font = ImageFont.load_default()
1195 return font
1198def get_frame_with_text(
1199 width: int,
1200 height: int,
1201 text: str,
1202 output_type: str = "numpy",
1203 font_color: str = "white",
1204 font_size: int = 32,
1205 background_color: str = "black",
1206) -> Union[np.ndarray, torch.Tensor, Image.Image]:
1207 """
1208 Create a blank frame with multi-line text (including emojis) centered.
1209 Frame shape: [h, w, RGB].
1210 """
1211 # Convert to PIL for better text rendering
1212 img_pil = Image.new("RGB", (width, height), color=background_color)
1213 draw = ImageDraw.Draw(img_pil)
1215 font = get_font(font_size=font_size)
1217 # Handle multi-line text
1218 lines = text.split("\n")
1219 line_sizes = []
1220 for line in lines:
1221 bbox = draw.textbbox((0, 0), line, font=font)
1222 w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
1223 line_sizes.append((w, h))
1225 total_height = sum(h for _, h in line_sizes) + (len(lines) - 1) * 5 # line spacing
1227 # Start y so that block is vertically centered
1228 y = (height - total_height) // 2
1229 for line, (tw, th) in zip(lines, line_sizes):
1230 x = (width - tw) // 2
1231 draw.text((x, y), line, font=font, fill=font_color)
1232 y += th + 5
1234 # Convert back to desired format
1235 if output_type == "torch":
1236 return torch.from_numpy(np.array(img_pil)).permute(2, 0, 1)
1237 if output_type == "pil":
1238 return img_pil.convert("RGB")
1239 return np.array(img_pil) # numpy
1242async def get_video_with_text(
1243 width: int,
1244 height: int,
1245 text: str,
1246 duration_seconds: float = 2.0,
1247 fps: float = 30.0,
1248 font_size: int = 32,
1249 font_color: str = "white",
1250) -> bytes:
1251 """
1252 Create a video with a single frame containing text.
1253 Returns the video as bytes.
1254 """
1255 video_frame = get_frame_with_text(
1256 width, height,
1257 text,
1258 font_size=font_size,
1259 font_color=font_color)
1260 # num_frames = int(math.ceil(duration_seconds * fps))
1261 num_frames = int(round(duration_seconds * fps))
1262 video_frames = [video_frame] * num_frames
1263 video_path = await save_video_frames(
1264 video_frames=video_frames,
1265 fps=fps)
1266 video_bytes = await read_file_bytes(video_path)
1267 if await aiofiles.os.path.exists(video_path):
1268 await aiofiles.os.unlink(video_path)
1269 return video_bytes
1272def get_text_position_coordinates(
1273 position: str,
1274 img_width: int,
1275 img_height: int,
1276 text_width: int,
1277 text_height: int,
1278 margin: int = 10
1279) -> Tuple[int, int]:
1280 """Get the (x, y) coordinates for placing text on an image based on the specified position."""
1281 if position == "top-left":
1282 x, y = margin, margin
1283 elif position == "top-right":
1284 x = img_width - text_width - margin
1285 y = margin
1286 elif position == "bottom-left":
1287 x = margin
1288 y = img_height - text_height - margin
1289 elif position == "bottom-right":
1290 x = img_width - text_width - margin
1291 y = img_height - text_height - margin
1292 elif position == "bottom-center":
1293 x = (img_width - text_width) // 2
1294 y = img_height - text_height - margin
1295 elif position == "center":
1296 x = (img_width - text_width) // 2
1297 y = (img_height - text_height) // 2
1298 else:
1299 raise ValueError(
1300 f"Invalid position: {position}. "
1301 f"Must be one of 'top-left', 'top-right', 'bottom-left', 'bottom-right', 'center'.")
1302 return x, y
1305def split_text_lines(
1306 input_text: str,
1307 max_line_length: int
1308) -> List[str]:
1309 """Split text into multiple lines with a maximum line length."""
1310 return textwrap.wrap(
1311 input_text,
1312 width=max_line_length)
1315def add_text_to_frame(
1316 frame: Union[np.ndarray, torch.Tensor, Image.Image],
1317 text: str,
1318 font_size: int = 32,
1319 font_color: str = "white",
1320 position: Union[str, Tuple[int, int]] = "top-left",
1321) -> Union[np.ndarray, torch.Tensor, Image.Image]:
1322 """
1323 Add text to the frame.
1324 Frame should be [h, w, RGB].
1325 """
1326 if not text:
1327 logging.warning("No text provided to add to frame.")
1328 return frame
1329 if "\n" in text:
1330 logging.warning("Text contains newlines.")
1331 text = text.replace("\n", " ")
1333 output_type = "numpy"
1334 if isinstance(frame, torch.Tensor):
1335 output_type = "torch"
1336 frame_torch: torch.Tensor = frame
1337 frame_np = frame_torch.numpy()
1338 frame = frame_np.astype(np.uint8)
1339 if isinstance(frame, Image.Image):
1340 output_type = "pil"
1341 frame = np.array(frame.convert("RGB")).astype(np.uint8)
1342 if not isinstance(frame, np.ndarray):
1343 raise TypeError(f"Unsupported frame type: {type(frame)}.")
1345 img_pil = Image.fromarray(frame)
1346 draw = ImageDraw.Draw(img_pil)
1348 font = get_font(font_size=font_size)
1350 if isinstance(position, str):
1351 x, y = get_text_position_coordinates(
1352 position,
1353 img_width=img_pil.width,
1354 img_height=img_pil.height,
1355 text_width=int(draw.textlength(text, font=font)),
1356 text_height=font_size,
1357 margin=10)
1358 elif isinstance(position, tuple) and len(position) == 2:
1359 x, y = position
1360 else:
1361 raise ValueError("Position must be a string or (x, y) coordinates.")
1362 draw.text((x, y), text, font=font, fill=font_color)
1364 ret_frame = np.array(img_pil).astype(np.uint8)
1365 if output_type == "torch":
1366 return torch.from_numpy(ret_frame).permute(2, 0, 1)
1367 if output_type == "pil":
1368 return Image.fromarray(ret_frame).convert("RGB")
1369 return np.array(ret_frame) # np.ndarray of shape (height, width, 3) with RGB values