|
3 | 3 | import struct
|
4 | 4 | import tempfile
|
5 | 5 | from contextlib import contextmanager
|
6 |
| -from typing import List, Union |
| 6 | +from typing import List, Optional, Union |
7 | 7 |
|
8 | 8 | import numpy as np
|
9 | 9 | import PIL.Image
|
@@ -139,8 +139,31 @@ def _legacy_export_to_video(
|
139 | 139 |
|
140 | 140 |
|
141 | 141 | def export_to_video(
|
142 |
| - video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 10 |
| 142 | + video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], |
| 143 | + output_video_path: str = None, |
| 144 | + fps: int = 10, |
| 145 | + quality: float = 5.0, |
| 146 | + bitrate: Optional[int] = None, |
| 147 | + macro_block_size: Optional[int] = 16, |
143 | 148 | ) -> str:
|
| 149 | + """ |
| 150 | + quality: |
| 151 | + Video output quality. Default is 5. Uses variable bit rate. Highest quality is 10, lowest is 0. Set to None to |
| 152 | + prevent variable bitrate flags to FFMPEG so you can manually specify them using output_params instead. |
| 153 | + Specifying a fixed bitrate using `bitrate` disables this parameter. |
| 154 | +
|
| 155 | + bitrate: |
| 156 | + Set a constant bitrate for the video encoding. Default is None causing `quality` parameter to be used instead. |
| 157 | + Better quality videos with smaller file sizes will result from using the `quality` variable bitrate parameter |
| 158 | + rather than specifiying a fixed bitrate with this parameter. |
| 159 | +
|
| 160 | + macro_block_size: |
| 161 | + Size constraint for video. Width and height, must be divisible by this number. If not divisible by this number |
| 162 | + imageio will tell ffmpeg to scale the image up to the next closest size divisible by this number. Most codecs |
| 163 | + are compatible with a macroblock size of 16 (default), some can go smaller (4, 8). To disable this automatic |
| 164 | + feature set it to None or 1, however be warned many players can't decode videos that are odd in size and some |
| 165 | + codecs will produce poor results or fail. See https://en.wikipedia.org/wiki/Macroblock. |
| 166 | + """ |
144 | 167 | # TODO: Dhruv. Remove by Diffusers release 0.33.0
|
145 | 168 | # Added to prevent breaking existing code
|
146 | 169 | if not is_imageio_available():
|
@@ -177,7 +200,9 @@ def export_to_video(
|
177 | 200 | elif isinstance(video_frames[0], PIL.Image.Image):
|
178 | 201 | video_frames = [np.array(frame) for frame in video_frames]
|
179 | 202 |
|
180 |
| - with imageio.get_writer(output_video_path, fps=fps) as writer: |
| 203 | + with imageio.get_writer( |
| 204 | + output_video_path, fps=fps, quality=quality, bitrate=bitrate, macro_block_size=macro_block_size |
| 205 | + ) as writer: |
181 | 206 | for frame in video_frames:
|
182 | 207 | writer.append_data(frame)
|
183 | 208 |
|
|
0 commit comments