Salia_Load_Lora_Background.py

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",
}