File size: 5,386 Bytes
5d2a0b7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
164
165
166
167
168
169
170
171
172
import logging
import os
import shutil
import subprocess
import urllib.parse
import urllib.request

import comfy.sd
import comfy.utils
import folder_paths


HF_REPO_BASE = "https://huggingface.co/saliacoel/background/resolve/main"


class Salia_Load_Lora_Background:
    """
    Downloads a single LoRA from the fixed public Hugging Face repo
    `saliacoel/background` when missing, then applies it model-only.

    Input examples that all resolve to the same file:
      - cars
      - cars.safetensors
    """

    CATEGORY = "loaders/saliacoel"
    FUNCTION = "load_single_lora"
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("loaded_out",)
    DESCRIPTION = (
        "Downloads a missing LoRA from saliacoel/background and applies it "
        "with model-only LoRA loading."
    )
    OUTPUT_TOOLTIPS = (
        "model_in with the requested LoRA applied.",
    )

    def __init__(self):
        self.loaded_lora = None

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "name": (
                    "STRING",
                    {
                        "default": "",
                        "multiline": False,
                        "placeholder": "cars",
                        "tooltip": (
                            "LoRA name from saliacoel/background. "
                            "The node resolves it to <name>.safetensors."
                        ),
                    },
                ),
                "model_in": (
                    "MODEL",
                    {"tooltip": "The MODEL input that will receive the LoRA."},
                ),
                "strength": (
                    "FLOAT",
                    {
                        "default": 1.0,
                        "min": -100.0,
                        "max": 100.0,
                        "step": 0.01,
                        "tooltip": "Strength used when applying the LoRA.",
                    },
                ),
            }
        }

    @staticmethod
    def _normalize_name(name: str) -> str:
        base = (name or "").strip()
        if not base:
            raise ValueError("name cannot be empty.")

        # Prevent path traversal or accidental subfolders.
        base = os.path.basename(base)

        if base.lower().endswith(".safetensors"):
            base = base[: -len(".safetensors")]

        if not base:
            raise ValueError("name resolves to an empty base name.")

        return base

    @classmethod
    def _build_filename(cls, name: str) -> str:
        return f"{cls._normalize_name(name)}.safetensors"

    @staticmethod
    def _download_file(url: str, target_path: str) -> None:
        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        tmp_path = target_path + ".download"

        if os.path.exists(tmp_path):
            os.remove(tmp_path)

        wget_path = shutil.which("wget")

        try:
            if wget_path:
                subprocess.run(
                    [wget_path, "-O", tmp_path, url],
                    check=True,
                    cwd=os.path.dirname(target_path),
                )
            else:
                request = urllib.request.Request(
                    url,
                    headers={"User-Agent": "ComfyUI-SaliacoelSingleRepoLoraModelOnly/1.0"},
                )
                with urllib.request.urlopen(request) as response, open(tmp_path, "wb") as out_file:
                    shutil.copyfileobj(response, out_file)

            os.replace(tmp_path, target_path)
        except Exception:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
            raise

    @classmethod
    def _ensure_lora_available(cls, lora_name: str) -> str:
        existing_path = folder_paths.get_full_path("loras", lora_name)
        if existing_path is not None:
            return existing_path

        lora_dirs = folder_paths.get_folder_paths("loras")
        if not lora_dirs:
            raise RuntimeError("No ComfyUI 'loras' folder is configured.")

        target_dir = lora_dirs[0]
        target_path = os.path.join(target_dir, lora_name)
        url = f"{HF_REPO_BASE}/{urllib.parse.quote(lora_name)}"

        logging.info("[SaliacoelSingleRepoLoraModelOnly] Downloading missing LoRA: %s", url)
        cls._download_file(url, target_path)

        resolved_path = folder_paths.get_full_path("loras", lora_name)
        return resolved_path if resolved_path is not None else target_path

    def _get_or_load_lora(self, lora_path: str):
        if self.loaded_lora is not None and self.loaded_lora[0] == lora_path:
            return self.loaded_lora[1]

        lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
        self.loaded_lora = (lora_path, lora)
        return lora

    def load_single_lora(self, name, model_in, strength):
        if strength == 0:
            return (model_in,)

        lora_name = self._build_filename(name)
        lora_path = self._ensure_lora_available(lora_name)
        lora = self._get_or_load_lora(lora_path)
        model_out, _ = comfy.sd.load_lora_for_models(model_in, None, lora, strength, 0)
        return (model_out,)


NODE_CLASS_MAPPINGS = {
    "Salia_Load_Lora_Background": Salia_Load_Lora_Background,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Salia_Load_Lora_Background": "Saliacoel LoRA Loader Background",
}