lora_loader.py

"""
LoRA loader for WAN 2.2 I2V — gay/MM content catalog.

All LoRAs are I2V format with HIGH/LOW noise pairs unless noted.
HIGH noise = transformer (motion/content); LOW noise = transformer_2 (consistency/detail).

Sources:
  jasbloom/Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128-lora  — mmxxii
  lopi999/Wan2.2-I2V_General-NSFW-LoRA                          — General NSFW
  lkzd7/WAN2.2_LoraSet_NSFW                                     — PENISLORA
  yeqiu168182/wan2.2_i2v_highnoise_pov_missionary_v1.0          — POV Missionary HIGH
  yeqiu168182/wan2.2_i2v_lownoise_pov_missionary_v1.0           — POV Missionary LOW
  UnifiedHorusRA/TheWan2.2I2V                                   — community mirror collection
"""
import re

from huggingface_hub import hf_hub_download

LORA_REPO          = "jasbloom/Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128-lora"
LORA_REPO_NSFW     = "lkzd7/WAN2.2_LoraSet_NSFW"
LORA_REPO_GEN_NSFW = "lopi999/Wan2.2-I2V_General-NSFW-LoRA"
LORA_REPO_POV_H    = "yeqiu168182/wan2.2_i2v_highnoise_pov_missionary_v1.0"
LORA_REPO_POV_L    = "yeqiu168182/wan2.2_i2v_lownoise_pov_missionary_v1.0"
LORA_REPO_UNIFIED  = "UnifiedHorusRA/TheWan2.2I2V"  # community mirror — hundreds of Wan 2.2 LoRAs
HF_TOKEN = None

# Path prefix inside UnifiedHorusRA/TheWan2.2I2V for all LoRAs from that collection
_U = "Wan_Video_2_2_I2V-A14B"  # subfolder inside each LoRA folder

# ── Catalog ──────────────────────────────────────────────────────────────────
#
# Each entry:
#   HIGH / LOW    : filename relative to repo root (or subfolder path for LORA_REPO_UNIFIED)
#   repo          : HF repo ID (defaults to LORA_REPO if absent)
#   repo_high     : override repo for HIGH file (for split-repo LoRAs)
#   repo_low      : override repo for LOW file
#   trigger       : trigger word(s) to include in prompt ("" = none needed)
#   source        : human-readable provenance note
#
LORA_CATALOG = {

    # ── MWM / Gay-specific ────────────────────────────────────────────────
    "Gay Rimjob": {
        "HIGH": f"Wan2_2-i2v-Gay_Rimjob/{_U}/Wan2.2-Gay_Rimjob-High_noise.safetensors",
        "LOW":  f"Wan2_2-i2v-Gay_Rimjob/{_U}/Wan2.2-Gay_Rimjob-low_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V, gay-specific",
        "trigger": "",
    },
    "Stranger (add male to scene)": {
        "HIGH": f"Wan_2_2_Stranger_by_MQ_Lab/{_U}/mql_stranger_erect_wan22_i2v_v1_high_noise.safetensors",
        "LOW":  f"Wan_2_2_Stranger_by_MQ_Lab/{_U}/mql_stranger_erect_wan22_i2v_v1_low_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "MQ Lab / CivitAI 2048121, UnifiedHorusRA mirror",
        "trigger": "",
    },

    # ── mmxxii — dedicated MWM LoRA (3 checkpoints) ──────────────────────
    "mmxxii (step 2750 — strongest)": {
        "HIGH": "Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128_000002750_high_noise.safetensors",
        "LOW":  "Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128_000002750_low_noise.safetensors",
        "source": "jasbloom, I2V rank-128",
        "trigger": "mmxxii",
    },
    "mmxxii (step 2500)": {
        "HIGH": "Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128_000002500_high_noise.safetensors",
        "LOW":  "Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128_000002500_low_noise.safetensors",
        "source": "jasbloom, I2V rank-128",
        "trigger": "mmxxii",
    },
    "mmxxii (step 2250 — lightest touch)": {
        "HIGH": "Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128_000002250_high_noise.safetensors",
        "LOW":  "Wan2.2-I2V-A14B-Diffusers-bf16-mmxxii-rank128_000002250_low_noise.safetensors",
        "source": "jasbloom, I2V rank-128",
        "trigger": "mmxxii",
    },

    # ── Male-specific explicit ────────────────────────────────────────────
    "Meatspin": {
        "HIGH": f"Meatspin_-_Wan_2_2_I2V/{_U}/meatspin_v2_wan22_14b_i2v_high_noise.safetensors",
        "LOW":  f"Meatspin_-_Wan_2_2_I2V/{_U}/meatspin_v2_wan22_14b_i2v_low_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V v2, male anatomy + motion",
        "trigger": "",
    },
    "Jerking Off": {
        "LOW":  f"Jerking_Off/{_U}/Wan2.2 - I2V - Jerking Off - LOW 14B.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V, LOW-noise only",
        "trigger": "",
    },
    "Morning Wood": {
        "HIGH": f"Discovering_a_morning_wood_Real__Wan_2_2_I2V_T2V/{_U}/morning_wood_000000200_high_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V/T2V, HIGH-noise only",
        "trigger": "",
    },
    "Anal Insertion": {
        "HIGH": f"Anal_insertion_I2V/{_U}/anal_insertion_HIGH_V01.safetensors",
        "LOW":  f"Anal_insertion_I2V/{_U}/anal_insertion_LOW_V01.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V v1",
        "trigger": "",
    },

    # ── General NSFW quality base ─────────────────────────────────────────
    "WAN General NSFW": {
        "HIGH": "NSFW-22-H-e8.safetensors",
        "LOW":  "NSFW-22-L-e8.safetensors",
        "repo": LORA_REPO_GEN_NSFW,
        "source": "lopi999/CivitAI 1307155, I2V general NSFW",
        "trigger": "nsfwsks",
    },
    "Mystic XXX": {
        "HIGH": f"NSFW_FluxWan_2_2Qwen_Mystic_XXX/{_U}/Wan2.2-I2V-High-MysticXXX.safetensors",
        "LOW":  f"NSFW_FluxWan_2_2Qwen_Mystic_XXX/{_U}/Wan2.2-I2V-Low-MysticXXX.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V, realism + smooth motion",
        "trigger": "",
    },

    # ── Intimate / romantic ───────────────────────────────────────────────
    "Kissing": {
        "HIGH": f"Kissing_Passionately_T2V__I2V/{_U}/Wan2.2 - I2V - Kissing - HIGH 14B.safetensors",
        "LOW":  f"Kissing_Passionately_T2V__I2V/{_U}/Wan2.2 - I2V - Kissing - LOW 14B.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V 14B",
        "trigger": "",
    },
    "POV Intimate Contact": {
        "HIGH": f"POV_Intimate_Contact/{_U}/povintimatecontact_WAN22_I2V_high_noise.safetensors",
        "LOW":  f"POV_Intimate_Contact/{_U}/povintimatecontact_WAN22_I2V_low_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V",
        "trigger": "",
    },
    "Orgasm": {
        "HIGH": f"Orgasm_-_Wan_2_2_Official_Release/{_U}/Wan2.2 - I2V - Orgasm v2 - 14B_high_noise.safetensors",
        "LOW":  f"Orgasm_-_Wan_2_2_Official_Release/{_U}/Wan2.2 - I2V - Orgasm v2 - 14B_low_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V v2 14B",
        "trigger": "",
    },

    # ── Motion / perspective specialists ──────────────────────────────────
    "POV Missionary": {
        "HIGH": "wan2.2_i2v_highnoise_pov_missionary_v1.0.safetensors",
        "LOW":  "wan2.2_i2v_lownoise_pov_missionary_v1.0.safetensors",
        "repo_high": LORA_REPO_POV_H,
        "repo_low":  LORA_REPO_POV_L,
        "source": "yeqiu168182, I2V v1.0, split-repo upload",
        "trigger": "",
    },

    # ── Quality / realism helpers ─────────────────────────────────────────
    "NSFW Motion Enhancer": {
        "HIGH": f"wan2_2_NSFW_Motion_Enhancer/{_U}/CLK_NSFW_high_v2.1.safetensors",
        "LOW":  f"wan2_2_NSFW_Motion_Enhancer/{_U}/CLK_NSFW_low_v2.1.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, CLK v2.1 — temporal consistency for explicit content",
        "trigger": "",
    },
    "Handheld Camera": {
        "HIGH": f"LuisaP_WAN_2_2_EXTREME_HANDHELD_CAMERA/{_U}/wan2.2handheldcamxtreme_high_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "LuisaP / UnifiedHorusRA mirror, I2V, HIGH-noise only",
        "trigger": "",
    },

    # ── Anatomy helpers ───────────────────────────────────────────────────
    "Penis Play (anatomy)": {
        "HIGH": f"Penis_Play_-_WAN_2_2/{_U}/pworship_high_noise.safetensors",
        "LOW":  f"Penis_Play_-_WAN_2_2/{_U}/pworship_low_noise.safetensors",
        "repo": LORA_REPO_UNIFIED,
        "source": "UnifiedHorusRA mirror, I2V anatomy helper",
        "trigger": "",
    },
    "PENISLORA_22_i2v (anatomy helper)": {
        "HIGH": "PENISLORA_22_i2v_HIGH_e320.safetensors",
        "LOW":  "PENISLORA_22_i2v_LOW_e496.safetensors",
        "repo": LORA_REPO_NSFW,
        "source": "Civitai 1476909, I2V anatomy helper",
        "trigger": "PENISLORA",
    },
}

# ── Role sets ──────────────────────────────────────────────────────────────
LORA_ALIASES = {}
MAX_LORAS = 3

HELPER_LORAS = {
    "PENISLORA_22_i2v (anatomy helper)",
    "Penis Play (anatomy)",
    "NSFW Motion Enhancer",
}
STYLE_LORAS = {
    "Mystic XXX",
    "Handheld Camera",
}
MOTION_LORAS = {
    "POV Missionary",
    "Meatspin",
    "POV Intimate Contact",
}
HIGH_ONLY_LORAS = {
    "Morning Wood",
    "Handheld Camera",
}
LOW_ONLY_LORAS = {
    "Jerking Off",
}
CLOSEUP_ACTION_LORAS = set()

# ── Blend weights ─────────────────────────────────────────────────────────
SINGLE_WEIGHTS = {
    # MWM specific
    "Gay Rimjob": 1.0,
    "Stranger (add male to scene)": 1.0,
    "mmxxii (step 2750 — strongest)": 1.0,
    "mmxxii (step 2500)": 1.0,
    "mmxxii (step 2250 — lightest touch)": 0.9,
    # Male explicit
    "Meatspin": 1.0,
    "Jerking Off": 1.0,
    "Morning Wood": 0.9,
    "Anal Insertion": 1.0,
    # General NSFW base
    "WAN General NSFW": 0.85,
    "Mystic XXX": 0.9,
    # Intimate
    "Kissing": 0.9,
    "POV Intimate Contact": 0.9,
    "Orgasm": 1.0,
    # Motion / perspective
    "POV Missionary": 0.8,
    # Quality helpers
    "NSFW Motion Enhancer": 0.6,
    "Handheld Camera": 0.7,
    "Penis Play (anatomy)": 0.65,
    "PENISLORA_22_i2v (anatomy helper)": 0.7,
}

HELPER_MIX_WEIGHTS = {
    "PENISLORA_22_i2v (anatomy helper)": 0.55,
    "Penis Play (anatomy)": 0.55,
    "NSFW Motion Enhancer": 0.5,
}

# ── Derived lookups ───────────────────────────────────────────────────────
LORA_PAIRS = {
    name: {"HIGH": spec.get("HIGH"), "LOW": spec.get("LOW")}
    for name, spec in LORA_CATALOG.items()
}
LORA_FILES = sorted(
    {
        filename
        for spec in LORA_CATALOG.values()
        for filename in (spec.get("HIGH"), spec.get("LOW"))
        if filename
    }
)


def _resolve_group_name(group_name):
    if not group_name:
        return None
    clean_name = re.sub(r"\s*\(HIGH only\)|\s*\(LOW only\)", "", str(group_name)).strip()
    clean_name = LORA_ALIASES.get(clean_name, clean_name)
    clean_name = LORA_ALIASES.get(group_name, clean_name)
    if clean_name in LORA_CATALOG:
        return clean_name
    lower_map = {name.lower(): name for name in LORA_CATALOG}
    return lower_map.get(clean_name.lower())


def get_lora_choices():
    return list(LORA_CATALOG.keys())


def get_lora_info(group_name):
    resolved = _resolve_group_name(group_name)
    return LORA_CATALOG.get(resolved) if resolved else None


def _coerce_selection(group_names):
    if not group_names:
        return []
    if isinstance(group_names, str):
        group_names = [group_names]
    selected = []
    seen = set()
    for raw_name in group_names:
        if not raw_name or raw_name == "(None)":
            continue
        resolved = _resolve_group_name(raw_name)
        if not resolved:
            raise ValueError(f"Unknown LoRA: {raw_name}")
        if resolved in seen:
            continue
        selected.append(resolved)
        seen.add(resolved)
    return selected


def get_lora_role(group_name):
    resolved = _resolve_group_name(group_name)
    if not resolved:
        return "unknown"
    if resolved in HELPER_LORAS:
        return "helper"
    if resolved in STYLE_LORAS:
        return "style"
    if resolved in HIGH_ONLY_LORAS:
        return "high_only"
    if resolved in LOW_ONLY_LORAS:
        return "low_only"
    if resolved in CLOSEUP_ACTION_LORAS:
        return "closeup_action"
    if resolved in MOTION_LORAS:
        return "motion"
    return "action"


def get_lora_blend_plan(group_names):
    selected = _coerce_selection(group_names)
    if len(selected) > MAX_LORAS:
        raise ValueError(f"Select at most {MAX_LORAS} LoRAs")
    multi = len(selected) > 1
    primary_action_assigned = False
    plan = []
    for name in selected:
        role = get_lora_role(name)
        if not multi:
            weight = SINGLE_WEIGHTS.get(name, 1.0)
        elif role == "helper":
            weight = HELPER_MIX_WEIGHTS.get(name, 0.55)
        elif role == "style":
            weight = 0.5
        elif role == "high_only":
            weight = 0.65
        elif role == "motion":
            weight = SINGLE_WEIGHTS.get(name, 0.8)
        elif not primary_action_assigned:
            weight = 0.85
            primary_action_assigned = True
        else:
            weight = 0.5
        plan.append({"name": name, "role": role, "weight": weight})
    return plan


def download_lora(group_name):
    resolved = _resolve_group_name(group_name)
    if not resolved:
        return None, None
    pair = LORA_CATALOG[resolved]
    default_repo = pair.get("repo", LORA_REPO)
    high_path, low_path = None, None
    if pair.get("HIGH"):
        high_repo = pair.get("repo_high", default_repo)
        high_path = hf_hub_download(high_repo, pair["HIGH"], token=HF_TOKEN)
    if pair.get("LOW"):
        low_repo = pair.get("repo_low", default_repo)
        low_path = hf_hub_download(low_repo, pair["LOW"], token=HF_TOKEN)
    return high_path, low_path


def _load_adapter(pipe, path, adapter_name, component, weight=1.0, group_name=None):
    kwargs = {}
    if component == "transformer_2":
        if not hasattr(pipe, "transformer_2") or pipe.transformer_2 is None:
            raise RuntimeError("This pipeline has no transformer_2 for LOW LoRA loading")
        kwargs["load_into_transformer_2"] = True
    pipe.load_lora_weights(path, adapter_name=adapter_name, **kwargs)
    return {
        "name": adapter_name,
        "component": component,
        "weight": float(weight),
        "group": group_name,
    }


def activate_loras(pipe, adapters):
    if not adapters:
        return
    by_component = {}
    for adapter in adapters:
        by_component.setdefault(adapter["component"], []).append(adapter["name"])
    for component, names in by_component.items():
        model = getattr(pipe, component, None)
        if model is None:
            raise RuntimeError(f"Pipeline component missing: {component}")
        if not hasattr(model, "set_adapters"):
            raise RuntimeError(f"{component} does not expose set_adapters()")
        weights = [float(a.get("weight", 1.0)) for a in adapters if a["component"] == component]
        model.set_adapters(names, weights=weights)
        if hasattr(model, "enable_lora"):
            model.enable_lora()
    if hasattr(pipe, "get_list_adapters"):
        print(f"LoRA adapters loaded: {pipe.get_list_adapters()}")
    if hasattr(pipe, "get_active_adapters"):
        print(f"LoRA adapters active: {pipe.get_active_adapters()}")


def load_lora_to_pipe(pipe, group_name, adapter_name="lora", weight=1.0):
    resolved = _resolve_group_name(group_name)
    high_path, low_path = download_lora(resolved)
    adapters = []
    if high_path and low_path:
        adapters.append(_load_adapter(pipe, high_path, f"{adapter_name}_high", "transformer", weight, resolved))
        adapters.append(_load_adapter(pipe, low_path, f"{adapter_name}_low", "transformer_2", weight, resolved))
        print(f"Loaded LoRA pair: {resolved} weight={weight:.2f}")
        return adapters
    if high_path:
        adapters.append(_load_adapter(pipe, high_path, adapter_name, "transformer", weight, resolved))
        return adapters
    if low_path:
        adapters.append(_load_adapter(pipe, low_path, adapter_name, "transformer_2", weight, resolved))
        return adapters
    raise ValueError(f"LoRA not found: {group_name}")


def unload_lora(pipe):
    try:
        pipe.unload_lora_weights()
    except Exception as exc:
        print(f"LoRA unload warning: {exc}")