From cabe0f60c941dd45f7f9948f878265e680f84fda Mon Sep 17 00:00:00 2001 From: Srijan Roy Date: Fri, 9 Oct 2026 09:51:18 +0530 Subject: [PATCH] Fix temporary file cleanup in load_video --- src/diffusers/utils/loading_utils.py | 109 ++++++++++++++------------- tests/others/test_loading_utils.py | 70 +++++++++++++++++ 2 files changed, 126 insertions(+), 53 deletions(-) create mode 100644 tests/others/test_loading_utils.py diff --git a/src/diffusers/utils/loading_utils.py b/src/diffusers/utils/loading_utils.py index ba528c7f857c..1ce2f5b55087 100644 --- a/src/diffusers/utils/loading_utils.py +++ b/src/diffusers/utils/loading_utils.py @@ -85,61 +85,64 @@ def load_video( f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {video} is not a valid path." ) - if is_url: - response = requests.get(video, stream=True) - if response.status_code != 200: - raise ValueError(f"Failed to download video. Status code: {response.status_code}") - - parsed_url = urlparse(video) - file_name = os.path.basename(unquote(parsed_url.path)) - - suffix = os.path.splitext(file_name)[1] or ".mp4" - video_path = tempfile.NamedTemporaryFile(suffix=suffix, delete=False).name - - was_tempfile_created = True - - video_data = response.iter_content(chunk_size=8192) - with open(video_path, "wb") as f: - for chunk in video_data: - f.write(chunk) - - video = video_path - - pil_images = [] - fps = None - if video.endswith(".gif"): - gif = PIL.Image.open(video) - # Milliseconds this frame is displayed for; GIFs are not obliged to record it. - frame_duration = gif.info.get("duration") - fps = 1000 / frame_duration if frame_duration else None - try: - while True: - pil_images.append(gif.copy()) - gif.seek(gif.tell() + 1) - except EOFError: - pass + try: + if is_url: + response = requests.get(video, stream=True) + if response.status_code != 200: + raise ValueError(f"Failed to download video. Status code: {response.status_code}") + + parsed_url = urlparse(video) + file_name = os.path.basename(unquote(parsed_url.path)) + + suffix = os.path.splitext(file_name)[1] or ".mp4" + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp: + video_path = tmp.name + + was_tempfile_created = True + + video_data = response.iter_content(chunk_size=8192) + with open(video_path, "wb") as f: + for chunk in video_data: + f.write(chunk) + + video = video_path + + pil_images = [] + fps = None + if video.endswith(".gif"): + with PIL.Image.open(video) as gif: + # Milliseconds this frame is displayed for; GIFs are not obliged to record it. + frame_duration = gif.info.get("duration") + fps = 1000 / frame_duration if frame_duration else None + try: + while True: + pil_images.append(gif.copy()) + gif.seek(gif.tell() + 1) + except EOFError: + pass - else: - if is_imageio_available(): - import imageio else: - raise ImportError(BACKENDS_MAPPING["imageio"][1].format("load_video")) - - try: - imageio.plugins.ffmpeg.get_exe() - except AttributeError: - raise AttributeError( - "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg" - ) - - with imageio.get_reader(video) as reader: - fps = reader.get_meta_data().get("fps") - # Read all frames - for frame in reader: - pil_images.append(PIL.Image.fromarray(frame)) - - if was_tempfile_created: - os.remove(video_path) + if is_imageio_available(): + import imageio + else: + raise ImportError(BACKENDS_MAPPING["imageio"][1].format("load_video")) + + try: + imageio.plugins.ffmpeg.get_exe() + except AttributeError: + raise AttributeError( + "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg" + ) + + with imageio.get_reader(video) as reader: + fps = reader.get_meta_data().get("fps") + # Read all frames + for frame in reader: + pil_images.append(PIL.Image.fromarray(frame)) + + finally: + if was_tempfile_created: + os.remove(video_path) if convert_method is not None: pil_images = convert_method(pil_images) diff --git a/tests/others/test_loading_utils.py b/tests/others/test_loading_utils.py new file mode 100644 index 000000000000..c8ce39ee5461 --- /dev/null +++ b/tests/others/test_loading_utils.py @@ -0,0 +1,70 @@ +from io import BytesIO + +import pytest +from PIL import Image, UnidentifiedImageError + +from diffusers.utils import load_video + + +class FakeResponse: + status_code = 200 + + def __init__(self, data, interrupt=False): + self.data = data + self.interrupt = interrupt + + def iter_content(self, chunk_size=8192): + yield self.data + if self.interrupt: + raise ConnectionError("Download interrupted") + + +def mock_remote(monkeypatch, tmp_path, response): + monkeypatch.setattr( + "diffusers.utils.loading_utils.requests.get", + lambda *args, **kwargs: response, + ) + monkeypatch.setattr( + "diffusers.utils.loading_utils.tempfile.tempdir", + str(tmp_path), + ) + + +@pytest.mark.parametrize("interrupt", [False, True]) +def test_remote_video_failure_cleans_tempfile(tmp_path, monkeypatch, interrupt): + response = FakeResponse(b"invalid-gif-data", interrupt=interrupt) + mock_remote(monkeypatch, tmp_path, response) + + expected = ConnectionError if interrupt else UnidentifiedImageError + + with pytest.raises(expected): + load_video("https://example.com/broken.gif") + + assert not list(tmp_path.iterdir()) + + +def test_remote_video_success_cleans_tempfile(tmp_path, monkeypatch): + buffer = BytesIO() + Image.new("RGB", (4, 4), "red").save( + buffer, format="GIF", duration=100 + ) + + mock_remote(monkeypatch, tmp_path, FakeResponse(buffer.getvalue())) + + frames, fps = load_video( + "https://example.com/valid.gif", return_fps=True + ) + + assert len(frames) == 1 + assert fps == pytest.approx(10.0) + assert not list(tmp_path.iterdir()) + + +def test_local_video_is_preserved(tmp_path): + video_path = tmp_path / "local.gif" + Image.new("RGB", (4, 4), "red").save(video_path, format="GIF") + + frames = load_video(str(video_path)) + + assert len(frames) == 1 + assert video_path.exists()