"""LoRA catalog + per-session custom LoRA loader (JSON-first). Single shared catalog lives in an external JSON file (``PERSISTENT_LORA_CATALOG_PATH``, default ``/loras-flux/config/lorasplayground.json``). Layers: * EXTERNAL_LORA_STYLES — full catalog from JSON (including archived). UI listing only shows entries with ``active: true``. * dynamic_loras (gr.State) — session-private try-outs. * LOADED_ADAPTERS — names already attached to the shared pipe. JSON entry fields: title, adapter_name, repo, weights, default_prompt, default_weight, admin_approved (bool), active (bool; false = archived/hidden from UI), sha256 (optional hex digest of the weight file), image (optional). """ from __future__ import annotations import hashlib import json import os import threading import uuid from pathlib import Path import gradio as gr from config import MAX_LORA_SLOTS, PERSISTENT_LORA_CATALOG_PATH FACE_SWAP_PROMPT = """head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. Remove the head from Picture 1 completely and replace it with the head from Picture 2. FROM PICTURE 1 (strictly preserve): - Scene: lighting conditions, shadows, highlights, color temperature, environment, background - Head positioning: exact rotation angle, tilt, direction the head is facing - Expression: facial expression, micro-expressions, eye gaze direction, mouth position, emotion FROM PICTURE 2 (strictly preserve identity): - Facial structure: face shape, bone structure, jawline, chin - All facial features: eye color, eye shape, nose structure, lip shape and fullness, eyebrows - Hair: color, style, texture, hairline - Skin: texture, tone, complexion The replaced head must seamlessly match Picture 1's lighting and expression while maintaining the complete identity from Picture 2. High quality, photorealistic, sharp details, 4k.""" _DEFAULT_LORA_IMAGE = ( "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/" "resolve/main/examples/image.webp" ) # Seed used only when the external JSON is missing/empty. After first write, # the JSON is the single source of truth — nothing here is merged at runtime. _SEED_LORA_STYLES = [ { "title": "Klein-Delight-Style", "adapter_name": "klein-delight", "repo": "linoyts/Flux2-Klein-Delight-LoRA", "weights": "pytorch_lora_weights.safetensors", "default_prompt": ( "Relight the image to remove all existing lighting conditions and replace them " "with neutral, uniform illumination. Apply soft, evenly distributed lighting with " "no directional shadows, no harsh highlights, and no dramatic contrast. Maintain " "the original identity of all subjects exactly—preserve facial structure, skin tone, " "proportions, expressions, hair, clothing, and textures. Do not alter pose, camera " "angle, background geometry, or image composition. Lighting should appear balanced, " "and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure " "consistent exposure across the entire image with realistic depth and subtle shading " "only where necessary for form." ), "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "Klein-Consistency", "adapter_name": "klein-consistency", "repo": "dx8152/Flux2-Klein-9B-Consistency", "weights": "Klein-consistency.safetensors", "default_prompt": None, "default_weight": 0.3, "admin_approved": True, "active": True, }, { "title": "Best-Face-Swap", "adapter_name": "face-swap", "repo": "Alissonerdx/BFS-Best-Face-Swap", "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors", "default_prompt": FACE_SWAP_PROMPT, "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "NSFW v2", "adapter_name": "nsfw-v2", "repo": "diroverflo/FLux_Klein_9B_NSFW", "weights": "Flux Klein - NSFW v2.safetensors", "default_prompt": None, "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "Ultimate Upscaler Klein-9b", "adapter_name": "Ultimate Upscaler", "repo": "loras", "weights": "Flux2-Klein-Image-RestoreV1.safetensors", "default_prompt": ( "restore the image quality, remove any compression artefacts, remove any haze " "and soft edges, enrich the original with new intricate detail in all textures " "and surfaces creating a professional photorealistic photograph with natural " "lighting and skin texture." ), "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "High Resolution", "adapter_name": "High Resolution", "repo": "loras", "weights": "HighResolution9B.safetensors", "default_prompt": "High Resolution", "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "InstaPic", "adapter_name": "InstaPic V3", "repo": "loras", "weights": "InstaPic V3.safetensors", "default_prompt": "instapic", "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "Realistic Nudes", "adapter_name": "Realistic Nudes", "repo": "loras", "weights": "realistic_nudes_klein_v3.safetensors", "default_prompt": None, "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "Perky Pointy Puffy Breasts", "adapter_name": "Perky Pointy Puffy Breasts", "repo": "loras", "weights": "PerkyPointyPuffy_v1.1_small_pointy_breasts_large_puffy_nipples.safetensors", "default_prompt": "Small pointy breasts with large puffy nipples", "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "Flat Chested", "adapter_name": "Flat Chested", "repo": "loras", "weights": "Flux2-Klein-9b-FlatChested-v1.safetensors", "default_prompt": "flat chested", "default_weight": 1.5, "admin_approved": True, "active": True, }, { "title": "Controllight", "adapter_name": "Controllight", "repo": "ControlLight/ControlLight", "weights": "controllight.safetensors", "default_prompt": None, "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "RefControl - Depth", "adapter_name": "RefConDep", "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-depth-lora", "weights": "flux2_klein_9b_refcontrol_depth.safetensors", "default_prompt": "refcontrol", "default_weight": 1.0, "admin_approved": True, "active": True, }, { "title": "RefControl - Pose", "adapter_name": "RefConPos", "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-pose-lora", "weights": "refcontrol_v2_poses.safetensors", "default_prompt": "apply pose from image 1 with reference from image 2", "default_weight": 1.0, "admin_approved": True, "active": True, }, ] # Back-compat alias — older imports still work; runtime catalog is JSON-only. LORA_STYLES: list[dict] = [] LOADED_ADAPTERS: set[str] = set() _CATALOG_LOCK = threading.Lock() EXTERNAL_LORA_STYLES: list[dict] = [] # full JSON list (active + archived) _WEIGHT_EXTS = (".safetensors", ".bin") _DEFAULT_WEIGHT_CANDIDATES = ( "pytorch_lora_weights.safetensors", "lora.safetensors", "adapter_model.safetensors", ) _HASH_CHUNK = 1024 * 1024 def _sanitize_name(value: str, fallback: str = "custom") -> str: cleaned = "".join(c if c.isalnum() or c in "-_" else "_" for c in (value or "")) return cleaned.strip("_") or fallback def _as_bool(value, default: bool = False) -> bool: if value is None: return default if isinstance(value, bool): return value if isinstance(value, (int, float)): return bool(value) if isinstance(value, str): return value.strip().lower() in {"1", "true", "yes", "y", "on"} return default def _file_sha256(path: str | Path) -> str | None: p = Path(path) if not p.is_file(): return None h = hashlib.sha256() try: with open(p, "rb") as f: while True: chunk = f.read(_HASH_CHUNK) if not chunk: break h.update(chunk) return h.hexdigest() except OSError as e: print(f"[lora_registry] sha256 failed for {p}: {e}") return None def _weight_path(repo: str, weights: str) -> Path | None: if not repo or not weights: return None if str(repo).startswith("/"): candidate = Path(repo) / weights if candidate.is_file(): return candidate candidate = Path(repo) / weights if candidate.is_file(): return candidate if str(repo).endswith(_WEIGHT_EXTS) and Path(repo).is_file(): return Path(repo) return None def _compute_entry_sha256(entry: dict) -> str | None: if entry.get("sha256"): return str(entry["sha256"]).lower() path = _weight_path(entry.get("repo", ""), entry.get("weights", "")) if path is None: return None return _file_sha256(path) def _parse_trigger_list(value) -> list[str]: """Normalize known_triggers from JSON list or free text (commas/newlines).""" if value is None: return [] if isinstance(value, list): out = [] for item in value: s = str(item).strip() if s: out.append(s) return out text = str(value).replace(",", "\n") return [line.strip() for line in text.splitlines() if line.strip()] def _normalize_catalog_entry(raw) -> dict | None: if not isinstance(raw, dict): return None title = (raw.get("title") or "").strip() repo = raw.get("repo") weights = raw.get("weights") if not title or not repo or not weights: return None adapter = (raw.get("adapter_name") or "").strip() or _sanitize_name(title) try: default_weight = float(raw.get("default_weight", 1.0)) except (TypeError, ValueError): default_weight = 1.0 prompt = raw.get("default_prompt") if isinstance(prompt, str): prompt = prompt.strip() or None else: prompt = None if "active" in raw: active = _as_bool(raw.get("active"), True) else: active = not _as_bool(raw.get("archived"), False) sha = raw.get("sha256") or raw.get("hash") or raw.get("sha256_hex") if isinstance(sha, str): sha = sha.strip().lower() or None else: sha = None notes = raw.get("notes") if isinstance(notes, str): notes = notes.strip() or None else: notes = None return { "image": raw.get("image") or _DEFAULT_LORA_IMAGE, "title": title, "adapter_name": adapter, "repo": str(repo), "weights": str(weights), "default_prompt": prompt, "default_weight": default_weight, "admin_approved": _as_bool(raw.get("admin_approved"), False), "active": active, "compatible_with_playground": _as_bool( raw.get("compatible_with_playground"), True ), "notes": notes, "known_triggers": _parse_trigger_list( raw.get("known_triggers") or raw.get("triggers") ), "sha256": sha, } def _serialize_entry(e: dict) -> dict: out = { "title": e["title"], "adapter_name": e["adapter_name"], "repo": e["repo"], "weights": e["weights"], "default_prompt": e.get("default_prompt"), "default_weight": float(e.get("default_weight", 1.0)), "admin_approved": bool(e.get("admin_approved", False)), "active": bool(e.get("active", True)), "compatible_with_playground": bool(e.get("compatible_with_playground", True)), } notes = (e.get("notes") or "").strip() if isinstance(e.get("notes"), str) else e.get("notes") if notes: out["notes"] = notes triggers = _parse_trigger_list(e.get("known_triggers")) if triggers: out["known_triggers"] = triggers if e.get("sha256"): out["sha256"] = e["sha256"] if e.get("image") and e["image"] != _DEFAULT_LORA_IMAGE: out["image"] = e["image"] return out def _read_catalog_file(path: str | None = None) -> list[dict]: catalog_path = Path(path or PERSISTENT_LORA_CATALOG_PATH) if not catalog_path.is_file(): return [] try: with open(catalog_path, "r", encoding="utf-8") as f: data = json.load(f) except Exception as e: print(f"[lora_registry] Could not read catalog {catalog_path}: {e}") return [] if isinstance(data, dict): items = data.get("loras") or data.get("styles") or data.get("items") or [] elif isinstance(data, list): items = data else: return [] out, seen_titles = [], set() for raw in items: entry = _normalize_catalog_entry(raw) if not entry or entry["title"] in seen_titles: continue seen_titles.add(entry["title"]) out.append(entry) return out def _write_catalog_file(entries: list[dict], path: str | None = None) -> str: catalog_path = Path(path or PERSISTENT_LORA_CATALOG_PATH) catalog_path.parent.mkdir(parents=True, exist_ok=True) payload = { "version": 3, "loras": [_serialize_entry(e) for e in entries], } tmp_path = catalog_path.with_suffix(catalog_path.suffix + ".tmp") with open(tmp_path, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2, ensure_ascii=False) f.write("\n") os.replace(tmp_path, catalog_path) return str(catalog_path) def _seed_catalog_if_needed(path: str | None = None) -> list[dict]: catalog_path = Path(path or PERSISTENT_LORA_CATALOG_PATH) entries = _read_catalog_file(str(catalog_path)) if entries: dirty = False for e in entries: if not e.get("sha256"): digest = _compute_entry_sha256(e) if digest: e["sha256"] = digest dirty = True if dirty: try: _write_catalog_file(entries, str(catalog_path)) except Exception as ex: print(f"[lora_registry] Could not backfill sha256: {ex}") return entries seeded = [] for raw in _SEED_LORA_STYLES: entry = _normalize_catalog_entry(raw) if not entry: continue digest = _compute_entry_sha256(entry) if digest: entry["sha256"] = digest seeded.append(entry) try: written = _write_catalog_file(seeded, str(catalog_path)) print(f"[lora_registry] Seeded catalog with {len(seeded)} LoRA(s) → {written}") except Exception as e: print(f"[lora_registry] Could not seed catalog at {catalog_path}: {e}") return seeded def load_external_catalog(path: str | None = None) -> list[dict]: global EXTERNAL_LORA_STYLES, LORA_STYLES with _CATALOG_LOCK: EXTERNAL_LORA_STYLES = _seed_catalog_if_needed(path) LORA_STYLES = [ e for e in EXTERNAL_LORA_STYLES if e.get("active", True) and e.get("compatible_with_playground", True) ] print( f"[lora_registry] Catalog: {len(LORA_STYLES)} active / " f"{len(EXTERNAL_LORA_STYLES)} total from " f"{path or PERSISTENT_LORA_CATALOG_PATH}" ) return list(EXTERNAL_LORA_STYLES) def reload_catalog_from_disk(path: str | None = None) -> list[dict]: """Re-read JSON from disk into memory without re-seeding over an existing file. Used on every browser page load so newly saved catalog entries (and active/inactive edits) appear without restarting the Space. """ global EXTERNAL_LORA_STYLES, LORA_STYLES catalog_path = Path(path or PERSISTENT_LORA_CATALOG_PATH) with _CATALOG_LOCK: if catalog_path.is_file(): entries = _read_catalog_file(str(catalog_path)) # Only fall back to seed when the file is missing, not when it is # a deliberate empty list — empty list means "show nothing extra". EXTERNAL_LORA_STYLES = entries else: EXTERNAL_LORA_STYLES = _seed_catalog_if_needed(str(catalog_path)) LORA_STYLES = [ e for e in EXTERNAL_LORA_STYLES if e.get("active", True) and e.get("compatible_with_playground", True) ] print( f"[lora_registry] Reloaded catalog: {len(LORA_STYLES)} active / " f"{len(EXTERNAL_LORA_STYLES)} total from {catalog_path}" ) return list(EXTERNAL_LORA_STYLES) def refresh_catalog_ui(dynamic_loras_state, currently_selected=None): """Reload disk catalog and return Gradio updates for selector + remove list. Drops session-only entries whose title now exists in the active catalog (they were saved). Preserves still-valid checkbox selections. """ reload_catalog_from_disk() dynamic_loras = dict(dynamic_loras_state or {}) active_titles = {s["title"] for s in _active_catalog()} # Session copies that were persisted no longer need to live in gr.State. stale_keys = [ k for k, v in dynamic_loras.items() if v.get("title") in active_titles ] for k in stale_keys: dynamic_loras.pop(k, None) choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] choice_set = set(choices) new_sel = [t for t in (currently_selected or []) if t in choice_set] save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() return ( gr.update(choices=choices, value=new_sel), gr.update(choices=save_choices), gr.update(choices=remove_choices), dynamic_loras, ) load_external_catalog() def _active_catalog() -> list[dict]: """Entries shown in the main selector: active and playground-compatible.""" return [ e for e in EXTERNAL_LORA_STYLES if e.get("active", True) and e.get("compatible_with_playground", True) ] def _norm_repo_key(repo: str | None) -> str: """Normalize repo/path keys for duplicate comparison. Local absolute paths are resolved so macOS /var vs /private/var (and trailing slashes) still match. HF ids are lower-cased lightly only for exact string compare after strip. """ if not repo: return "" repo = str(repo).strip() if repo.startswith("/"): try: return str(Path(repo).resolve()) except OSError: return repo.rstrip("/") return repo def _find_duplicate(entries: list[dict], *, title=None, sha256=None, repo=None, weights=None): if sha256: sha = str(sha256).lower() for e in entries: if e.get("sha256") and e["sha256"].lower() == sha: return e if title: for e in entries: if e.get("title") == title: return e if repo and weights: repo_key = _norm_repo_key(repo) weight_key = str(weights).strip() for e in entries: if ( _norm_repo_key(e.get("repo")) == repo_key and str(e.get("weights") or "").strip() == weight_key ): return e return None def get_all_styles(dynamic_loras): all_styles = list(_active_catalog()) for dynamic_lora in (dynamic_loras or {}).values(): all_styles.append(dynamic_lora) return all_styles def get_selectable_styles(dynamic_loras): return [s for s in get_all_styles(dynamic_loras) if s.get("adapter_name") is not None] def get_style_by_title(title, dynamic_loras): for style in get_all_styles(dynamic_loras): if style["title"] == title: return style return None def get_style_by_adapter_name(adapter_name, dynamic_loras): for style in get_all_styles(dynamic_loras): if style["adapter_name"] == adapter_name: return style return None def format_selected_lora_details(selected_titles, dynamic_loras) -> str: """Plain-text detail block for currently ticked LoRAs (repo/triggers/notes).""" styles = [ get_style_by_title(t, dynamic_loras) for t in (selected_titles or []) if get_style_by_title(t, dynamic_loras) ] if not styles: return "" blocks = [] for s in styles: lines = [ f"**{s['title']}**", f"- repo: `{s.get('repo')}`", f"- weights: `{s.get('weights')}`", f"- default weight: {s.get('default_weight', 1.0)}", ] if s.get("default_prompt"): lines.append(f"- default prompt: {s['default_prompt']}") triggers = _parse_trigger_list(s.get("known_triggers")) if triggers: lines.append("- known triggers: " + ", ".join(f"`{t}`" for t in triggers)) if s.get("notes"): lines.append(f"- notes: {s['notes']}") blocks.append("\n".join(lines)) return "\n\n".join(blocks) _PROMPT_SEP = "\n\n" _PROMPT_SET_PREFIX = "__set__:" def _style_default_prompt(style) -> str: return (str(style.get("default_prompt") or "")).strip() def _prompt_set_key(titles) -> str: return _PROMPT_SET_PREFIX + "\x1f".join(titles or []) def _capture_lora_prompts(prev_titles, current_text, prompt_memory): """Persist the on-screen LoRA prompt box into session memory.""" mem = dict(prompt_memory or {}) prev = list(prev_titles or []) if not prev: return mem text = "" if current_text is None else str(current_text) # Always remember the exact previous selection's combined text. mem[_prompt_set_key(prev)] = text parts = text.split(_PROMPT_SEP) if len(parts) == len(prev): for title, part in zip(prev, parts): mem[str(title)] = part elif len(prev) == 1: mem[str(prev[0])] = text return mem def _build_lora_prompt(new_titles, prompt_memory, dynamic_loras, prev_titles=None): """Rebuild the combined LoRA prompt, preserving edits when possible.""" mem = dict(prompt_memory or {}) titles = list(new_titles or []) if not titles: return "", mem set_key = _prompt_set_key(titles) if set_key in mem: return mem[set_key], mem prev = list(prev_titles or []) prev_key = _prompt_set_key(prev) if prev else None prev_set = set(prev) new_set = set(titles) # Pure add onto a previous (possibly freely-edited) combined prompt. if prev_key and prev_key in mem and prev_set and prev_set.issubset(new_set) and new_set != prev_set: base = mem[prev_key] extras = [] for title in titles: if title in prev_set: continue if title in mem: frag = mem[title] else: style = get_style_by_title(title, dynamic_loras) or {} frag = _style_default_prompt(style) mem[title] = frag frag = (frag or "").strip() if frag: extras.append(frag) parts = [] if (base or "").strip(): parts.append(base.rstrip()) parts.extend(extras) combined = _PROMPT_SEP.join(parts) mem[set_key] = combined return combined, mem # Pure remove: if the previous combined text still splits cleanly into one # fragment per previous title, drop the removed titles' fragments. # Otherwise the user free-edited the box — keep that text rather than # guessing from stale per-title defaults. if prev_key and prev_key in mem and new_set and new_set.issubset(prev_set) and new_set != prev_set: prev_text = mem[prev_key] prev_parts = prev_text.split(_PROMPT_SEP) if len(prev_parts) == len(prev): title_to_part = dict(zip(prev, prev_parts)) parts = [] for title in titles: frag = (title_to_part.get(title) or mem.get(title) or "").strip() if frag: parts.append(frag) combined = _PROMPT_SEP.join(parts) else: combined = prev_text mem[set_key] = combined return combined, mem # Default path: per-title memory, else catalog default. parts = [] for title in titles: if title in mem and not str(title).startswith(_PROMPT_SET_PREFIX): frag = mem[title] else: style = get_style_by_title(title, dynamic_loras) or {} frag = _style_default_prompt(style) mem[title] = frag frag = (frag or "").strip() if frag: parts.append(frag) combined = _PROMPT_SEP.join(parts) mem[set_key] = combined return combined, mem def update_weight_sliders( selected_titles, dynamic_loras, weight_memory=None, prev_selected=None, prompt_memory=None, current_lora_prompt=None, *current_slider_values, ): """Show/hide weight sliders for the current LoRA selection. Preserves user-adjusted weights and LoRA prompt text: - Reads live slider values for the previous selection into `weight_memory` - Reuses remembered weights for still-selected (or re-selected) titles - Only new titles fall back to catalog `default_weight` - Same idea for the editable LoRA prompt box (per-title + set memory) """ selected_styles = [] for t in (selected_titles or []): style = get_style_by_title(t, dynamic_loras) if style is not None: selected_styles.append(style) memory = dict(weight_memory or {}) prev = list(prev_selected or []) # Capture current on-screen slider values before rebuilding slots. for i, title in enumerate(prev): if i >= len(current_slider_values): break val = current_slider_values[i] if val is None or title is None: continue try: memory[str(title)] = float(val) except (TypeError, ValueError): pass # Capture current LoRA prompt box before selection rebuild. p_memory = _capture_lora_prompts(prev, current_lora_prompt, prompt_memory) slider_updates = [] for i in range(MAX_LORA_SLOTS): if i < len(selected_styles): style = selected_styles[i] title = style["title"] default_w = float(style.get("default_weight", 1.0)) weight = memory.get(title, default_w) try: weight = float(weight) except (TypeError, ValueError): weight = default_w memory[title] = weight slider_updates.append(gr.update( visible=True, interactive=True, label=f"{title} — weight", value=weight, )) else: # Keep label stable when hiding — fewer DOM thrash / stuck-progress cases. slider_updates.append(gr.update( visible=False, interactive=False, value=1.0, )) new_selected = [s["title"] for s in selected_styles] combined, p_memory = _build_lora_prompt( new_selected, p_memory, dynamic_loras, prev_titles=prev, ) # Keep the box always visible so selection races can't hide it while the # prompt is still applied at generate time. Empty when nothing selected. lora_prompt_update = gr.update( value=combined or "", visible=True, interactive=True, ) details = format_selected_lora_details(new_selected, dynamic_loras) if not details: details = ( "*Tick one or more LoRAs above to see full repo paths, " "known triggers, and notes.*" ) # Always keep the Advanced accordion body populated (no visibility toggle). details_update = gr.update(value=details) return slider_updates + [ lora_prompt_update, details_update, memory, new_selected, p_memory, ] def _looks_like_local_path(value: str) -> bool: return bool(value) and value.startswith("/") def _is_weight_filename(name: str | None) -> bool: if not name: return False lower = str(name).strip().lower() return any(lower.endswith(ext) for ext in _WEIGHT_EXTS) def _split_hf_repo_and_weight(repo_id: str, weight_name: str | None) -> tuple[str, str | None]: """Split owner/repo[/nested/weight.safetensors] into hub repo id + weight path. Supports nested weights inside the repo, e.g.: user/repo/sub/dir/model.safetensors -> repo=user/repo, weight=sub/dir/model.safetensors user/repo + weight=sub/dir/model.safetensors (unchanged) """ repo = (repo_id or "").strip().strip("/") weight = weight_name.strip() if weight_name and str(weight_name).strip() else None # If weight already given: hub repo is owner/name; optional extra path # segments are a subfolder prefix (user/repo/sub + file.safetensors). if weight: weight = weight.lstrip("/") parts = [p for p in repo.split("/") if p] if len(parts) >= 2: hub = f"{parts[0]}/{parts[1]}" extra = parts[2:] # user/repo/subfolder + model.safetensors -> subfolder/model.safetensors if extra and not _is_weight_filename(extra[-1]): prefix = "/".join(extra) if not (weight == prefix or weight.startswith(prefix + "/")): weight = f"{prefix}/{weight}" return hub, weight return repo, weight parts = [p for p in repo.split("/") if p] if len(parts) <= 2: return repo, None # owner/repo/ owner, name, *rest = parts hub_repo = f"{owner}/{name}" rest_path = "/".join(rest) # user/repo/file.safetensors OR user/repo/sub/file.safetensors if _is_weight_filename(rest[-1]): return hub_repo, rest_path # user/repo/subfolder (prefix inside repo; weight still unknown) # Keep as repo + None so auto-detect can filter siblings under this prefix. return hub_repo, None if not rest_path else f"{rest_path}/" # trailing slash = prefix marker def _resolve_local_lora(path_str: str, weight_name: str | None): path = Path(os.path.expanduser(path_str)).resolve() requested = weight_name.strip() if weight_name and weight_name.strip() else None if path.is_file(): if path.suffix.lower() not in _WEIGHT_EXTS: raise ValueError(f"Not a LoRA weight file: {path.name}") # Keep nested filename only for local load_lora_weights(dir, weight_name=file) return str(path.parent), path.name, path.stem if not path.is_dir(): raise FileNotFoundError(f"Local path not found: {path}") if requested: # Allow nested relative weight paths: sub/dir/model.safetensors candidate = (path / requested).resolve() try: candidate.relative_to(path) except ValueError as e: raise FileNotFoundError( f"Weight path escapes directory {path}: {requested}" ) from e if not candidate.is_file(): available = sorted( str(p.relative_to(path)) for p in path.rglob("*") if p.is_file() and p.suffix.lower() in _WEIGHT_EXTS )[:20] raise FileNotFoundError( f"'{requested}' not under {path}. Available: {', '.join(available) or 'None'}" ) # diffusers local: repo=dir containing file tree root we pass, weights=relpath rel = str(candidate.relative_to(path)).replace("\\", "/") return str(path), rel, Path(rel).stem for name in _DEFAULT_WEIGHT_CANDIDATES: if (path / name).is_file(): return str(path), name, path.name # Prefer top-level weights; fall back to a single nested weight if unique. top = sorted( p.name for p in path.iterdir() if p.is_file() and p.suffix.lower() in _WEIGHT_EXTS ) if len(top) == 1: return str(path), top[0], Path(top[0]).stem if top: raise FileNotFoundError( f"Multiple weights in {path}; set Weight filename. Available: {', '.join(top)}" ) nested = sorted( str(p.relative_to(path)).replace("\\", "/") for p in path.rglob("*") if p.is_file() and p.suffix.lower() in _WEIGHT_EXTS ) if len(nested) == 1: return str(path), nested[0], Path(nested[0]).stem if not nested: raise FileNotFoundError(f"No .safetensors/.bin weights found in {path}") raise FileNotFoundError( f"Multiple nested weights in {path}; set Weight path. Available: {', '.join(nested[:20])}" ) def _resolve_hf_lora(repo_id: str, weight_name: str | None): from huggingface_hub import model_info hub_repo, weight_or_prefix = _split_hf_repo_and_weight(repo_id, weight_name) # Trailing slash marks "directory prefix inside repo" from user/repo/subfolder prefix = None actual_weight = weight_or_prefix if actual_weight and actual_weight.endswith("/") and not _is_weight_filename(actual_weight): prefix = actual_weight.lstrip("/") actual_weight = None elif actual_weight: actual_weight = actual_weight.lstrip("/") info = model_info(hub_repo) siblings = list(info.siblings or []) def _weight_siblings(pref: str | None = None): out = [] for f in siblings: name = getattr(f, "filename", None) or "" if not name.endswith(_WEIGHT_EXTS): continue if pref and not name.startswith(pref): continue out.append(name) return out if not actual_weight: # Auto-pick under optional subfolder prefix. search_prefix = prefix or "" for name in _DEFAULT_WEIGHT_CANDIDATES: candidate = f"{search_prefix}{name}" if search_prefix else name if any(getattr(f, "filename", None) == candidate for f in siblings): actual_weight = candidate break if not actual_weight: available = _weight_siblings(search_prefix or None) # If prefix was a folder and defaults missing, unique weight under prefix if len(available) == 1: actual_weight = available[0] elif not available and not search_prefix: available = _weight_siblings(None) if len(available) == 1: actual_weight = available[0] if not actual_weight: shown = available[:30] if available else _weight_siblings(None)[:30] where = f" under '{search_prefix.rstrip('/')}'" if search_prefix else "" raise FileNotFoundError( f"No weight found in {hub_repo}{where}. " f"Available: {', '.join(shown) or 'None'}" ) # Validate nested path exists in repo file list when possible sibling_names = {getattr(f, "filename", None) for f in siblings} if actual_weight not in sibling_names: # allow if list incomplete; still try exact match after strip alt = actual_weight.lstrip("./") if alt in sibling_names: actual_weight = alt else: available = _weight_siblings(None) # helpful: show nested matches by basename base = Path(actual_weight).name nested_hits = [a for a in available if a == actual_weight or a.endswith("/" + base)] hint = nested_hits[:10] if nested_hits else available[:20] raise FileNotFoundError( f"Weight '{actual_weight}' not in {hub_repo}. " f"Try nested path like 'subfolder/{base}'. Available: {', '.join(hint) or 'None'}" ) sha = None for sib in siblings: if getattr(sib, "filename", None) == actual_weight: lfs = getattr(sib, "lfs", None) or {} if isinstance(lfs, dict): sha = lfs.get("sha256") or lfs.get("oid") break display = Path(actual_weight).stem if actual_weight else hub_repo.split("/")[-1] return hub_repo, actual_weight, display, (str(sha).lower() if sha else None) def session_custom_lora_titles(dynamic_loras) -> list[str]: return [s["title"] for s in (dynamic_loras or {}).values() if s.get("title")] def removable_catalog_titles() -> list[str]: return [ e["title"] for e in EXTERNAL_LORA_STYLES if not e.get("admin_approved", False) ] def _unique_adapter_name(base: str) -> str: existing = { s["adapter_name"] for s in EXTERNAL_LORA_STYLES if s.get("adapter_name") } | set(LOADED_ADAPTERS) if base not in existing: return base for i in range(2, 1000): candidate = f"{base}_{i}" if candidate not in existing: return candidate return f"{base}_{uuid.uuid4().hex[:6]}" def _empty_add_result(msg, dynamic_loras): save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() return ( msg, gr.update(), dynamic_loras, gr.update(choices=save_choices), gr.update(choices=remove_choices), ) def add_custom_lora(repo_id, weight_name, adapter_name, dynamic_loras_state): dynamic_loras = dict(dynamic_loras_state or {}) if not repo_id or not repo_id.strip(): return _empty_add_result( "Please enter a HuggingFace repo ID or a local path " "(e.g. /loras-flux/my.safetensors).", dynamic_loras, ) repo_id = repo_id.strip() requested_name = adapter_name.strip() if adapter_name and adapter_name.strip() else None try: sha = None if _looks_like_local_path(repo_id): resolved_repo, actual_weight, auto_name = _resolve_local_lora(repo_id, weight_name) source_label = resolved_repo sha = _file_sha256(Path(resolved_repo) / actual_weight) else: resolved_repo, actual_weight, auto_name, sha = _resolve_hf_lora(repo_id, weight_name) source_label = resolved_repo # Re-read disk first so "already in catalog" matches what new browsers should see. reload_catalog_from_disk() dup = _find_duplicate( EXTERNAL_LORA_STYLES, sha256=sha, repo=resolved_repo, weights=actual_weight, title=None, ) if dup: where = "active catalog" if dup.get("active", True) else "archived catalog" # Always refresh selector choices so the existing entry becomes visible # (fixes "saved but stuck / not listed in a new session"). choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() is_active = bool(dup.get("active", True)) msg = ( f"⚠️ Already in {where} as '{dup['title']}'" + (f" (sha256 {dup['sha256'][:12]}…)" if dup.get("sha256") else "") ) if is_active: msg += ". Catalog list refreshed — you can select it above." sel_upd = gr.update(choices=choices, value=[dup["title"]]) else: msg += ". It is archived (active=false); set active=true in JSON to show it." sel_upd = gr.update(choices=choices) return ( msg, sel_upd, dynamic_loras, gr.update(choices=save_choices), gr.update(choices=remove_choices), ) for s in dynamic_loras.values(): if sha and s.get("sha256") and s["sha256"] == sha: return ( f"⚠️ Already added this session as '{s['title']}' (same sha256).", gr.update(), dynamic_loras, gr.update(choices=session_custom_lora_titles(dynamic_loras)), gr.update(choices=session_custom_lora_titles(dynamic_loras) + removable_catalog_titles()), ) if s.get("repo") == resolved_repo and s.get("weights") == actual_weight: return ( f"⚠️ Already added this session as '{s['title']}'.", gr.update(), dynamic_loras, gr.update(choices=session_custom_lora_titles(dynamic_loras)), gr.update(choices=session_custom_lora_titles(dynamic_loras) + removable_catalog_titles()), ) base_name = _sanitize_name(requested_name or auto_name, fallback="custom") static_names = { s["adapter_name"] for s in EXTERNAL_LORA_STYLES if s.get("adapter_name") } final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}" while final_adapter_name in static_names or final_adapter_name in LOADED_ADAPTERS: final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}" custom_style = { "image": _DEFAULT_LORA_IMAGE, "title": f"Custom: {base_name}", "adapter_name": final_adapter_name, "repo": resolved_repo, "weights": actual_weight, "default_prompt": None, "default_weight": 1.0, "admin_approved": False, "active": True, "compatible_with_playground": True, "notes": None, "known_triggers": [], "sha256": sha, "session_only": True, } dynamic_loras[final_adapter_name] = custom_style new_choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() hash_note = f", sha256={sha[:12]}…" if sha else "" return ( f"✅ Added (session only): {base_name} from {source_label} " f"({actual_weight}{hash_note}). Try it, then save or remove below.", gr.update(choices=new_choices), dynamic_loras, gr.update(choices=save_choices, value=custom_style["title"]), gr.update(choices=remove_choices, value=custom_style["title"]), ) except Exception as e: return _empty_add_result(f"❌ Failed: {e}", dynamic_loras) def remove_lora(selected_title, dynamic_loras_state, currently_selected): dynamic_loras = dict(dynamic_loras_state or {}) if not selected_title: choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] remove_choices = session_custom_lora_titles(dynamic_loras) + removable_catalog_titles() return ( "Pick a LoRA to remove.", gr.update(choices=choices, value=currently_selected or []), gr.update(choices=session_custom_lora_titles(dynamic_loras)), gr.update(choices=remove_choices), dynamic_loras, ) session_key = None for k, v in dynamic_loras.items(): if v.get("title") == selected_title: session_key = k break if session_key is not None: dynamic_loras.pop(session_key, None) new_sel = [t for t in (currently_selected or []) if t != selected_title] choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() return ( f"🗑️ Removed session LoRA '{selected_title}'.", gr.update(choices=choices, value=new_sel), gr.update(choices=save_choices, value=None), gr.update(choices=remove_choices, value=None), dynamic_loras, ) with _CATALOG_LOCK: entries = _read_catalog_file() global EXTERNAL_LORA_STYLES EXTERNAL_LORA_STYLES = list(entries) idx = next((i for i, e in enumerate(entries) if e["title"] == selected_title), None) if idx is None: choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] remove_choices = session_custom_lora_titles(dynamic_loras) + removable_catalog_titles() return ( f"❌ '{selected_title}' not found in session or catalog.", gr.update(choices=choices), gr.update(choices=session_custom_lora_titles(dynamic_loras)), gr.update(choices=remove_choices), dynamic_loras, ) if entries[idx].get("admin_approved", False): choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] remove_choices = session_custom_lora_titles(dynamic_loras) + removable_catalog_titles() return ( f"❌ '{selected_title}' is admin-approved and cannot be removed from the UI. " f"Set active=false in the JSON to archive it.", gr.update(choices=choices), gr.update(choices=session_custom_lora_titles(dynamic_loras)), gr.update(choices=remove_choices), dynamic_loras, ) entries.pop(idx) try: path = _write_catalog_file(entries) except Exception as e: return ( f"❌ Failed to write catalog: {e}", gr.update(), gr.update(), gr.update(), dynamic_loras, ) EXTERNAL_LORA_STYLES = list(entries) LORA_STYLES[:] = [ e for e in EXTERNAL_LORA_STYLES if e.get("active", True) and e.get("compatible_with_playground", True) ] new_sel = [t for t in (currently_selected or []) if t != selected_title] choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() return ( f"🗑️ Removed catalog LoRA '{selected_title}' → {path}", gr.update(choices=choices, value=new_sel), gr.update(choices=save_choices, value=None), gr.update(choices=remove_choices, value=None), dynamic_loras, ) def _filter_selector_value(currently_selected, choices, *, rename_from=None, rename_to=None): """Keep CheckboxGroup value valid after choices change. Optionally rename one selected title (session "Custom: x" → catalog "x"). """ choice_set = set(choices or []) out = [] seen = set() for t in currently_selected or []: mapped = rename_to if (rename_from is not None and t == rename_from) else t if mapped in choice_set and mapped not in seen: out.append(mapped) seen.add(mapped) return out def save_session_lora_to_catalog( selected_title, catalog_title, default_weight, default_prompt, dynamic_loras_state, known_triggers=None, notes=None, currently_selected=None, ): dynamic_loras = dict(dynamic_loras_state or {}) save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() selector_choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] cur_sel = list(currently_selected or []) def _fail(msg): # Always re-assert a valid value so a stale Custom: title can't brick the UI. safe_sel = _filter_selector_value(cur_sel, selector_choices) return ( msg, gr.update(choices=selector_choices, value=safe_sel), gr.update(choices=save_choices), gr.update(choices=remove_choices), dynamic_loras, ) if not selected_title: return _fail("Pick a session LoRA to save (add one above first).") source = None source_key = None for k, style in dynamic_loras.items(): if style.get("title") == selected_title: source = style source_key = k break if source is None: return _fail(f"❌ '{selected_title}' is not a session custom LoRA.") title = (catalog_title or "").strip() or selected_title.removeprefix("Custom: ").strip() if not title: title = source.get("adapter_name") or "Custom LoRA" try: weight = float(default_weight) if default_weight is not None else float( source.get("default_weight", 1.0) ) except (TypeError, ValueError): weight = 1.0 prompt = (default_prompt or "").strip() or source.get("default_prompt") or None triggers = _parse_trigger_list(known_triggers) if not triggers: triggers = _parse_trigger_list(source.get("known_triggers")) note_text = (notes or "").strip() or source.get("notes") or None adapter_base = _sanitize_name( title, fallback=_sanitize_name(source.get("adapter_name", "custom")) ) sha = source.get("sha256") or _compute_entry_sha256(source) with _CATALOG_LOCK: entries = _read_catalog_file() global EXTERNAL_LORA_STYLES EXTERNAL_LORA_STYLES = list(entries) dup = _find_duplicate( entries, sha256=sha, title=title, repo=source["repo"], weights=source["weights"], ) if dup and dup["title"] != title: return _fail( f"❌ Duplicate of existing catalog entry '{dup['title']}'" + (" (sha256 match)" if sha and dup.get("sha256") == sha else "") ) if dup and dup.get("admin_approved", False) and dup["title"] == title: return _fail( f"❌ '{title}' is admin-approved — edit the JSON directly to change it." ) existing_idx = next((i for i, e in enumerate(entries) if e["title"] == title), None) if existing_idx is None: adapter_name = _unique_adapter_name(adapter_base) prev_approved = False prev_active = True else: adapter_name = entries[existing_idx].get("adapter_name") or _unique_adapter_name(adapter_base) prev_approved = bool(entries[existing_idx].get("admin_approved", False)) prev_active = bool(entries[existing_idx].get("active", True)) prev_compatible = True if existing_idx is not None: prev_compatible = bool( entries[existing_idx].get("compatible_with_playground", True) ) entry = { "image": source.get("image") or _DEFAULT_LORA_IMAGE, "title": title, "adapter_name": adapter_name, "repo": source["repo"], "weights": source["weights"], "default_prompt": prompt, "default_weight": weight, "admin_approved": prev_approved, "active": prev_active, "compatible_with_playground": prev_compatible, "notes": note_text, "known_triggers": triggers, "sha256": sha, } if existing_idx is None: entries.append(entry) action = "Saved" else: entries[existing_idx] = entry action = "Updated" try: path = _write_catalog_file(entries) except Exception as e: return _fail(f"❌ Failed to write catalog: {e}") EXTERNAL_LORA_STYLES = list(entries) LORA_STYLES[:] = [ e for e in EXTERNAL_LORA_STYLES if e.get("active", True) and e.get("compatible_with_playground", True) ] if source_key is not None: dynamic_loras.pop(source_key, None) selector_choices = [s["title"] for s in get_selectable_styles(dynamic_loras)] save_choices = session_custom_lora_titles(dynamic_loras) remove_choices = save_choices + removable_catalog_titles() # Session title disappears; remap selection to the new catalog title so the # CheckboxGroup never keeps "Custom: …" against choices that only have "…". new_sel = _filter_selector_value( cur_sel, selector_choices, rename_from=selected_title, rename_to=title, ) hash_note = f", sha256={sha[:12]}…" if sha else "" trig_note = f", {len(triggers)} trigger(s)" if triggers else "" return ( f"✅ {action} '{title}' → {path} (weight={weight}{hash_note}{trig_note})", gr.update(choices=selector_choices, value=new_sel), gr.update(choices=save_choices, value=None), gr.update(choices=remove_choices, value=None), dynamic_loras, ) def fill_catalog_save_form(selected_title, dynamic_loras_state): empty = gr.update(), gr.update(), gr.update(), gr.update(), gr.update() if not selected_title: return empty style = None for s in (dynamic_loras_state or {}).values(): if s.get("title") == selected_title: style = s break if not style: return empty suggested = selected_title.removeprefix("Custom: ").strip() or selected_title triggers = _parse_trigger_list(style.get("known_triggers")) return ( gr.update(value=suggested), gr.update(value=float(style.get("default_weight", 1.0))), gr.update(value=style.get("default_prompt") or ""), gr.update(value="\n".join(triggers)), gr.update(value=style.get("notes") or ""), )