FilesExpand file tree

main

/

loading_utils.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

163 lines (132 loc) · 5.18 KB

main

/

loading_utils.py

Copy path

Top

File metadata and controls

• Code
• Blame

163 lines (132 loc) · 5.18 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

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

import os

import tempfile

from typing import Any, Callable, List, Optional, Tuple, Union

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: Union[str, PIL.Image.Image], convert_method: Optional[Callable[[PIL.Image.Image], PIL.Image.Image]] = 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: Optional[Callable[[List[PIL.Image.Image]], List[PIL.Image.Image]]] = None,

) -> List[PIL.Image.Image]:

"""

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

Returns:

`List[PIL.Image.Image]`:

The video as a list of PIL images.

"""

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 = []

if video.endswith(".gif"):

gif = PIL.Image.open(video)

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:

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

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