Compare commits

..

1 Commits

Author SHA1 Message Date
Jong Wook Kim 551aca3475
Fixes #396 2023-10-16 15:47:21 -07:00
2 changed files with 5 additions and 5 deletions

View File

@ -2,8 +2,8 @@ import hashlib
import os import os
import urllib import urllib
import warnings import warnings
from packaging import version from typing import Any, Union, List
from typing import Union, List from pkg_resources import packaging
import torch import torch
from PIL import Image from PIL import Image
@ -20,7 +20,7 @@ except ImportError:
BICUBIC = Image.BICUBIC BICUBIC = Image.BICUBIC
if version.parse(torch.__version__) < version.parse("1.7.1"): if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
warnings.warn("PyTorch version 1.7.1 or higher is recommended") warnings.warn("PyTorch version 1.7.1 or higher is recommended")
@ -133,6 +133,7 @@ def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_a
if jit: if jit:
warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead") warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
jit = False jit = False
opened_file.seek(0)
state_dict = torch.load(opened_file, map_location="cpu") state_dict = torch.load(opened_file, map_location="cpu")
if not jit: if not jit:
@ -228,7 +229,7 @@ def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: b
sot_token = _tokenizer.encoder["<|startoftext|>"] sot_token = _tokenizer.encoder["<|startoftext|>"]
eot_token = _tokenizer.encoder["<|endoftext|>"] eot_token = _tokenizer.encoder["<|endoftext|>"]
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
if version.parse(torch.__version__) < version.parse("1.8.0"): if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
else: else:
result = torch.zeros(len(all_tokens), context_length, dtype=torch.int) result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)

View File

@ -1,5 +1,4 @@
ftfy ftfy
packaging
regex regex
tqdm tqdm
torch torch