Spaces:
Runtime error
Runtime error
File size: 2,512 Bytes
5fbd25d |
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 |
import hashlib
import os
import requests
import tarfile
def _hash_url(url):
"""Generates a hash value for a given URL."""
return hashlib.md5(url.encode('utf-8')).hexdigest()
class LoraManager:
"""
Manager loras from url
"""
def __init__(self):
self.cache_dir = os.path.join(
os.path.dirname(os.path.realpath(__file__)),
'../../',
'repositories/Fooocus/models/loras')
def _download_lora(self, url):
"""
Downloads a LoRa from a URL, saves it in the cache, and if it's a .tar file, extracts it and returns the .safetensors file.
"""
url_hash = _hash_url(url)
file_ext = url.split('.')[-1]
filepath = os.path.join(self.cache_dir, f"{url_hash}.{file_ext}")
if not os.path.exists(filepath):
print(f"Start download for: {url}")
try:
response = requests.get(url, timeout=10, stream=True)
response.raise_for_status()
with open(filepath, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
if file_ext == "tar":
print("Extracting the tar file...")
with tarfile.open(filepath, 'r:*') as tar:
tar.extractall(path=self.cache_dir)
print("Extraction completed.")
return self._find_safetensors_file(self.cache_dir)
print(f"Download successfully, saved as {filepath}")
except Exception as e:
raise Exception(f"Error downloading {url}: {e}") from e
else:
print(f"LoRa already downloaded {url}")
return filepath
def _find_safetensors_file(self, directory):
"""
Finds the first .safetensors file in the specified directory.
"""
print("Searching for .safetensors file.")
for root, dirs, files in os.walk(directory):
for file in files:
if file.endswith('.safetensors'):
return os.path.join(root, file)
raise FileNotFoundError("No .safetensors file found in the extracted files.")
def check(self, urls):
"""Manages the specified LoRAs: downloads missing ones and returns their file names."""
paths = []
for url in urls:
path = self._download_lora(url)
paths.append(path)
return paths
|