|  | 
| 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