""" 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 = "
⚠️ Please enter a valid LoRA download URL first.
" 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"
" f"✅ Custom LoRA Successfully Downloaded & Ready!
" f"" f"📦 File: {fname} ({size_mb} MB) • Cached to CPU memory with 0 GPU Quota consumed. " f"Will automatically fuse on your next video generation!
" ) return msg, gr.update(choices=list_cached_loras(), value=fname) except Exception as e: msg = f"
❌ Error downloading LoRA: {e}
" return msg, gr.update(choices=list_cached_loras())