lora_loader.py

"""
Custom LoRA Loader for WAN 2.2 I2V.
Add your custom LoRA models (Hugging Face or direct Civitai URLs) in the EXTRA dictionary below.
"""
import os
import urllib.parse
import urllib.request
import urllib.error
import re
import hashlib
import inspect
from huggingface_hub import hf_hub_download

# Monkey patch for peft TorchaoLoraLinear bug (compatibility between peft 0.19.1 and diffusers)
try:
    import peft.tuners.lora.torchao as peft_torchao
    if hasattr(peft_torchao, "TorchaoLoraLinear"):
        _orig_torchao_init = peft_torchao.TorchaoLoraLinear.__init__
        _sig = inspect.signature(_orig_torchao_init)
        if "get_apply_tensor_subclass" in _sig.parameters:
            _param = _sig.parameters["get_apply_tensor_subclass"]
            if _param.default is inspect.Parameter.empty:
                def _patched_torchao_init(self, *args, **kwargs):
                    if "get_apply_tensor_subclass" not in kwargs:
                        base_layer = args[0] if args else kwargs.get("base_layer", None)
                        get_subclass_fn = getattr(base_layer, "get_apply_tensor_subclass", None) if base_layer else None
                        kwargs["get_apply_tensor_subclass"] = get_subclass_fn
                    return _orig_torchao_init(self, *args, **kwargs)
                peft_torchao.TorchaoLoraLinear.__init__ = _patched_torchao_init
                print("✅ Applied peft TorchaoLoraLinear compatibility patch.")
except Exception as patch_err:
    print(f"TorchaoLoraLinear patch notice: {patch_err}")

HF_TOKEN = os.environ.get("HF_TOKEN")  # authenticated downloads (covers private repos)
CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", "")

# Pinned commit hashes if needed (optional)
PINNED_REVISIONS = {}

LORA_FILES = []

# group -> {"HIGH": (repo, file) | url | None, "LOW": (repo, file) | url | None}
LORA_PAIRS = {}
for f in LORA_FILES:
    name = urllib.parse.unquote(f).replace(".safetensors", "")
    is_high = bool(re.search(r'(high|HN|_H\b)', name, re.IGNORECASE))
    is_low = bool(re.search(r'(low|LN|_L\b)', name, re.IGNORECASE))
    group = re.sub(r'[\s_-]*(high|low|noise|HN|LN)([\s_-]*noise)?[\s_-]*(v?\d+(\.\d+)?)?\s*$', '', name, flags=re.IGNORECASE).strip()
    group = re.sub(r'[\s_]+$', '', group)
    LORA_PAIRS.setdefault(group, {"HIGH": None, "LOW": None})

# Custom LoRAs dictionary (label -> URL or (repo_id, high_file, low_file|None))
# Examples:
# EXTRA = {
#     "My Civitai LoRA": "https://civitai.red/api/download/models/2098405?fileId=1994044",
#     "My HF Dual LoRA": ("username/my-lora-repo", "motion_high.safetensors", "motion_low.safetensors"),
# }
EXTRA = {
    "lopi999 - Wan2.2 I2V General NSFW LoRA (Trigger: nsfwsks)": (
        "lopi999/Wan2.2-I2V_General-NSFW-LoRA",
        "NSFW-22-H-e8.safetensors",
        "NSFW-22-L-e8.safetensors"
    ),
}

for label, item in EXTRA.items():
    LORA_PAIRS.setdefault(label, {"HIGH": None, "LOW": None})
    if isinstance(item, str):
        LORA_PAIRS[label]["HIGH"] = item
    elif isinstance(item, (tuple, list)):
        if len(item) == 3:
            repo, hi, lo = item
            if isinstance(repo, str) and (repo.startswith("http://") or repo.startswith("https://")):
                LORA_PAIRS[label]["HIGH"] = repo
                if hi and isinstance(hi, str) and (hi.startswith("http://") or hi.startswith("https://")):
                    LORA_PAIRS[label]["LOW"] = hi
            else:
                if hi:
                    LORA_PAIRS[label]["HIGH"] = (repo, hi)
                if lo:
                    LORA_PAIRS[label]["LOW"] = (repo, lo)
        elif len(item) == 2:
            hi, lo = item
            if hi:
                LORA_PAIRS[label]["HIGH"] = hi
            if lo:
                LORA_PAIRS[label]["LOW"] = lo


def get_loras_dir():
    data_dir = "/data/loras"
    if os.path.exists("/data") and os.path.isdir("/data"):
        os.makedirs(data_dir, exist_ok=True)
        return data_dir
    os.makedirs("loras", exist_ok=True)
    return "loras"


def download_file_from_url(url, custom_name=None):
    target_dir = get_loras_dir()
    civitai_tok = os.environ.get("CIVITAI_TOKEN", "") or CIVITAI_TOKEN
    
    url_to_fetch = url
    if "huggingface.co" in url_to_fetch.lower() and "/blob/" in url_to_fetch.lower():
        url_to_fetch = url_to_fetch.replace("/blob/", "/resolve/")

    if "civitai" in url.lower() and civitai_tok and "token=" not in url.lower():
        sep = "&" if "?" in url else "?"
        url_to_fetch = f"{url}{sep}token={civitai_tok}"

    if not custom_name:
        url_hash = hashlib.md5(url.encode()).hexdigest()[:8]
        custom_name = f"lora_{url_hash}.safetensors"

    local_path = os.path.join(target_dir, custom_name)
    if os.path.exists(local_path) and os.path.getsize(local_path) > 1000:
        return local_path

    print(f"📥 Downloading LoRA from URL: {url_to_fetch} ...")
    req = urllib.request.Request(url_to_fetch, headers={"User-Agent": "Mozilla/5.0"})
    
    try:
        with urllib.request.urlopen(req) as response:
            cd = response.headers.get("Content-Disposition", "")
            if "filename=" in cd:
                fname = re.findall(r'filename="?([^";]+)"?', cd)
                if fname:
                    real_name = fname[0].strip()
                    if not real_name.endswith(".safetensors"):
                        real_name += ".safetensors"
                    alt_path = os.path.join(target_dir, real_name)
                    if os.path.exists(alt_path) and os.path.getsize(alt_path) > 1000:
                        return alt_path
                    local_path = alt_path

            os.makedirs(os.path.dirname(local_path), exist_ok=True)

            with open(local_path, "wb") as f:
                while True:
                    chunk = response.read(8192)
                    if not chunk:
                        break
                    f.write(chunk)
    except urllib.error.HTTPError as e:
        if e.code == 401:
            raise Exception(
                "Download failed (401 Unauthorized). Civitai requires an API Token. "
                "Add ?token=YOUR_CIVITAI_API_KEY to the download URL or set CIVITAI_TOKEN environment variable."
            )
        raise Exception(f"Failed to download LoRA from URL (Status {e.code}): {e.reason}")
    except Exception as e:
        raise Exception(f"Failed to download LoRA from URL: {e}")

    print(f"✅ Downloaded LoRA successfully: {local_path} ({os.path.getsize(local_path)} bytes)")
    return local_path


def get_lora_choices():
    choices = []
    for group in sorted(LORA_PAIRS.keys()):
        p = LORA_PAIRS[group]
        if p["HIGH"] and p["LOW"]:
            choices.append(group)
        elif p["HIGH"]:
            choices.append(f"{group} (HIGH only)")
        elif p["LOW"]:
            choices.append(f"{group} (LOW only)")
    return choices


def download_lora(group_name):
    if not group_name:
        return None, None
    clean_name = re.sub(r'\s*\(HIGH only\)|\s*\(LOW only\)', '', group_name)
    if clean_name not in LORA_PAIRS:
        return None, None
    pair = LORA_PAIRS[clean_name]
    
    def resolve_entry(entry):
        if not entry:
            return None
        if isinstance(entry, str) and (entry.startswith("http://") or entry.startswith("https://")):
            return download_file_from_url(entry)
        elif isinstance(entry, (tuple, list)) and len(entry) == 2:
            repo, fn = entry
            if isinstance(repo, str) and (repo.startswith("http://") or repo.startswith("https://")):
                return download_file_from_url(repo)
            rev = PINNED_REVISIONS.get(repo)
            return hf_hub_download(repo, fn, token=HF_TOKEN, revision=rev) if rev else hf_hub_download(repo, fn, token=HF_TOKEN)
        return None

    high_path = resolve_entry(pair["HIGH"])
    low_path = resolve_entry(pair["LOW"])
    return high_path, low_path


def load_lora_to_pipe(pipe, group_name, adapter_name="lora"):
    high_path, low_path = download_lora(group_name)
    if high_path and low_path:
        pipe.load_lora_weights(high_path, adapter_name=f"{adapter_name}_high")
        pipe.load_lora_weights(low_path, adapter_name=f"{adapter_name}_low")
        print(f"Loaded LoRA pair: {group_name}")
        return True
    elif high_path:
        pipe.load_lora_weights(high_path, adapter_name=adapter_name)
        print(f"Loaded LoRA: {group_name}")
        return True
    elif low_path:
        pipe.load_lora_weights(low_path, adapter_name=adapter_name)
        print(f"Loaded LoRA (low): {group_name}")
        return True
    return False


def list_cached_loras():
    target_dir = get_loras_dir()
    files = [f for f in os.listdir(target_dir) if f.endswith(".safetensors")]
    files.sort(key=lambda x: os.path.getmtime(os.path.join(target_dir, x)), reverse=True)
    return ["(None / Disable)"] + files


def unload_lora(pipe):
    try:
        pipe.unload_lora_weights()
    except:
        pass


def load_custom_url_lora(pipe, url_or_path, adapter_name="custom_lora", scale=1.0):
    if not url_or_path or not str(url_or_path).strip() or str(url_or_path).strip() == "(None / Disable)":
        return False
    target = str(url_or_path).strip()
    target_dir = get_loras_dir()
    try:
        if os.path.exists(target) and target.endswith(".safetensors"):
            file_path = target
        elif os.path.exists(os.path.join(target_dir, target)):
            file_path = os.path.join(target_dir, target)
        elif os.path.exists(os.path.join("loras", target)):
            file_path = os.path.join("loras", target)
        elif target.startswith("http://") or target.startswith("https://"):
            file_path = download_file_from_url(target)
        else:
            candidates = [os.path.join(target_dir, f) for f in os.listdir(target_dir) if target.lower() in f.lower()] if os.path.exists(target_dir) else []
            if candidates:
                file_path = candidates[0]
            else:
                return False

        if file_path and os.path.exists(file_path):
            try:
                pipe.load_lora_weights(file_path, adapter_name=adapter_name)
                if hasattr(pipe, "set_adapters"):
                    try:
                        pipe.set_adapters([adapter_name], adapter_weights=[float(scale)])
                    except Exception:
                        pass
                print(f"✅ Loaded Custom LoRA: {file_path} (scale={scale})")
                return True
            except Exception as load_err:
                print(f"⚠️ Incompatible or corrupted LoRA detected ({file_path}): {load_err}")
                print(f"🗑️ Automatically deleting incompatible LoRA file: {file_path}...")
                try:
                    os.remove(file_path)
                except Exception as del_err:
                    print(f"Warning deleting file: {del_err}")
                return False
    except Exception as e:
        print(f"❌ Failed to load custom LoRA: {e}")
    return False


def download_custom_lora_ui_action(url):
    import gradio as gr
    if not url or not str(url).strip():
        msg = "<div style='background: rgba(239, 68, 68, 0.15); border: 1.5px solid rgba(239, 68, 68, 0.4); border-radius: 12px; padding: 12px 16px; color: #f87171; font-weight: 600;'>⚠️ Please enter a valid LoRA download URL first.</div>"
        return msg, gr.update(choices=list_cached_loras())
    try:
        path = download_file_from_url(str(url).strip())
        size_mb = round(os.path.getsize(path) / (1024 * 1024), 2)
        fname = os.path.basename(path)
        msg = (
            f"<div style='background: linear-gradient(135deg, rgba(16, 185, 129, 0.22) 0%, rgba(5, 150, 105, 0.12) 100%); "
            f"border: 1.5px solid #10b981; border-radius: 14px; padding: 14px 18px; color: #34d399; font-weight: 700; "
            f"box-shadow: 0 6px 20px rgba(16, 185, 129, 0.25); backdrop-filter: blur(10px); margin-top: 10px;'>"
            f"✅ <b>Custom LoRA Successfully Downloaded & Ready!</b><br>"
            f"<span style='font-weight: 500; font-size: 0.88rem; color: #e2e8f0; margin-top: 4px; display: inline-block;'>"
            f"📦 File: <code>{fname}</code> ({size_mb} MB) • Cached to CPU memory with <b>0 GPU Quota consumed</b>. "
            f"Will automatically fuse on your next video generation!</span></div>"
        )
        return msg, gr.update(choices=list_cached_loras(), value=fname)
    except Exception as e:
        msg = f"<div style='background: rgba(239, 68, 68, 0.15); border: 1.5px solid rgba(239, 68, 68, 0.4); border-radius: 12px; padding: 12px 16px; color: #f87171; font-weight: 600;'>❌ Error downloading LoRA: {e}</div>"
        return msg, gr.update(choices=list_cached_loras())