diffusers/src/diffusers/utils/loading_utils.py at main · huggingface/diffusers · GitHub

History

177 lines (145 loc) · 6.01 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

import os

import tempfile

from typing import Any, Callable

from urllib.parse import unquote, urlparse

import PIL.Image

import PIL.ImageOps

import requests

from .constants import DIFFUSERS_REQUEST_TIMEOUT

from .import_utils import BACKENDS_MAPPING, is_imageio_available

def load_image(

image: str | PIL.Image.Image, convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] | None = None

) -> PIL.Image.Image:

"""

Loads `image` to a PIL Image.

Args:

image (`str` or `PIL.Image.Image`):

The image to convert to the PIL Image format.

convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*):

A conversion method to apply to the image after loading it. When set to `None` the image will be converted

"RGB".

Returns:

`PIL.Image.Image`:

A PIL Image.

"""

if isinstance(image, str):

if image.startswith("http://") or image.startswith("https://"):

image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)

elif os.path.isfile(image):

image = PIL.Image.open(image)

else:

raise ValueError(

f"Incorrect path or URL. URLs must start with `http://` or `https://`, and {image} is not a valid path."

)

elif isinstance(image, PIL.Image.Image):

image = image

else:

raise ValueError(

"Incorrect format used for the image. Should be a URL linking to an image, a local path, or a PIL image."

)

image = PIL.ImageOps.exif_transpose(image)

if convert_method is not None:

image = convert_method(image)

else:

image = image.convert("RGB")

return image

def load_video(

video: str,

convert_method: Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]] | None = None,

return_fps: bool = False,

) -> list[PIL.Image.Image] | tuple[list[PIL.Image.Image], float]:

"""

Loads `video` to a list of PIL Image.

Args:

video (`str`):

A URL or Path to a video to convert to a list of PIL Image format.

convert_method (Callable[[list[PIL.Image.Image]], list[PIL.Image.Image]], *optional*):

A conversion method to apply to the video after loading it. When set to `None` the images will be converted

to "RGB".

return_fps (`bool`, *optional*, defaults to `False`):

Whether to also return the frame rate the video was encoded at. Needed by pipelines that resample the input

to the frame rate their model works at, since a list of frames does not carry that information.

Returns:

`list[PIL.Image.Image]` or `tuple[list[PIL.Image.Image], float]`:

The video as a list of PIL images, and its frame rate if `return_fps` is set.

"""

is_url = video.startswith("http://") or video.startswith("https://")

is_file = os.path.isfile(video)

was_tempfile_created = False

if not (is_url or is_file):

raise ValueError(

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

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 convert_method is not None:

pil_images = convert_method(pil_images)

if return_fps:

if fps is None:

raise ValueError(f"Could not read a frame rate from {video}, so `return_fps=True` cannot be honoured.")

return pil_images, fps

return pil_images

# Taken from `transformers`.

def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]:

if "." in tensor_name:

splits = tensor_name.split(".")

for split in splits[:-1]:

new_module = getattr(module, split)

if new_module is None:

raise ValueError(f"{module} has no attribute {split}.")

module = new_module

tensor_name = splits[-1]

return module, tensor_name

def get_submodule_by_name(root_module, module_path: str):

current = root_module

parts = module_path.split(".")

for part in parts:

if part.isdigit():

idx = int(part)

current = current[idx] # e.g., for nn.ModuleList or nn.Sequential

else:

current = getattr(current, part)

return current

diffusers/src/diffusers/utils/loading_utils.py at main · huggingface/diffusers · GitHub

Uh oh!

Expand file tree

Breadcrumbs

loading_utils.py

Latest commit

History

Breadcrumbs

loading_utils.py

File metadata and controls