Skip to content

Commit 0ab8fe4

Browse files
authored
Quality options in export_to_video (#11090)
* Quality options in `export_to_video` * make style
1 parent 3be6706 commit 0ab8fe4

File tree

1 file changed

+28
-3
lines changed

1 file changed

+28
-3
lines changed

src/diffusers/utils/export_utils.py

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import struct
44
import tempfile
55
from contextlib import contextmanager
6-
from typing import List, Union
6+
from typing import List, Optional, Union
77

88
import numpy as np
99
import PIL.Image
@@ -139,8 +139,31 @@ def _legacy_export_to_video(
139139

140140

141141
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,
143148
) -> 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+
"""
144167
# TODO: Dhruv. Remove by Diffusers release 0.33.0
145168
# Added to prevent breaking existing code
146169
if not is_imageio_available():
@@ -177,7 +200,9 @@ def export_to_video(
177200
elif isinstance(video_frames[0], PIL.Image.Image):
178201
video_frames = [np.array(frame) for frame in video_frames]
179202

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:
181206
for frame in video_frames:
182207
writer.append_data(frame)
183208

0 commit comments

Comments
 (0)