loras / __init__.py
ykvlvai's picture
Create __init__.py
b46de31 verified
Raw
History Blame Contribute Delete
2.04 kB
# custom_nodes/download_lora_node/__init__.py
import os
import shutil
import tarfile
import tempfile
import requests
class DownloadLoRANode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"lora_url": ("STRING", {"default": "https://example.com/flux-lora.tar"}),
"lora_name": ("STRING", {"default": "my_lora"}),
"loras_dir": ("STRING", {"default": "./ComfyUI/models/loras"}),
}
}
RETURN_TYPES = ("STRING",) # путь к safetensors
FUNCTION = "download_lora"
CATEGORY = "Utils/LoRA"
NODE_NAME = "Download LoRA (.tar URL)"
def download_lora(self, lora_url, lora_name, loras_dir):
os.makedirs(loras_dir, exist_ok=True)
with tempfile.TemporaryDirectory() as tmpdir:
tar_path = os.path.join(tmpdir, f"{lora_name}.tar")
# Скачиваем архив
with requests.get(lora_url, stream=True, timeout=300) as r:
r.raise_for_status()
with open(tar_path, "wb") as f:
for chunk in r.iter_content(1 << 20):
f.write(chunk)
# Распаковываем
with tarfile.open(tar_path, "r") as t:
t.extractall(tmpdir)
# ищем .safetensors внутри архива
safetensor_file = None
for root, _, files in os.walk(tmpdir):
for fn in files:
if fn.endswith(".safetensors"):
safetensor_file = os.path.join(root, fn)
break
if safetensor_file:
break
if not safetensor_file:
raise RuntimeError("❌ Не найден .safetensors в архиве")
dst = os.path.join(loras_dir, f"{lora_name}.safetensors")
shutil.move(safetensor_file, dst)
return (dst,)
NODE_CLASS_MAPPINGS = {
"Download LoRA (.tar URL)": DownloadLoRANode
}