lora_loader.py

"""
LoRA Loader for WAN 2.2 I2V - multi-repo.
Baseline: lkzd7/WAN2.2_LoraSet_NSFW. Plus curated I2V act/helper LoRAs from
Lythiga/WAN2.2_I2V_Lora, rolechar/rc_wan2.2_i2v_loraset_nsfw_private,
lopi999/Wan2.2-I2V_General-NSFW-LoRA. T2V / lightning / furry / dance / male-char filtered out.

Also resolves plain download links (CivitAI and anything else serving a
.safetensors), so a slot can take a URL instead of a repo file.
"""
import hashlib
import json
import os
import struct
import tempfile
import urllib.parse
import re
from huggingface_hub import hf_hub_download, list_repo_files

LORA_REPO = "lkzd7/WAN2.2_LoraSet_NSFW"
HF_TOKEN = os.environ.get("HF_TOKEN")

PINNED_REVISIONS = {
    "lkzd7/WAN2.2_LoraSet_NSFW": "ccfa867ce4b1",
    "Lythiga/WAN2.2_I2V_Lora": "67ef542e6faa",
    "rolechar/rc_wan2.2_i2v_loraset_nsfw_private": "4855a2f216ac",
    "lopi999/Wan2.2-I2V_General-NSFW-LoRA": "aeef17d7fa51",
}

LORA_FILES = [
    "Blink_Squatting_Cowgirl_Position_I2V_HIGH.safetensors",
    "Blink_Squatting_Cowgirl_Position_I2V_LOW.safetensors",
    "PENISLORA_22_i2v_HIGH_e320.safetensors",
    "PENISLORA_22_i2v_LOW_e496.safetensors",
    "Pornmaster_wan 2.2_14b_I2V_bukkake_v1.4_high_noise.safetensors",
    "Pornmaster_wan 2.2_14b_I2V_bukkake_v1.4_low_noise.safetensors",
    "W22_Multiscene_Photoshoot_Softcore_i2v_HN.safetensors",
    "W22_Multiscene_Photoshoot_Softcore_i2v_LN.safetensors",
    "WAN-2.2-I2V-Double-Blowjob-HIGH-v1.safetensors",
    "WAN-2.2-I2V-Double-Blowjob-LOW-v1.safetensors",
    "WAN-2.2-I2V-HandjobBlowjobCombo-HIGH-v1.safetensors",
    "WAN-2.2-I2V-HandjobBlowjobCombo-LOW-v1.safetensors",
    "WAN-2.2-I2V-SensualTeasingBlowjob-HIGH-v1.safetensors",
    "WAN-2.2-I2V-SensualTeasingBlowjob-LOW-v1.safetensors",
    "iGOON_Blink_Blowjob_I2V_HIGH.safetensors",
    "iGOON_Blink_Blowjob_I2V_LOW.safetensors",
    "iGoon - Blink_Front_Doggystyle_I2V_HIGH.safetensors",
    "iGoon - Blink_Front_Doggystyle_I2V_LOW.safetensors",
    "iGoon - Blink_Missionary_I2V_HIGH.safetensors",
    "iGoon - Blink_Missionary_I2V_LOW v2.safetensors",
    "iGoon - Blink_Missionary_I2V_LOW.safetensors",
    "iGoon%20-%20Blink_Back_Doggystyle_HIGH.safetensors",
    "iGoon%20-%20Blink_Back_Doggystyle_LOW.safetensors",
    "iGoon%20-%20Blink_Facial_I2V_HIGH.safetensors",
    "iGoon%20-%20Blink_Facial_I2V_LOW.safetensors",
    "iGoon_Blink_Missionary_I2V_HIGH v2.safetensors",
    "iGoon_Blink_Titjob_I2V_HIGH.safetensors",
    "iGoon_Blink_Titjob_I2V_LOW.safetensors",
    "lips-bj_high_noise.safetensors",
    "lips-bj_low_noise.safetensors",
    "mql_casting_sex_doggy_kneel_diagonally_behind_vagina_wan22_i2v_v1_high_noise.safetensors",
    "mql_casting_sex_doggy_kneel_diagonally_behind_vagina_wan22_i2v_v1_low_noise.safetensors",
    "mql_casting_sex_reverse_cowgirl_lie_front_vagina_wan22_i2v_v1_high_noise.safetensors",
    "mql_casting_sex_reverse_cowgirl_lie_front_vagina_wan22_i2v_v1_low_noise.safetensors",
    "mql_casting_sex_spoon_wan22_i2v_v1_high_noise.safetensors",
    "mql_casting_sex_spoon_wan22_i2v_v1_low_noise.safetensors",
    "mql_massage_tits_wan22_i2v_v1_high_noise.safetensors",
    "mql_massage_tits_wan22_i2v_v1_low_noise.safetensors",
    "mql_panties_aside_wan22_i2v_v1_high_noise.safetensors",
    "mql_panties_aside_wan22_i2v_v1_low_noise.safetensors",
    "sfbehind_v2.1_high_noise.safetensors",
    "sfbehind_v2.1_low_noise.safetensors",
    "sid3l3g_transition_v2.0_H.safetensors",
    "sid3l3g_transition_v2.0_L.safetensors",
    "wan2.2_i2v_high_ulitmate_pussy_asshole.safetensors",
    "wan2.2_i2v_low_ulitmate_pussy_asshole.safetensors",
    "wan22-mouthfull-140epoc-high-k3nk.safetensors",
    "wan22-mouthfull-152epoc-low-k3nk.safetensors",
]

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})
    if is_high:
        LORA_PAIRS[group]["HIGH"] = (LORA_REPO, f)
    elif is_low:
        LORA_PAIRS[group]["LOW"] = (LORA_REPO, f)

L = "Lythiga/WAN2.2_I2V_Lora"
R = "rolechar/rc_wan2.2_i2v_loraset_nsfw_private"
P = "lopi999/Wan2.2-I2V_General-NSFW-LoRA"

EXTRA = {
    "Ultimate Deepthroat": (L, "wan22-ultimatedeepthroat-i2v-102epoc-high-k3nk.safetensors", "wan22-ultimatedeepthroat-I2V-101epoc-low-k3nk.safetensors"),
    "BBC Deepthroat": (L, "wan22-bbcdeepthroat-115epoc-high-k3nk.safetensors", "wan22-bbcdeepthroat-155epoc-low-720-k3nk.safetensors"),
    "Throat V2": (L, "Wan22_ThroatV2_High.safetensors", "Wan22_ThroatV2_Low.safetensors"),
    "JFJ Deepthroat": (L, "jfj-deepthroat-W22-I2V-HN.safetensors", None),
    "BBC Blowjob Extreme": (L, "BBC Blowjobs Extreme_high_noise.safetensors", "BBC Blowjobs Extreme_low_noise.safetensors"),
    "POV Blowjob": (L, "Pov_Blowjob_Wan2.2_high_I2V_v1.0.safetensors", None),
    "Cunnilingus": (L, "wan22-cunilingus-I2V-106epoc-high.safetensors", "wan22-cunilingus-I2V-72epoc-low.safetensors"),
    "Reverse Cowgirl v2": (L, "wan22.r3v3rs3_c0wg1rl-14b-High-i2v_e70.safetensors", "wan22.r3v3rs3_c0wg1rl-14b-Low-i2v_e70.safetensors"),
    "Doggy Front View v2": (L, "doggy_style_sex_front_view_Wan2.2_I2V_high_v1.0.safetensors", "doggy_style_sex_front_view_Wan2.2_I2V_Low_v1.0.safetensors"),
    "Fucked From Behind": (L, "derpfuckedfrombehind_000002125_high_noise.safetensors", "derpfuckedfrombehind_000002125_low_noise.safetensors"),
    "Prone Bone": (L, "Pronebone_high_noise.safetensors", "Pronebone_low_noise.safetensors"),
    "Reverse Suspended Congress": (L, "reverse_suspended_congress_I2V_high.safetensors", "reverse_suspended_congress_I2V_low.safetensors"),
    "POV Missionary": (L, "wan2.2_i2v_highnoise_pov_missionary_v1.0.safetensors", "wan2.2_i2v_lownoise_pov_missionary_v1.0.safetensors"),
    "Mating Press": (L, "mating_press_high.safetensors", "mating_press_low.safetensors"),
    "Standing Upright Sex": (L, "Upright_sex_Wan2.2_14B_TI2V_HIGH_v1.0.safetensors", "Upright_sex_Wan2.2_14B_TI2V_LOW_v1.0.safetensors"),
    "Oral Insertion": (L, "wan2.2-i2v-high-oral-insertion-v1.0.safetensors", "wan2.2-i2v-low-oral-insertion-v1.0.safetensors"),
    "Handjob": (L, "WAN-2.2-I2V-Handjob-HIGH-v1.safetensors", "WAN-2.2-I2V-Handjob-LOW-v1.safetensors"),
    "Balls Sucking": (L, "WAN-2.2-I2V_Balls_sucking_High_noise.safetensors", "WAN-2.2-I2V_Balls_sucking_Low_noise.safetensors"),
    "Fingering": (R, "Sensual_fingering_v1_high_noise.safetensors", "Sensual_fingering_v1_low_noise.safetensors"),
    "Facesplash Cumshot": (L, "wan22-f4c3spl4sh-100epoc-high-k3nk.safetensors", "wan22-f4c3spl4sh-154epoc-low-k3nk.safetensors"),
    "Creampie CRM": (L, "Creampie CRM-FULL-EPOCH-80-HIGH.safetensors", "Creampie CRM-FULL-EPOCH-80-LOW.safetensors"),
    "Pornmaster Creampie": (L, "Pornmaster_wan 2.2_14b_I2V_Creampie_v1_high_noise.safetensors", "Pornmaster_wan 2.2_14b_I2V_Creampie_v1_low_noise.safetensors"),
    "Cum (generic)": (L, "Wan22_Cum_high_noise_1.V1.safetensors", "Wan22_Cum_low_noise_1.V1.safetensors"),
    "French Kiss": (L, "WAN2.2-FrenchKiss_HighNoise.safetensors", "WAN2.2-FrenchKiss_LowNoise.safetensors"),
    "Twerking": (L, "Twerking_I2V_high_noise.safetensors", "Twerking_I2V_low_noise.safetensors"),
    "Ahegao Drooling": (R, "middlefinger_drooling_ahegao_high.safetensors", "middlefinger_drooling_ahegao_low.safetensors"),
    "General NSFW Booster": (P, "NSFW-22-H-e8.safetensors", "NSFW-22-L-e8.safetensors"),
    "Penetration Insert": (L, "PenInsert_high_noise.safetensors", "PenInsert_low_noise.safetensors"),
    "Biggest Cock (size)": (L, "WAN-2.2-I2V_BiggestCock_high_noise_V1.safetensors", "WAN-2.2-I2V_BiggestCock_low_noise_V1.safetensors"),
    "Doggy Slider": (L, "I2V_doggyslider_high.safetensors", "I2V_doggyslider_low.safetensors"),
    "Pussy Helper": (R, "wan2.2_i2v_ulitmate_pussy_helper_high.safetensors", "wan2.2_i2v_ulitmate_pussy_helper_low.safetensors"),
    "FOV Slider (camera)": (L, "wan2.2-i2v-high-sex-fov-slider-v1.0.safetensors", "wan2.2-i2v-low-sex-fov-slider-v1.0.safetensors"),
    "Smashcut (camera)": (R, "wan2.2-i2v-sex-smashcut-v1.0-high.safetensors", "wan2.2-i2v-sex-smashcut-v1.0-low.safetensors"),
}
for label, (repo, hi, lo) in EXTRA.items():
    LORA_PAIRS.setdefault(label, {"HIGH": None, "LOW": None})
    if hi:
        LORA_PAIRS[label]["HIGH"] = (repo, hi)
    if lo:
        LORA_PAIRS[label]["LOW"] = (repo, lo)


def get_lora_choices():
    choices = ["(None)"]
    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 parse_repo_link(link: str):
    if not link:
        return None
    link = link.strip().strip('"').strip("'")
    m = re.match(r"https?://huggingface\.co/(?:models/)?([^/]+)/([^/]+)", link)
    if m:
        return f"{m.group(1)}/{m.group(2)}"
    if not link.startswith("http") and link.count("/") == 1:
        return link
    return None


def parse_hf_file_link(link: str):
    if not link:
        return None, None
    link = link.strip().strip('"').strip("'")
    m = re.match(
        r"https?://huggingface\.co/([^/]+)/([^/]+)/(?:resolve|blob)/[^/]+/(.+)",
        link,
    )
    if m:
        return f"{m.group(1)}/{m.group(2)}", urllib.parse.unquote(m.group(3).split("?")[0])
    if not link.startswith("http") and link.count("/") >= 2:
        parts = link.split("/")
        return f"{parts[0]}/{parts[1]}", "/".join(parts[2:])
    return None, None


def list_loras_from_repo(repo_link: str):
    """Files in a Hugging Face repo. A direct download link is handed back as the
    only choice, so pasting a CivitAI URL into the repo box fills the dropdowns."""
    link = (repo_link or "").strip()
    if is_direct_link(link):
        return ["(None)", link]

    repo_id = parse_repo_link(link)
    if not repo_id:
        return ["(None)"]
    try:
        files = list_repo_files(repo_id, token=HF_TOKEN, revision=PINNED_REVISIONS.get(repo_id))
        safes = sorted([f for f in files if f.lower().endswith(".safetensors")])
        return ["(None)"] + safes if safes else ["(None)"]
    except Exception as e:
        print(f"list_loras_from_repo error: {e}")
        return ["(None)"]


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]
    high_path, low_path = None, None
    if pair["HIGH"]:
        repo, fn = pair["HIGH"]
        high_path = hf_hub_download(repo, fn, token=HF_TOKEN, revision=PINNED_REVISIONS.get(repo))
    if pair["LOW"]:
        repo, fn = pair["LOW"]
        low_path = hf_hub_download(repo, fn, token=HF_TOKEN, revision=PINNED_REVISIONS.get(repo))
    return high_path, low_path


# ----------------------------------------------------------------------------
# Direct links (CivitAI and any other plain .safetensors URL)
# ----------------------------------------------------------------------------
# The download runs on the Space's machine, not in the visitor's browser, so a
# CivitAI session in a browser tab does not authenticate it: a gated model
# answers a server with an HTML login page instead of weights. CIVITAI_TOKEN
# (Settings -> Variables and secrets) is appended automatically when it is set,
# and the file header is checked, so a login page fails with a sentence that
# says what to do rather than a parse error deep inside safetensors.

CIVITAI_HOSTS = ("civitai.com", "civitai.red", "civitai.green", "civitai.work")
CIVITAI_TOKEN = os.environ.get("CIVITAI_TOKEN", "").strip()


def is_direct_link(text: str) -> bool:
    """A URL that is not a huggingface.co link, so a plain file download."""
    text = (text or "").strip().strip('"').strip("'")
    if not text.startswith(("http://", "https://")):
        return False
    host = (urllib.parse.urlparse(text).hostname or "").lower()
    return not host.endswith("huggingface.co")


def download_direct_link(url: str) -> str:
    import requests

    url = url.strip().strip('"').strip("'")
    host = (urllib.parse.urlparse(url).hostname or "").lower()
    request_url = url
    if CIVITAI_TOKEN and any(host.endswith(known) for known in CIVITAI_HOSTS) and "token=" not in url:
        request_url = url + ("&" if "?" in url else "?") + f"token={CIVITAI_TOKEN}"

    cache_dir = os.path.join(tempfile.gettempdir(), "url-loras")
    os.makedirs(cache_dir, exist_ok=True)
    cached = os.path.join(cache_dir, hashlib.sha256(url.encode()).hexdigest()[:16] + ".safetensors")
    if os.path.exists(cached) and os.path.getsize(cached) > 1_000_000:
        return cached

    response = requests.get(request_url, stream=True, timeout=180,
                            headers={"User-Agent": "Mozilla/5.0"})
    response.raise_for_status()

    if "text/html" in (response.headers.get("content-type") or "").lower():
        raise ValueError(
            "That link answered with a web page instead of a file. The model is gated, so the "
            "Space needs its own key: add CIVITAI_TOKEN under Settings -> Variables and secrets."
        )

    written = 0
    with open(cached, "wb") as handle:
        for chunk in response.iter_content(chunk_size=1 << 20):
            if chunk:
                handle.write(chunk)
                written += len(chunk)

    if written < 1_000_000:
        os.remove(cached)
        raise ValueError(
            "That link returned only a few kilobytes - almost always a login or error page rather "
            "than weights. Check the link, or add CIVITAI_TOKEN to the Space."
        )

    # safetensors starts with an 8-byte little-endian header length, then that much JSON.
    with open(cached, "rb") as handle:
        length = struct.unpack("<Q", handle.read(8))[0]
        if not (0 < length < 100_000_000):
            os.remove(cached)
            raise ValueError("The downloaded file is not a .safetensors (bad header).")
        try:
            json.loads(handle.read(length).decode("utf-8"))
        except Exception:
            os.remove(cached)
            raise ValueError("The downloaded file is not a .safetensors (unreadable header).")

    print(f"Downloaded {written / 1e6:.0f} MB from {host}")
    return cached


def download_repo_file(repo_link: str, filename: str):
    if not filename or filename == "(None)":
        return None

    # A local file (already downloaded, or uploaded into the Space) is used as it is.
    if os.path.exists(filename):
        return filename

    # A plain download link, in either box.
    if is_direct_link(filename):
        return download_direct_link(filename)
    if is_direct_link(repo_link) and filename == repo_link:
        return download_direct_link(repo_link)

    repo_from_file, weight = parse_hf_file_link(filename)
    if repo_from_file and weight:
        return hf_hub_download(
            repo_from_file, weight, token=HF_TOKEN,
            revision=PINNED_REVISIONS.get(repo_from_file),
        )
    repo_id = parse_repo_link(repo_link)
    if not repo_id:
        raise ValueError(f"Invalid repo/file: repo={repo_link}, file={filename}")
    return hf_hub_download(
        repo_id, filename, token=HF_TOKEN,
        revision=PINNED_REVISIONS.get(repo_id),
    )


def _load_path(pipe, path, adapter_name, into_t2=False):
    kwargs = {"adapter_name": adapter_name}
    if into_t2:
        kwargs["load_into_transformer_2"] = True
    pipe.load_lora_weights(path, **kwargs)
    return adapter_name


def load_lora_to_pipe(pipe, group_name, adapter_name="lora", high_scale=1.0, low_scale=1.0):
    if not group_name or group_name == "(None)":
        return [], []
    high_path, low_path = download_lora(group_name)
    names, weights = [], []
    if high_path:
        n = f"{adapter_name}_high"
        _load_path(pipe, high_path, n, into_t2=False)
        names.append(n)
        weights.append(float(high_scale))
    if low_path:
        n = f"{adapter_name}_low"
        _load_path(pipe, low_path, n, into_t2=True)
        names.append(n)
        weights.append(float(low_scale))
    if names:
        print(f"Loaded preset LoRA: {group_name} scales={weights}")
    return names, weights


def load_lora_slots(pipe, slots):
    all_names, all_weights = [], []
    for idx, slot in enumerate(slots or []):
        if not slot:
            continue
        preset = slot.get("preset")
        if preset and preset != "(None)":
            names, weights = load_lora_to_pipe(
                pipe, preset, adapter_name=f"lora_{idx}",
                high_scale=slot.get("high_scale", 1.0),
                low_scale=slot.get("low_scale", 1.0),
            )
            all_names.extend(names)
            all_weights.extend(weights)
            continue
        repo = slot.get("repo")
        fh = slot.get("file_high")
        fl = slot.get("file_low")
        hs = float(slot.get("high_scale", 1.0) or 1.0)
        ls = float(slot.get("low_scale", 1.0) or 1.0)
        if fh and fh != "(None)":
            path = download_repo_file(repo, fh)
            n = f"link_{idx}_high"
            _load_path(pipe, path, n, into_t2=False)
            all_names.append(n)
            all_weights.append(hs)
            print(f"Loaded link HIGH: {fh} scale={hs}")
        if fl and fl != "(None)":
            path = download_repo_file(repo, fl)
            n = f"link_{idx}_low"
            _load_path(pipe, path, n, into_t2=True)
            all_names.append(n)
            all_weights.append(ls)
            print(f"Loaded link LOW: {fl} scale={ls}")
    if all_names:
        pipe.set_adapters(all_names, adapter_weights=all_weights)
        print(f"Active adapters: {list(zip(all_names, all_weights))}")
        return True
    return False


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