Skip to content

Commit 3e46043

Browse files
authored
Small improvements for video loading (#9183)
* update * update
1 parent 1a92bc0 commit 3e46043

File tree

1 file changed

+14
-3
lines changed

1 file changed

+14
-3
lines changed

src/diffusers/utils/loading_utils.py

+14-3
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import os
22
import tempfile
33
from typing import Callable, List, Optional, Union
4+
from urllib.parse import unquote, urlparse
45

56
import PIL.Image
67
import PIL.ImageOps
@@ -80,12 +81,22 @@ def load_video(
8081
)
8182

8283
if is_url:
83-
video_data = requests.get(video, stream=True).raw
84-
suffix = os.path.splitext(video)[1] or ".mp4"
84+
response = requests.get(video, stream=True)
85+
if response.status_code != 200:
86+
raise ValueError(f"Failed to download video. Status code: {response.status_code}")
87+
88+
parsed_url = urlparse(video)
89+
file_name = os.path.basename(unquote(parsed_url.path))
90+
91+
suffix = os.path.splitext(file_name)[1] or ".mp4"
8592
video_path = tempfile.NamedTemporaryFile(suffix=suffix, delete=False).name
93+
8694
was_tempfile_created = True
95+
96+
video_data = response.iter_content(chunk_size=8192)
8797
with open(video_path, "wb") as f:
88-
f.write(video_data.read())
98+
for chunk in video_data:
99+
f.write(chunk)
89100

90101
video = video_path
91102

0 commit comments

Comments
 (0)