import os import gc import csv import time import random import uuid import zipfile import gradio as gr import spaces import torch import numpy as np from PIL import Image # ── Local modules — single source of truth for each concern ───────────────── from config import ( MODEL_VARIANT, MODEL_REPO, MAX_SEED, MAX_LORA_SLOTS, PERSISTENT_LORA_CATALOG_PATH, UNCENSORED_TE_REPO, UNCENSORED_TE_FILE, ) from ui_theme import orange_red_theme from upscale import UPSCALE_MODELS, apply_realesrgan from lora_registry import ( LORA_STYLES, LOADED_ADAPTERS, get_selectable_styles, get_style_by_title, update_weight_sliders, add_custom_lora, save_session_lora_to_catalog, fill_catalog_save_form, remove_lora, removable_catalog_titles, refresh_catalog_ui, ) from image_utils import ( fix_orientation, compute_canvas_dimensions, fit_to_canvas, on_base_image_change, make_solid_base_image, collect_reference_images, reference_info_text, compact_reference_slots, move_base_down, move_ref1_up, move_ref1_down, move_ref2_up, move_ref2_down, move_ref3_up, process_images, reencode_upload, send_editor_to_base, send_editor_to_reference, load_heic_to_editor, send_output_to_base, send_output_to_reference, save_with_metadata, build_pnginfo, push_pil_to_base, push_pil_to_reference, extend_editor_canvas, render_extend_schematic, resolve_output_download_path, latest_gallery_download_path, ) from control_tools import ( generate_depthmap, detect_pose, render_pose_skeleton, render_pose_overlay, move_joint, hide_joint, clear_all_joints, default_pose_template, person_choices, parse_person_idx, joint_name_to_index, OPENPOSE_KEYPOINT_NAMES, ) DEFAULT_PROJECT_NAME = "f2klora" MAX_PROJECT_NAME_LEN = 12 def sanitize_project_name(name: str | None) -> str: """Alphanumeric only, max 12 chars. Default f2klora.""" raw = (name or "").strip() cleaned = "".join(c for c in raw if c.isalnum()) cleaned = cleaned[:MAX_PROJECT_NAME_LEN] return cleaned or DEFAULT_PROJECT_NAME def make_run_stamp() -> str: """yymmddhhmmss — no separators.""" return time.strftime("%y%m%d%H%M%S") def make_download_basename(project: str | None, stamp: str | None = None, batch_index: int | None = None) -> str: stem = f"{sanitize_project_name(project)}{stamp or make_run_stamp()}" if batch_index is not None: stem = f"{stem}{int(batch_index):02d}" return stem def _tmp_named(basename: str, ext: str) -> str: ext = ext if ext.startswith(".") else f".{ext}" return f"/tmp/{basename}{ext}" def save_simple_image(image: Image.Image, basename: str | None = None) -> str: """Save image without any metadata for privacy.""" path = _tmp_named(basename or f"gen{uuid.uuid4().hex[:8]}", ".png") image.save(path, format="PNG") return path def _build_full_prompt(prompt, lora_prompt_text, custom_prompt_text) -> str: return "\n".join( p for p in [ (prompt or "").strip(), (lora_prompt_text or "").strip(), (custom_prompt_text or "").strip(), ] if p ) def save_webp_from_image(image: Image.Image, basename: str | None = None) -> str: path = _tmp_named(basename or f"gen{uuid.uuid4().hex[:8]}", ".webp") image.convert("RGB").save(path, format="WEBP", quality=90, method=4) return path def save_webp_from_path(path, basename: str | None = None) -> str | None: from image_utils import _gallery_item_path p = _gallery_item_path(path) or (path if isinstance(path, str) else None) if not p: return None try: img = Image.open(p).convert("RGB") except Exception: return None return save_webp_from_image(img, basename=basename) def save_prompt_txt(text: str | None, basename: str | None = None) -> str | None: text = (text or "").strip() if not text: return None path = _tmp_named(basename or f"prompt{uuid.uuid4().hex[:8]}", ".txt") with open(path, "w", encoding="utf-8") as f: f.write(text) if not text.endswith("\n"): f.write("\n") return path def copy_as_named_png(src_path, basename: str) -> str | None: """Copy an existing PNG to a project+timestamp name for download.""" from image_utils import _gallery_item_path import shutil p = _gallery_item_path(src_path) or (src_path if isinstance(src_path, str) else None) if not p or not os.path.isfile(p): return None dest = _tmp_named(basename, ".png") if os.path.abspath(p) == os.path.abspath(dest): return p try: shutil.copy2(p, dest) return dest except Exception: return p def resolve_download_bundle(selected_path, gallery_value, last_prompt_text, project_name): """PNG + WebP + prompt txt named project+yymmddhhmmss.(png|webp|txt).""" png_src = resolve_output_download_path(selected_path, gallery_value) base = make_download_basename(project_name) png = copy_as_named_png(png_src, base) if png_src else None webp = save_webp_from_path(png_src, basename=base) if png_src else None prompt_file = save_prompt_txt(last_prompt_text, basename=base) return png, webp, prompt_file # ── Download tracking (warn before overwriting undownloaded outputs) ───────── def _norm_img_path(item) -> str | None: from image_utils import _gallery_item_path p = _gallery_item_path(item) if p: return os.path.abspath(p) if isinstance(item, str) and item.strip(): return os.path.abspath(item) if os.path.isabs(item) else item.strip() return None def gallery_image_paths(gallery_value) -> list[str]: paths: list[str] = [] seen: set[str] = set() for item in gallery_value or []: p = _norm_img_path(item) if p and p not in seen: seen.add(p) paths.append(p) return paths def format_download_status(pending_list, gallery_value) -> str: paths = gallery_image_paths(gallery_value) pending = set(pending_list or []) undownloaded = [p for p in paths if p in pending] if not paths: return "*No generated images yet.*" if not undownloaded: return f"✅ All **{len(paths)}** gallery image(s) marked downloaded." names = ", ".join(f"`{os.path.basename(p)}`" for p in undownloaded[:4]) extra = f" +{len(undownloaded) - 4} more" if len(undownloaded) > 4 else "" return ( f"⚠️ **{len(undownloaded)}/{len(paths)}** not downloaded yet: {names}{extra}. " f"Use ⬇️ PNG/WebP or the gallery ↓ icon." ) def sync_download_tracking(gallery_value, pending_list, downloaded_list): """Keep pending in sync with gallery contents. New gallery paths not yet marked downloaded become pending. Paths that left the gallery drop out of pending. """ pending = set(pending_list or []) downloaded = set(downloaded_list or []) paths = gallery_image_paths(gallery_value) path_set = set(paths) pending = {p for p in pending if p in path_set} for p in paths: if p not in downloaded: pending.add(p) pending_out = [p for p in paths if p in pending] # stable order downloaded_out = sorted(downloaded) return pending_out, downloaded_out, format_download_status(pending_out, gallery_value) def _mark_paths_downloaded(paths_to_mark, gallery_value, pending_list, downloaded_list): pending = set(pending_list or []) downloaded = set(downloaded_list or []) marked = [] for raw in paths_to_mark or []: p = _norm_img_path(raw) if not p: # bare filename / URL fragment from gallery JS s = str(raw or "").strip() if not s: continue base = os.path.basename(s.split("?")[0].split("#")[0]) for gp in gallery_image_paths(gallery_value): if os.path.basename(gp) == base or base in gp or gp.endswith(base): p = gp break if not p and base: # still record basename key so status can clear if paths match later p = base if not p: continue downloaded.add(p) pending.discard(p) # also clear any gallery path sharing basename base = os.path.basename(p) for gp in list(pending): if os.path.basename(gp) == base: pending.discard(gp) downloaded.add(gp) marked.append(p) paths = gallery_image_paths(gallery_value) pending_out = [p for p in paths if p in pending] return pending_out, sorted(downloaded), format_download_status(pending_out, gallery_value) def mark_current_output_downloaded(selected_path, gallery_value, pending_list, downloaded_list): """Mark the selected (or latest) gallery image as downloaded.""" p = resolve_output_download_path(selected_path, gallery_value) return _mark_paths_downloaded([p] if p else [], gallery_value, pending_list, downloaded_list) def mark_from_gallery_signal(signal, gallery_value, pending_list, downloaded_list): """Mark download from gallery ↓ icon (JS writes URL/filename into signal).""" raw = (signal or "").strip() if not raw: return ( list(pending_list or []), list(downloaded_list or []), format_download_status(pending_list, gallery_value), gr.update(value=""), ) # JS appends "|timestamp" so repeated clicks still fire .change token = raw.split("|", 1)[0].strip() pending, downloaded, status = _mark_paths_downloaded( [token], gallery_value, pending_list, downloaded_list, ) return pending, downloaded, status, gr.update(value="") def warn_undownloaded_before_generate(pending_list, gallery_value): """Disabled for now — undownloaded tracking will return later. Still clears selected_output_state when chained before generate. """ return None # ── Model load ────────────────────────────────────────────────────────────── # Pipeline class depends on MODEL_VARIANT and is the only thing here that # can't live in config.py (config must stay torch/diffusers-free). if MODEL_VARIANT == "9B-KV": from diffusers import Flux2KleinKVPipeline as _PipeClass else: from diffusers import Flux2KleinPipeline as _PipeClass device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Loading FLUX.2 Klein {MODEL_VARIANT} from {MODEL_REPO}...") pipe = _PipeClass.from_pretrained(MODEL_REPO, torch_dtype=torch.bfloat16).to(device) print(f"Model loaded successfully: FLUX.2 Klein {MODEL_VARIANT}") # ── Replace text encoder with abliterated (uncensored) version ─────────── try: from huggingface_hub import hf_hub_download from safetensors.torch import load_file print(f"Downloading abliterated text encoder from {UNCENSORED_TE_REPO}...") te_path = hf_hub_download(repo_id=UNCENSORED_TE_REPO, filename=UNCENSORED_TE_FILE) print(f"Loading abliterated weights from {te_path}...") state_dict = load_file(te_path) pipe.text_encoder.load_state_dict(state_dict, strict=True) pipe.text_encoder.to(device=device, dtype=torch.bfloat16) print("Abliterated text encoder loaded — safety filters removed.") except Exception as e: print(f"Abliterated text encoder unavailable ({e}) — using stock encoder.") # ── UI helper callbacks ────────────────────────────────────────────────────── def on_canvas_mode_change(mode): """Custom W/H sliders only relevant when mode == Custom.""" is_custom = (mode == "Custom") return gr.update(visible=is_custom), gr.update(visible=is_custom) def on_fit_mode_change(fit_mode): """Pad colour swatch only relevant for Pad (color).""" return gr.update(visible=(fit_mode == "Pad (color)")) def on_batch_vary_change(vary_mode): """Sweep range only relevant for the LoRA sweep mode.""" is_sweep = (vary_mode == "Sweep first LoRA weight") return gr.update(visible=is_sweep), gr.update(visible=is_sweep) def on_gallery_select(evt: gr.SelectData, gallery_value): """Remember which gallery item is selected so Send→* uses it.""" if evt is None or gallery_value is None or evt.index is None: return None try: item = gallery_value[evt.index] except (IndexError, TypeError): return None return item[0] if isinstance(item, (list, tuple)) else item # ── Logging (disabled) ─────────────────────────────────────────────────────── def _spawn_log(*_args, **_kwargs): """No-op — logging intentionally disabled.""" return # ── GPU step (shared by single, batch, and bulk) ───────────────────────────── @spaces.GPU def _infer_gpu( pil_images, prompt, lora_prompt_text, custom_prompt_text, selected_titles, seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, dynamic_loras, *slider_values, progress=gr.Progress(track_tqdm=True), ): if "Best-Face-Swap" in selected_titles: if len(pil_images) < 2: raise gr.Error("Face Swap requires 2 images: a Base image and one Reference image.") if len(pil_images) > 2: gr.Warning("Face Swap uses only the Base image and the first Reference image.") pil_images = pil_images[:2] active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None] weights = list(slider_values[:len(active_styles)]) if not active_styles: pipe.disable_lora() else: for style in active_styles: an = style["adapter_name"] if an not in LOADED_ADAPTERS: try: pipe.load_lora_weights(style["repo"], weight_name=style["weights"], adapter_name=an) LOADED_ADAPTERS.add(an) except Exception as e: raise gr.Error(f"Failed to load {style['title']}: {e}") pipe.set_adapters([s["adapter_name"] for s in active_styles], adapter_weights=[float(w) for w in weights]) full_prompt = "\n".join(p for p in [ (prompt or "").strip(), (lora_prompt_text or "").strip(), (custom_prompt_text or "").strip(), ] if p) width, height = compute_canvas_dimensions(pil_images[0], canvas_mode, custom_width, custom_height) print(f"Generating at: {width}×{height} (canvas={canvas_mode}, fit={canvas_fit_mode})") processed = [fit_to_canvas(img, width, height, canvas_fit_mode, pad_color) for img in pil_images] image_input = processed if len(processed) > 1 else processed[0] try: kwargs = dict(image=image_input, prompt=full_prompt, width=width, height=height, num_inference_steps=steps, generator=torch.Generator(device=device).manual_seed(seed)) if MODEL_VARIANT != "9B-KV": kwargs["guidance_scale"] = guidance_scale image = pipe(**kwargs).images[0] except Exception as e: raise gr.Error(f"Inference failed: {e}") if upscale_factor and upscale_factor != "None": gc.collect(); torch.cuda.synchronize(); torch.cuda.empty_cache() try: image = apply_realesrgan(image, upscale_factor, device) except Exception as e: gr.Warning(f"Upscaling failed, returning {width}×{height} result: {e}") gc.collect(); torch.cuda.empty_cache() return image, seed, width, height # ── Single / batch infer (generator → streams into gr.Gallery) ─────────────── def infer( base_image, ref1, ref2, ref3, prompt, lora_prompt_text, custom_prompt_text, selected_titles, seed, randomize_seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, batch_count, batch_vary, sweep_min, sweep_max, project_name, dynamic_loras, *slider_values, progress=gr.Progress(track_tqdm=True), ): """Generator. Streams a list of PNG paths into the output gallery.""" gc.collect(); torch.cuda.empty_cache() if not isinstance(upscale_factor, str) or upscale_factor not in UPSCALE_MODELS: upscale_factor = "None" if base_image is None: raise gr.Error("Please upload a base image.") reference_images = collect_reference_images(ref1, ref2, ref3) pil_images = process_images(base_image, reference_images) if not pil_images: raise gr.Error("Could not process uploaded images.") selected_titles = selected_titles or [] batch_count = max(1, int(batch_count)) project = sanitize_project_name(project_name) base_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) seeds, weight_overrides = [], [] for i in range(batch_count): if batch_vary == "Sequential seed (+1 each)": seeds.append((base_seed + i) % (MAX_SEED + 1)); weight_overrides.append(None) elif batch_vary == "Sweep first LoRA weight": seeds.append(base_seed) t = i / max(batch_count - 1, 1) weight_overrides.append((0, float(sweep_min) + t * (float(sweep_max) - float(sweep_min)))) else: # "Random seed each run" (default) seeds.append(random.randint(0, MAX_SEED)); weight_overrides.append(None) active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None] results = [] last_seed_text = "" full_prompt = _build_full_prompt(prompt, lora_prompt_text, custom_prompt_text) # One stamp for the whole batch; multi-run items get 00/01/... suffix (no hyphen). run_stamp = make_run_stamp() prompt_base = make_download_basename(project, run_stamp) prompt_file = save_prompt_txt(full_prompt, basename=prompt_base) for i in range(batch_count): sliders = list(slider_values) if weight_overrides[i] is not None: slot, val = weight_overrides[i] if slot < len(sliders): sliders[slot] = val cur_seed = seeds[i] t0 = time.perf_counter() try: image, used_seed, w, h = _infer_gpu( pil_images, prompt, lora_prompt_text, custom_prompt_text, selected_titles, cur_seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, dynamic_loras, *sliders, progress=progress, ) item_base = make_download_basename( project, run_stamp, batch_index=i if batch_count > 1 else None, ) png_path = save_simple_image(image, basename=item_base) webp_path = save_webp_from_image(image, basename=item_base) results.append(png_path) last_seed_text = str(used_seed) # Yield download paths explicitly. Relying only on output_gallery.change # fails on the first generate in a virgin session (buttons stay empty # until some later interaction re-triggers the change chain). yield results, last_seed_text, png_path, webp_path, prompt_file, full_prompt except Exception as e: png = latest_gallery_download_path(results) fail_base = make_download_basename(project) webp = save_webp_from_path(png, basename=fail_base) if png else None png_named = copy_as_named_png(png, fail_base) if png else None yield ( results, f"Batch {i+1}/{batch_count} failed: {e}", png_named, webp, prompt_file, full_prompt, ) # ── Bulk processing (one input image per iteration) ───────────────────────── def _new_bulk_workdir() -> str: sid = uuid.uuid4().hex[:8] path = f"/tmp/bulk_{sid}" os.makedirs(path, exist_ok=True) return path def bulk_infer( input_files, prompt, lora_prompt_text, custom_prompt_text, selected_titles, seed, randomize_seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, dynamic_loras, *slider_values, progress=gr.Progress(), ): """Process each uploaded image as its own GPU call.""" if not input_files: raise gr.Error("Upload at least one image first.") work_dir = _new_bulk_workdir() results = [] succeeded = 0 total = len(input_files) active_styles = [get_style_by_title(t, dynamic_loras) for t in (selected_titles or []) if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None] for i, path in enumerate(input_files): progress(i / total, desc=f"Image {i+1}/{total}") fname = os.path.basename(path) if isinstance(path, str) else f"input_{i}" t0 = time.perf_counter() try: img = fix_orientation(Image.open(path)).convert("RGB") cur_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) image, used_seed, w, h = _infer_gpu( [img], prompt, lora_prompt_text, custom_prompt_text, selected_titles or [], cur_seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, dynamic_loras, *slider_values, progress=progress, ) stem = os.path.splitext(fname)[0] out_path = os.path.join(work_dir, f"{i:03d}_{stem}.png") image.save(out_path, format="PNG") # plain save, no metadata results.append(out_path) succeeded += 1 except Exception as e: print(f"Bulk image {i+1} failed: {e}") return results # ── Custom prompt manager (session-local) ──────────────────────────────────── def add_custom_prompt(name, text, prompts_state, counter_state): prompts = dict(prompts_state); counter = int(counter_state) text = text.strip() if text else "" name = name.strip() if name else "" if not text: return "Please enter some prompt text.", prompts, counter, gr.update(), gr.update(), gr.update() if not name: counter += 1; name = f"Prompt {counter}" if name in prompts: return f"⚠️ '{name}' already exists.", prompts, counter, gr.update(), gr.update(), gr.update() prompts[name] = text choices = list(prompts.keys()) return (f"✅ Saved: '{name}'", prompts, counter, gr.update(choices=choices), gr.update(choices=choices), gr.update(value="", interactive=True)) def delete_custom_prompt(name, currently_selected, prompts_state): prompts = dict(prompts_state) msg = f"🗑️ Deleted: '{name}'" if name and name in prompts else "Nothing to delete." if name and name in prompts: del prompts[name] choices = list(prompts.keys()) new_sel = [n for n in (currently_selected or []) if n in prompts] return (msg, prompts, gr.update(choices=choices, value=new_sel), gr.update(choices=choices, value=None)) def update_custom_prompt_display(selected_names, prompts_state): if not selected_names: return gr.update(value="", visible=False) texts = [prompts_state[n] for n in selected_names if n in prompts_state] if texts: return gr.update(value="\n\n".join(texts), visible=True) return gr.update(value="", visible=False) # ── UI ─────────────────────────────────────────────────────────────────────── # Shared viewport height so Base / Reference / Output feel the same size. # Keep this moderate — oversized Gallery CSS previously split the reference # panel into a huge empty pane + tiny control strip. _IMAGE_BOX_H = 320 _OUTPUT_GALLERY_H = 520 css = f""" #col-container {{ margin: 0 auto; max-width: 1100px; }} #main-title h1 {{ font-size: 2.4em !important; }} .lora-weight-row {{ background: var(--block-background-fill); border-radius: 8px; padding: 4px 12px; margin-bottom: 4px; }} #used_seed textarea {{ min-height: 0 !important; height: 2.2rem !important; }} /* Output gallery: avoid huge empty preview chrome / forced scrollbars */ #output_gallery {{ min-height: {_OUTPUT_GALLERY_H}px; }} #output_gallery .grid-wrap, #output_gallery .gallery-container, #output_gallery .thumbnail-item, #output_gallery .preview-image, #output_gallery img {{ max-height: {_OUTPUT_GALLERY_H - 48}px !important; object-fit: contain !important; }} #output_gallery .preview {{ max-height: {_OUTPUT_GALLERY_H - 24}px !important; overflow: hidden !important; }} .slot-move-row button {{ min-width: 2.4rem !important; }} """ # Enter inserts newline in multi-line textboxes (Gradio default often submits). # Shift+Enter also inserts newline for muscle-memory parity with chat UIs. _TEXTBOX_NEWLINE_JS = """ () => { const isMulti = (el) => { if (!el || el.tagName !== 'TEXTAREA') return false; if (el.closest('#used_seed') || el.closest('#project_name')) return false; return true; }; const onKey = (e) => { if (e.key !== 'Enter' || e.isComposing) return; const t = e.target; if (!isMulti(t)) return; // Always keep newline behaviour; never submit the form from a prompt box. e.stopPropagation(); // Browser already inserts newline on plain Enter in textarea; // for Shift+Enter some hosts swallow it — insert manually if needed. if (e.shiftKey) { e.preventDefault(); const start = t.selectionStart ?? t.value.length; const end = t.selectionEnd ?? start; const v = t.value; t.value = v.slice(0, start) + '\\n' + v.slice(end); const pos = start + 1; t.selectionStart = t.selectionEnd = pos; t.dispatchEvent(new Event('input', { bubbles: true })); } }; document.addEventListener('keydown', onKey, true); } """ # Gradio 6.0: theme/css go on launch(), not Blocks() with gr.Blocks() as demo: custom_prompts_state = gr.State({}) custom_prompt_counter_state = gr.State(0) dynamic_loras_state = gr.State({}) selected_output_state = gr.State(None) last_prompt_state = gr.State("") # Paths still in the gallery that have not been marked downloaded this session. pending_download_state = gr.State([]) downloaded_images_state = gr.State([]) # Filled by JS when the gallery's built-in ↓ icon is clicked. gallery_dl_signal = gr.Textbox( value="", visible=False, elem_id="gallery_dl_signal", ) # Remember per-title LoRA weights / prompts across add/remove so existing # values don't snap back to catalog defaults when the selection changes. lora_weight_memory_state = gr.State({}) lora_prompt_memory_state = gr.State({}) lora_prev_selected_state = gr.State([]) with gr.Column(elem_id="col-container"): gr.Markdown("# **flux2klein lora playground**", elem_id="main-title") gr.Markdown( f"Apply one or more [LoRA](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B) " f"adapters using [FLUX.2-Klein-{MODEL_VARIANT}]({MODEL_REPO}). " f"**Model:** `{MODEL_VARIANT}`" ) with gr.Tabs() as main_tabs: # ── Generate tab ───────────────────────────────────────────────── with gr.Tab("🎨 Generate", id="tab_generate"): # equal_height=False: otherwise the short seed box stretches to match # the tall left column (base + reference). with gr.Row(equal_height=False): with gr.Column(scale=1): base_image = gr.Image( label="Base Image", type="pil", sources=["upload", "clipboard"], height=_IMAGE_BOX_H, elem_id="base_image", ) with gr.Row(elem_classes="slot-move-row"): base_down_btn = gr.Button("↓ Base → Ref1", size="sm") # t2i-style workflows still need a base image; solid colour # often becomes the background for i2i LoRAs used as t2i. with gr.Row(): base_solid_color = gr.ColorPicker( label="Solid base colour", value="#FFFFFF", scale=1, elem_id="base_solid_color", ) make_solid_base_btn = gr.Button( "⬜ Use solid base", size="sm", scale=1, elem_id="make_solid_base_btn", ) gr.Markdown( "*No photo? Use a solid base for t2i-style runs. " "Size follows Custom canvas W×H if set, else 1024×1024.*", ) size_info = gr.Markdown("*No image uploaded yet*") run_button_top = gr.Button( "▶ Generate", variant="primary", size="lg", elem_id="run_button_top", ) # Progressive single-image refs (max 3). Ref2 appears after # Ref1 is set; Ref3 after Ref2. Deleting a middle slot packs # remaining refs upward. Reorder with ↑/↓ (includes Base). ref1 = gr.Image( label="Reference 1 — optional", type="pil", sources=["upload", "clipboard"], height=_IMAGE_BOX_H, elem_id="ref1", ) with gr.Row(elem_classes="slot-move-row"): ref1_up_btn = gr.Button("↑", size="sm", scale=0) ref1_down_btn = gr.Button("↓", size="sm", scale=0) ref2 = gr.Image( label="Reference 2 — optional", type="pil", sources=["upload", "clipboard"], height=_IMAGE_BOX_H, visible=False, elem_id="ref2", ) with gr.Row(elem_classes="slot-move-row"): ref2_up_btn = gr.Button("↑", size="sm", scale=0, visible=False) ref2_down_btn = gr.Button("↓", size="sm", scale=0, visible=False) ref3 = gr.Image( label="Reference 3 — optional", type="pil", sources=["upload", "clipboard"], height=_IMAGE_BOX_H, visible=False, elem_id="ref3", ) with gr.Row(elem_classes="slot-move-row"): ref3_up_btn = gr.Button("↑", size="sm", scale=0, visible=False) reference_info = gr.Markdown("📷 No reference images") gr.Markdown( "*Up to 3 reference images. Next box appears after you fill the previous one. " "Clearing a slot shifts the others up. Use ↑/↓ to reorder Base + refs. " "Face Swap uses Base + Reference 1 only.*" ) prompt = gr.Textbox( label="Prompt", lines=3, max_lines=12, placeholder="Describe the edit, or leave blank for style-only LoRAs. Enter = new line.", ) lora_prompt_display = gr.Textbox( label="LoRA prompts (auto-filled from selection — editable for this run)", interactive=True, visible=True, lines=3, max_lines=16, value="", placeholder="Tick LoRAs above to auto-fill. Edits are kept when you change selection.", info="Filled when you tick LoRAs. Edit freely for the current generate; " "does not change the stored catalog default.", ) custom_prompt_display = gr.Textbox( label="Custom Prompts (auto-appended)", interactive=False, visible=False, lines=3, max_lines=12, ) run_button = gr.Button("▶ Generate", variant="primary", size="lg") with gr.Column(scale=1): output_gallery = gr.Gallery( label="Output", type="filepath", columns=2, rows=1, height=_OUTPUT_GALLERY_H, allow_preview=True, preview=False, object_fit="contain", show_label=True, elem_id="output_gallery", ) with gr.Row(): used_seed = gr.Textbox( label="🌱 Seed used (last run)", interactive=False, lines=1, max_lines=1, elem_id="used_seed", scale=2, ) project_name = gr.Textbox( label="Project short name", value=DEFAULT_PROJECT_NAME, max_lines=1, lines=1, max_length=MAX_PROJECT_NAME_LEN, placeholder=DEFAULT_PROJECT_NAME, info="Max 12 letters/digits. Used in download filenames.", scale=1, elem_id="project_name", ) with gr.Row(): download_png_btn = gr.DownloadButton( label="⬇️ PNG", value=None, variant="primary", size="sm", scale=1, elem_id="download_png_btn", ) download_webp_btn = gr.DownloadButton( label="⬇️ WebP", value=None, variant="secondary", size="sm", scale=1, elem_id="download_webp_btn", ) download_prompt_btn = gr.DownloadButton( label="⬇️ Prompt", value=None, variant="secondary", size="sm", scale=1, elem_id="download_prompt_btn", ) with gr.Row(): send_out_to_base_btn = gr.Button("↩ Send → Base", size="sm") send_out_to_ref_btn = gr.Button("↩ Send → Reference", size="sm") download_status = gr.Markdown( "*No generated images yet.*", elem_id="download_status", ) gr.Markdown( "*Downloads: `project` + `yymmddhhmmss` + `.png/.webp/.txt` " f"(default project `{DEFAULT_PROJECT_NAME}`). " "Click a gallery thumbnail before Download / Send→; otherwise latest. " "Prompt file = last run's combined user + LoRA + custom text. " "PNG/WebP or the gallery ↓ icon marks an image downloaded; " "Generate warns if undownloaded images would be replaced.*" ) with gr.Row(): gr.Markdown("### 🎨 Select LoRA(s)", elem_classes=["lora-heading"]) reload_catalog_btn = gr.Button( "🔄 Reload catalog", size="sm", scale=0, ) lora_selector = gr.CheckboxGroup( choices=[s["title"] for s in get_selectable_styles({})], value=[], label="Active LoRAs — tick one or more", ) catalog_load_status = gr.Markdown("", visible=True) gr.Markdown("#### Weights for selected LoRAs") weight_sliders = [] with gr.Group(): for i in range(MAX_LORA_SLOTS): with gr.Row(elem_classes="lora-weight-row"): weight_sliders.append(gr.Slider( minimum=0.0, maximum=2.0, step=0.05, value=1.0, label=f"LoRA slot {i+1}", visible=False, interactive=True, )) with gr.Accordion("➕ Load Custom LoRA (HF repo or local path)", open=False): gr.Markdown( "Add any FLUX.2-Klein-compatible LoRA from a **HuggingFace repo** " "(`user/repo`) or a **local path** (file or directory), e.g. " "`/loras-flux/my.safetensors` or `/loras-flux/foo/bar/male`. " "**Import is always session-only** — try it first, then optionally " f"save it to the catalog JSON (`{PERSISTENT_LORA_CATALOG_PATH}`). " "Duplicates are blocked by title, repo+filename, and sha256 when available." ) with gr.Row(): lora_repo_id = gr.Textbox( label="HF repo ID or local path", placeholder="user/repo or user/repo/sub/model.safetensors or /loras-flux/my.safetensors", info="Nested HF paths OK: user/repo/folder/model.safetensors", ) with gr.Row(): lora_weight_name = gr.Textbox( label="Weight path inside repo (optional)", placeholder="subfolder/model.safetensors", info="Use for nested files if not included in the repo field.", ) lora_adapter_name = gr.Textbox(label="Adapter name (optional)", placeholder="my-lora") with gr.Row(): add_lora_btn = gr.Button("Add LoRA (session only)", variant="primary") lora_status = gr.Textbox(label="Status", interactive=False) gr.Markdown("#### 💾 Save tried LoRA to catalog") gr.Markdown( "After testing a session LoRA, save it here so it appears for everyone " "on the next load. UI saves are **not** admin-approved; set " "`admin_approved: true` in the JSON yourself. Set `active: false` to " "archive/hide without deleting." ) catalog_save_select = gr.Dropdown( label="Session LoRA to save", choices=[], value=None, interactive=True, ) with gr.Row(): catalog_save_title = gr.Textbox( label="Catalog title", placeholder="My LoRA name", scale=2, ) catalog_save_weight = gr.Slider( label="Default weight", minimum=0.0, maximum=2.0, step=0.05, value=1.0, scale=1, ) catalog_save_prompt = gr.Textbox( label="Default prompt (optional)", lines=2, placeholder="Safe ready-to-go prompt auto-appended when selected", ) catalog_save_triggers = gr.Textbox( label="Known triggers (optional)", lines=3, placeholder=( "One per line or comma-separated. Docs only — not auto-appended.\n" "e.g. small penis, large penis, flaccid penis, erect penis" ), info="Can include mutually exclusive keywords; pick what you need in the prompt.", ) catalog_save_notes = gr.Textbox( label="Notes (optional)", lines=2, placeholder="Usage notes, caveats, pairing tips…", ) with gr.Row(): catalog_save_btn = gr.Button( "💾 Save to catalog", variant="secondary", ) catalog_save_status = gr.Textbox(label="Catalog status", interactive=False) gr.Markdown("#### 🗑️ Remove LoRA") gr.Markdown( "Remove a **session** custom LoRA, or a catalog entry that is **not** " "`admin_approved`. Admin-approved entries can only be archived via JSON " "(`active: false`)." ) with gr.Row(): catalog_remove_select = gr.Dropdown( label="LoRA to remove", choices=removable_catalog_titles(), value=None, interactive=True, scale=2, ) catalog_remove_btn = gr.Button("🗑️ Remove", variant="stop", scale=1) catalog_remove_status = gr.Textbox(label="Remove status", interactive=False) with gr.Accordion("📝 Custom Prompts", open=False): gr.Markdown("Save reusable prompt snippets for this session.") custom_prompt_selector = gr.CheckboxGroup( choices=[], value=[], label="Saved prompts — tick to append to generation", ) with gr.Row(): prompt_name_input = gr.Textbox(label="Name", placeholder="e.g. Skin detail enhancer", scale=1) with gr.Row(): prompt_text_input = gr.Textbox(label="Prompt text", lines=4, placeholder="Enter the prompt snippet you want to save…") with gr.Row(): add_prompt_btn = gr.Button("💾 Save Prompt", variant="primary") prompt_status = gr.Textbox(label="Status", interactive=False, scale=2) with gr.Row(): delete_prompt_name = gr.Dropdown(label="Delete a saved prompt", choices=[], value=None, interactive=True, scale=2) delete_prompt_btn = gr.Button("🗑️ Delete", variant="secondary", scale=1) # Full-width advanced block at the bottom of Generate tab with gr.Accordion("⚙️ Advanced Settings", open=False): with gr.Row(): with gr.Column(scale=1): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) guidance_scale = gr.Slider( label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0, visible=MODEL_VARIANT != "9B-KV", ) steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1) upscale_factor = gr.Dropdown( label="Upscale model", choices=list(UPSCALE_MODELS.keys()), value="None", ) with gr.Column(scale=1): gr.Markdown("#### 🖼️ Output canvas size") canvas_mode = gr.Radio( choices=["Auto (from base image)", "Custom"], value="Auto (from base image)", label="Canvas mode", info=("Auto matches the base image's aspect ratio (longest side 1024). " "Use Custom when base and references have very different proportions."), ) custom_width = gr.Slider( label="Width", minimum=512, maximum=2048, step=16, value=1024, visible=False, ) custom_height = gr.Slider( label="Height", minimum=512, maximum=2048, step=16, value=1024, visible=False, ) canvas_fit_mode = gr.Radio( choices=["Stretch", "Pad (color)", "Pad (blur)", "Crop (cover)"], value="Stretch", label="Canvas fit mode", info=("How input images are placed into the canvas. " "Stretch = current default (can squish). " "Pad keeps aspect; Crop fills by trimming edges."), ) pad_color = gr.ColorPicker( label="Pad colour", value="#000000", visible=False, ) gr.Markdown("#### 🔁 Batch") batch_count = gr.Slider( label="Number of runs", minimum=1, maximum=12, step=1, value=1, ) batch_vary = gr.Radio( choices=["Random seed each run", "Sequential seed (+1 each)", "Sweep first LoRA weight"], value="Random seed each run", label="Variation strategy", info=("Sweep linearly varies the weight of whichever LoRA is in " "slot 1 (first ticked) across the runs."), ) with gr.Row(): sweep_min = gr.Slider( label="Sweep min weight", minimum=0.0, maximum=2.0, step=0.05, value=0.4, visible=False, ) sweep_max = gr.Slider( label="Sweep max weight", minimum=0.0, maximum=2.0, step=0.05, value=1.4, visible=False, ) with gr.Accordion("📋 Selected LoRA details (repo / triggers / notes)", open=False): selected_lora_details = gr.Markdown( value="*Tick one or more LoRAs above to see full repo paths, " "known triggers, and notes.*", elem_id="selected_lora_details", ) # ── Crop / Fix Image tab ───────────────────────────────────────── with gr.Tab("✂️ Crop / Fix Image", id="tab_editor"): gr.Markdown( "Upload an image to crop / paint on it, then send the result to the Base " "Image or add it as a Reference. EXIF orientation is corrected on export." ) editor = gr.ImageEditor( label="Editor", type="pil", transforms=("crop",), brush=gr.Brush(default_size=12, colors=["#FF4500", "#FFFFFF", "#000000", "#FF0000", "#00FF00", "#0000FF"], color_mode="fixed"), eraser=gr.Eraser(default_size=20), layers=False, sources=["upload", "clipboard"], height=420, ) with gr.Row(): heic_uploader = gr.File( label="📸 Load HEIC / HEIF (iPhone photos)", file_types=[".heic", ".heif", ".HEIC", ".HEIF"], file_count="single", type="filepath", ) with gr.Row(): send_to_base_btn = gr.Button("→ Send to Base Image", variant="primary") send_to_ref_btn = gr.Button("→ Add to Reference Images") # ── Extend canvas section ──────────────────────────────────────────────── # Uses the editor's current composite as the source so cropping + painting # happen first, then we grow the canvas around the result. Output is # loaded back into the same editor — Send → Base / Reference from there. with gr.Accordion("📐 Extend canvas (add padding around image)", open=True): gr.Markdown( "Grow the editor image's canvas by a percentage in any combination " "of directions. Percentages are relative to the *current* image " "size — `Down = 100` doubles the height with the image on top. " "The result replaces the editor contents so you can crop again or " "send it to Base / Reference with the buttons above." ) with gr.Row(): ext_up = gr.Number(label="Up %", value=0, minimum=0, precision=2) ext_down = gr.Number(label="Down %", value=0, minimum=0, precision=2) ext_left = gr.Number(label="Left %", value=0, minimum=0, precision=2) ext_right = gr.Number(label="Right %", value=0, minimum=0, precision=2) with gr.Row(): ext_fill = gr.ColorPicker(label="Fill colour", value="#000000") extend_btn = gr.Button("📐 Extend canvas", variant="primary") with gr.Row(): ext_schematic = gr.Image( label="Layout preview (red outline = current image)", type="pil", interactive=False, height=220, ) ext_info = gr.Markdown("*Upload something into the editor first.*") # ── Bulk processing tab ────────────────────────────────────────── with gr.Tab("📦 Bulk Process", id="tab_bulk"): gr.Markdown( "Upload many images and process them with the **same settings as the " "Generate tab** (prompt, LoRAs, weights, canvas, upscaler, etc.). " "Outputs stream in one-by-one — each image is its own GPU call, so a " "ZeroGPU quota wall mid-run only loses the in-progress item. " "Earlier outputs stay in the gallery and on disk under `/tmp/bulk_/`." ) bulk_files = gr.File( label="Input images", file_count="multiple", type="filepath", file_types=["image", ".heic", ".heif"], ) with gr.Row(): bulk_run_btn = gr.Button("▶ Start bulk run", variant="primary") bulk_stop_btn = gr.Button("⏹ Stop", variant="stop") bulk_status = gr.Markdown("*Ready.*") bulk_gallery = gr.Gallery( label="Bulk outputs", type="filepath", columns=4, rows=2, height=480, allow_preview=True, object_fit="contain", ) bulk_zip = gr.File(label="📥 Download all (zip + manifest.csv)", interactive=False) # ── Depth / Pose tab ───────────────────────────────────────────── with gr.Tab("🦴 Depth / Pose", id="tab_control"): gr.Markdown( "Generate ControlNet-style **depthmaps** and editable **OpenPose** " "skeletons. The result feeds well into the **RefControl – Depth** / " "**RefControl – Pose** LoRAs on the Generate tab when sent as a " "Reference image." ) pose_source_state = gr.State(None) pose_keypoints_state = gr.State([]) with gr.Row(): with gr.Column(scale=1): ctrl_source = gr.Image( label="Source image", type="pil", sources=["upload", "clipboard"], height=320, ) with gr.Row(): detect_depth_btn = gr.Button("🌐 Generate depthmap", variant="primary") detect_pose_btn = gr.Button("🦴 Detect pose", variant="primary") insert_blank_btn = gr.Button("➕ Insert blank skeleton template") with gr.Column(scale=1): depth_output = gr.Image(label="Depthmap", type="pil", interactive=False, height=320, format="png") with gr.Row(): send_depth_ref_btn = gr.Button("→ Send depth to Reference", variant="primary") send_depth_base_btn = gr.Button("→ Send depth to Base") gr.Markdown("### ✏️ Pose editor") gr.Markdown( "Pick a person and a joint, then **click anywhere on the editor preview** " "to move that joint. Hidden joints can be re-added the same way — select " "them and click. Use the buttons below for delete / clear / re-detect." ) with gr.Row(): with gr.Column(scale=1): pose_overlay = gr.Image( label="Editor — click to place active joint", type="pil", interactive=False, height=420, format="png", ) with gr.Column(scale=1): pose_clean = gr.Image( label="Skeleton (sent to Reference / Base)", type="pil", interactive=False, height=420, format="png", ) with gr.Row(): active_person_dd = gr.Dropdown( label="Active person", choices=[], value=None, interactive=True, ) active_joint_dd = gr.Dropdown( label="Active joint", choices=list(OPENPOSE_KEYPOINT_NAMES), value=None, interactive=True, ) with gr.Row(): delete_joint_btn = gr.Button("🗑️ Hide active joint") reset_pose_btn = gr.Button("🔄 Re-detect from source") clear_pose_btn = gr.Button("🧹 Clear all joints") with gr.Row(): send_pose_ref_btn = gr.Button("→ Send pose to Reference", variant="primary") send_pose_base_btn = gr.Button("→ Send pose to Base") # ── Event wiring ───────────────────────────────────────────────────────── # Lightweight UI handlers: hide Gradio progress. On ZeroGPU/Spaces, the # default spinner often sticks on pure gr.update visibility changes # (LoRA weight sliders, canvas size text) until another event flushes UI. _ui = dict(show_progress="hidden") base_image.upload(fn=reencode_upload, inputs=[base_image], outputs=[base_image], **_ui) base_image.change(fn=on_base_image_change, inputs=[base_image], outputs=[size_info], **_ui) make_solid_base_btn.click( fn=make_solid_base_image, inputs=[base_solid_color, canvas_mode, custom_width, custom_height], outputs=[base_image], show_progress="hidden", ).then( fn=on_base_image_change, inputs=[base_image], outputs=[size_info], **_ui, ) # Progressive ref slots: pack non-empty images upward on any change so # deleting Ref1 shifts Ref2/3 up instead of wiping them. for _ref in (ref1, ref2, ref3): _ref.upload(fn=reencode_upload, inputs=[_ref], outputs=[_ref], **_ui) _ref_compact_outputs = [ ref1, ref2, ref3, reference_info, ref2_up_btn, ref2_down_btn, ref3_up_btn, ] for _ref in (ref1, ref2, ref3): _ref.change( fn=compact_reference_slots, inputs=[ref1, ref2, ref3], outputs=_ref_compact_outputs, **_ui, ) _slot_inputs = [base_image, ref1, ref2, ref3] _move_outputs = [ base_image, ref1, ref2, ref3, reference_info, ref2_up_btn, ref2_down_btn, ref3_up_btn, ] base_down_btn.click(fn=move_base_down, inputs=_slot_inputs, outputs=_move_outputs, **_ui) ref1_up_btn.click(fn=move_ref1_up, inputs=_slot_inputs, outputs=_move_outputs, **_ui) ref1_down_btn.click(fn=move_ref1_down, inputs=_slot_inputs, outputs=_move_outputs, **_ui) ref2_up_btn.click(fn=move_ref2_up, inputs=_slot_inputs, outputs=_move_outputs, **_ui) ref2_down_btn.click(fn=move_ref2_down, inputs=_slot_inputs, outputs=_move_outputs, **_ui) ref3_up_btn.click(fn=move_ref3_up, inputs=_slot_inputs, outputs=_move_outputs, **_ui) # Every browser open re-reads the bucket JSON and refreshes selector choices. # Without this, choices stay frozen at process start and newly saved LoRAs # (e.g. thickcum) are "already in catalog" but invisible in new sessions. def _on_page_load(dynamic_loras, selected): sel_upd, save_upd, rem_upd, dyn = refresh_catalog_ui(dynamic_loras, selected) # Count from live catalog after reload (gr.update is not always a plain dict). n = len(get_selectable_styles(dyn)) status = f"*Catalog loaded — **{n}** active LoRA(s).*" return sel_upd, save_upd, rem_upd, dyn, status demo.load( fn=_on_page_load, inputs=[dynamic_loras_state, lora_selector], outputs=[lora_selector, catalog_save_select, catalog_remove_select, dynamic_loras_state, catalog_load_status], show_progress="hidden", ) # Must be registered inside the Blocks context (Gradio rejects load outside). demo.load(fn=None, js=_TEXTBOX_NEWLINE_JS) # Capture clicks on Gradio Gallery's built-in download (↓) control. demo.load( fn=None, js=""" () => { if (window.__fluxGalleryDlHook) return; window.__fluxGalleryDlHook = true; const setSignal = (val) => { const root = document.getElementById('gallery_dl_signal'); if (!root) return; const ta = root.querySelector('textarea, input'); if (!ta) return; ta.value = val || ''; ta.dispatchEvent(new Event('input', { bubbles: true })); }; document.addEventListener('click', (e) => { const t = e.target; if (!t || !t.closest) return; const gal = t.closest('#output_gallery'); if (!gal) return; // Gradio download control: anchor with download attr, or button near download icon. const a = t.closest('a[download], a.download-link, a[href*="file="]'); const btn = t.closest('button'); let href = ''; if (a && a.href) { href = a.getAttribute('download') || a.href; } else if (btn) { const label = (btn.getAttribute('aria-label') || btn.title || btn.textContent || '').toLowerCase(); if (!(label.includes('download') || label.includes('save') || btn.innerHTML.includes('download'))) { // still allow if nested svg title looks like download const svgTitle = (btn.querySelector('title')?.textContent || '').toLowerCase(); if (!svgTitle.includes('download') && !btn.querySelector('[data-testid*="download"]')) { return; } } const nearA = btn.closest('a') || btn.querySelector('a') || gal.querySelector('a[download]'); href = (nearA && (nearA.getAttribute('download') || nearA.href)) || ''; if (!href) { // fallback: selected/preview image src basename const img = gal.querySelector('.preview img, .thumbnail-lg img, img'); href = (img && (img.currentSrc || img.src)) || 'gallery-download'; } } else { return; } if (!href) return; try { const u = href.startsWith('http') || href.startsWith('blob:') || href.startsWith('/') ? href : href; const base = (u.split('/').pop() || u).split('?')[0]; setSignal(base + '|' + Date.now()); } catch (_) { setSignal(String(href) + '|' + Date.now()); } }, true); } """, ) reload_catalog_btn.click( fn=_on_page_load, inputs=[dynamic_loras_state, lora_selector], outputs=[lora_selector, catalog_save_select, catalog_remove_select, dynamic_loras_state, catalog_load_status], show_progress="minimal", ) # update_weight_sliders is the one imported from lora_registry now. # Only wire .change — also binding .input/.select raced and could leave the # LoRA prompt box hidden/empty while still applying defaults at generate time. # Pass live slider/prompt values + memory so existing settings survive add/remove. _lora_slider_inputs = [ lora_selector, dynamic_loras_state, lora_weight_memory_state, lora_prev_selected_state, lora_prompt_memory_state, lora_prompt_display, ] + weight_sliders _lora_slider_outputs = ( weight_sliders + [lora_prompt_display, selected_lora_details, lora_weight_memory_state, lora_prev_selected_state, lora_prompt_memory_state] ) lora_selector.change( fn=update_weight_sliders, inputs=_lora_slider_inputs, outputs=_lora_slider_outputs, show_progress="hidden", trigger_mode="once", ) # add_custom_lora is also imported from lora_registry. # Always session-only; optional persist / remove are separate explicit actions. add_lora_btn.click( fn=add_custom_lora, inputs=[lora_repo_id, lora_weight_name, lora_adapter_name, dynamic_loras_state], outputs=[lora_status, lora_selector, dynamic_loras_state, catalog_save_select, catalog_remove_select], show_progress="minimal", ) catalog_save_select.change( fn=fill_catalog_save_form, inputs=[catalog_save_select, dynamic_loras_state], outputs=[catalog_save_title, catalog_save_weight, catalog_save_prompt, catalog_save_triggers, catalog_save_notes], show_progress="hidden", ) catalog_save_btn.click( fn=save_session_lora_to_catalog, inputs=[catalog_save_select, catalog_save_title, catalog_save_weight, catalog_save_prompt, dynamic_loras_state, catalog_save_triggers, catalog_save_notes, lora_selector], outputs=[catalog_save_status, lora_selector, catalog_save_select, catalog_remove_select, dynamic_loras_state], show_progress="minimal", # After save, force weight/prompt UI to follow the remapped selection # (Custom: x → catalog title) so nothing stays bound to a removed title. ).then( fn=update_weight_sliders, inputs=_lora_slider_inputs, outputs=_lora_slider_outputs, show_progress="hidden", ) catalog_remove_btn.click( fn=remove_lora, inputs=[catalog_remove_select, dynamic_loras_state, lora_selector], outputs=[catalog_remove_status, lora_selector, catalog_save_select, catalog_remove_select, dynamic_loras_state], show_progress="minimal", ) add_prompt_btn.click( fn=add_custom_prompt, inputs=[prompt_name_input, prompt_text_input, custom_prompts_state, custom_prompt_counter_state], outputs=[prompt_status, custom_prompts_state, custom_prompt_counter_state, custom_prompt_selector, delete_prompt_name, prompt_name_input], show_progress="hidden", ) delete_prompt_btn.click( fn=delete_custom_prompt, inputs=[delete_prompt_name, custom_prompt_selector, custom_prompts_state], outputs=[prompt_status, custom_prompts_state, custom_prompt_selector, delete_prompt_name], show_progress="hidden", ) custom_prompt_selector.change( fn=update_custom_prompt_display, inputs=[custom_prompt_selector, custom_prompts_state], outputs=[custom_prompt_display], show_progress="hidden", trigger_mode="always_last", ) canvas_mode.change(fn=on_canvas_mode_change, inputs=[canvas_mode], outputs=[custom_width, custom_height], **_ui) canvas_fit_mode.change(fn=on_fit_mode_change, inputs=[canvas_fit_mode], outputs=[pad_color], **_ui) batch_vary.change(fn=on_batch_vary_change, inputs=[batch_vary], outputs=[sweep_min, sweep_max], **_ui) output_gallery.select(fn=on_gallery_select, inputs=[output_gallery], outputs=[selected_output_state], show_progress="hidden") # Keep download buttons pointed at the selected gallery item (or latest) # plus the last-run prompt text. selected_output_state.change( fn=resolve_download_bundle, inputs=[selected_output_state, output_gallery, last_prompt_state, project_name], outputs=[download_png_btn, download_webp_btn, download_prompt_btn], show_progress="hidden", ) output_gallery.change( fn=resolve_download_bundle, inputs=[selected_output_state, output_gallery, last_prompt_state, project_name], outputs=[download_png_btn, download_webp_btn, download_prompt_btn], show_progress="hidden", ) # Track which gallery images still need downloading. output_gallery.change( fn=sync_download_tracking, inputs=[output_gallery, pending_download_state, downloaded_images_state], outputs=[pending_download_state, downloaded_images_state, download_status], show_progress="hidden", ) project_name.change( fn=resolve_download_bundle, inputs=[selected_output_state, output_gallery, last_prompt_state, project_name], outputs=[download_png_btn, download_webp_btn, download_prompt_btn], show_progress="hidden", ) # Open PNG/WebP in a new tab (in addition to the browser download) and mark # the current gallery image as downloaded for the pending-status tracker. _OPEN_DL_TAB_JS = """ (btnId) => { const openHref = (href) => { if (!href || href === '#' || href.endsWith('/')) return false; window.open(href, '_blank', 'noopener,noreferrer'); return true; }; const tryOpen = () => { const root = document.getElementById(btnId); if (!root) return false; const anchors = root.querySelectorAll('a.download-link, a[href], a[download]'); for (const a of anchors) { const href = a.href || a.getAttribute('href') || ''; if (openHref(href)) return true; } // Gradio sometimes nests the file link one tick later after value bind. return false; }; if (tryOpen()) return; // Retry briefly — DownloadButton href can lag the click on first bind. let n = 0; const t = setInterval(() => { n += 1; if (tryOpen() || n >= 8) clearInterval(t); }, 50); } """ download_png_btn.click( fn=None, # Concatenate (not f-string) so braces inside _OPEN_DL_TAB_JS stay literal JS. js="() => { (" + _OPEN_DL_TAB_JS + ")('download_png_btn'); }", ).then( fn=mark_current_output_downloaded, inputs=[selected_output_state, output_gallery, pending_download_state, downloaded_images_state], outputs=[pending_download_state, downloaded_images_state, download_status], show_progress="hidden", ) download_webp_btn.click( fn=None, js="() => { (" + _OPEN_DL_TAB_JS + ")('download_webp_btn'); }", ).then( fn=mark_current_output_downloaded, inputs=[selected_output_state, output_gallery, pending_download_state, downloaded_images_state], outputs=[pending_download_state, downloaded_images_state, download_status], show_progress="hidden", ) # Gallery built-in ↓ icon → JS signal → mark downloaded. gallery_dl_signal.change( fn=mark_from_gallery_signal, inputs=[gallery_dl_signal, output_gallery, pending_download_state, downloaded_images_state], outputs=[pending_download_state, downloaded_images_state, download_status, gallery_dl_signal], show_progress="hidden", ) # Reset any stale gallery selection before a new run starts, so Send→Base # / Send→Ref after this run can't accidentally reuse a path from the # previous run's gallery contents. Also warn if undownloaded images exist. _infer_inputs = [ base_image, ref1, ref2, ref3, prompt, lora_prompt_display, custom_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, batch_count, batch_vary, sweep_min, sweep_max, project_name, dynamic_loras_state, ] + weight_sliders _infer_outputs = [ output_gallery, used_seed, download_png_btn, download_webp_btn, download_prompt_btn, last_prompt_state, ] for _run_btn in (run_button, run_button_top): _run_btn.click( fn=warn_undownloaded_before_generate, inputs=[pending_download_state, output_gallery], outputs=[selected_output_state], show_progress="hidden", ) run_event = run_button.click( fn=infer, inputs=_infer_inputs, # Download bundle on every generate yield — required because # gallery.change alone misses the first virgin-session result. outputs=_infer_outputs, ) run_event_top = run_button_top.click( fn=infer, inputs=_infer_inputs, outputs=_infer_outputs, ) # ── Editor tab wiring ──────────────────────────────────────────────────── heic_uploader.upload(fn=load_heic_to_editor, inputs=[heic_uploader], outputs=[editor]) send_to_base_btn.click(fn=send_editor_to_base, inputs=[editor], outputs=[base_image]) \ .then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info]) \ .then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs]) _send_ref_outputs = [ ref1, ref2, ref3, reference_info, ref2_up_btn, ref2_down_btn, ref3_up_btn, ] send_to_ref_btn.click( fn=send_editor_to_reference, inputs=[editor, ref1, ref2, ref3], outputs=_send_ref_outputs, ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs]) send_out_to_base_btn.click( fn=send_output_to_base, inputs=[selected_output_state, output_gallery], outputs=[base_image], ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info]) send_out_to_ref_btn.click( fn=send_output_to_reference, inputs=[selected_output_state, output_gallery, ref1, ref2, ref3], outputs=_send_ref_outputs, ) # Extend-canvas wiring # The schematic previews the *current editor composite* so users see live # feedback as they nudge the percentages / fill colour. def _editor_source_for_preview(editor_value): if not editor_value or editor_value.get("composite") is None: return None comp = editor_value["composite"] if isinstance(comp, np.ndarray): from PIL import Image as _Image comp = _Image.fromarray(comp) return comp def _update_extend_preview(editor_value, up, down, left, right, fill): return render_extend_schematic( _editor_source_for_preview(editor_value), up, down, left, right, fill, ) _extend_preview_inputs = [editor, ext_up, ext_down, ext_left, ext_right, ext_fill] _extend_preview_outputs = [ext_schematic, ext_info] for _c in (ext_up, ext_down, ext_left, ext_right, ext_fill): _c.change(fn=_update_extend_preview, inputs=_extend_preview_inputs, outputs=_extend_preview_outputs) # Refresh the schematic when a NEW image lands in the editor — not on every # `change` event. `editor.change` fires very frequently on iOS Safari/Chrome # (once per stroke/layer/crop-preview) and the round-trips OOM'd the tab # even for small uploads. `.upload` fires only when a new image comes in # via the upload/clipboard sources, which is the case the preview actually # cares about (image dimensions changed → schematic scale needs redrawing). editor.upload(fn=_update_extend_preview, inputs=_extend_preview_inputs, outputs=_extend_preview_outputs) # HEIC uploads bypass the editor's own upload event because they come from # the separate File component, so wire that path in explicitly too. heic_uploader.upload(fn=_update_extend_preview, inputs=_extend_preview_inputs, outputs=_extend_preview_outputs) # Run: extend, then hand the new PIL back to the editor. `render_extend_ # schematic` re-runs via editor.change once the new image lands, so no # extra .then() is needed for the preview. extend_btn.click( fn=extend_editor_canvas, inputs=[editor, ext_up, ext_down, ext_left, ext_right, ext_fill], outputs=[editor], ) # ── Bulk tab wiring ────────────────────────────────────────────────────── bulk_event = bulk_run_btn.click( fn=bulk_infer, inputs=[bulk_files, prompt, lora_prompt_display, custom_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor, canvas_mode, custom_width, custom_height, canvas_fit_mode, pad_color, dynamic_loras_state] + weight_sliders, outputs=[bulk_gallery, bulk_status, bulk_zip], ) bulk_stop_btn.click(fn=lambda: gr.Info("Stop requested — finishing current image."), cancels=[bulk_event, run_event, run_event_top]) # ── Depth / Pose tab wiring ────────────────────────────────────────────── ctrl_source.change( fn=lambda img: img, inputs=[ctrl_source], outputs=[pose_source_state], ) detect_depth_btn.click( fn=generate_depthmap, inputs=[ctrl_source], outputs=[depth_output], ) def _on_detect_pose(source): if source is None: raise gr.Error("Upload a source image first.") poses, w, h = detect_pose(source) if not poses: gr.Warning("No people detected — try 'Insert blank skeleton template' " "or a different image.") return ([], gr.update(choices=[], value=None), gr.update(value=None), None, None) return ( poses, gr.update(choices=person_choices(poses), value="Person 1"), gr.update(value=OPENPOSE_KEYPOINT_NAMES[0]), render_pose_overlay(source, poses, 0, 0), render_pose_skeleton(poses, w, h), ) detect_pose_btn.click( fn=_on_detect_pose, inputs=[ctrl_source], outputs=[pose_keypoints_state, active_person_dd, active_joint_dd, pose_overlay, pose_clean], ) reset_pose_btn.click( fn=_on_detect_pose, inputs=[ctrl_source], outputs=[pose_keypoints_state, active_person_dd, active_joint_dd, pose_overlay, pose_clean], ) def _on_insert_blank(source): if source is None: raise gr.Error("Upload a source image first.") w, h = source.size poses = [default_pose_template(w, h)] return ( poses, gr.update(choices=["Person 1"], value="Person 1"), gr.update(value=OPENPOSE_KEYPOINT_NAMES[0]), render_pose_overlay(source, poses, 0, 0), render_pose_skeleton(poses, w, h), ) insert_blank_btn.click( fn=_on_insert_blank, inputs=[ctrl_source], outputs=[pose_keypoints_state, active_person_dd, active_joint_dd, pose_overlay, pose_clean], ) def _on_overlay_click(evt: gr.SelectData, poses, source, person_label, joint_name): if not poses or source is None or evt is None or evt.index is None: return gr.update(), gr.update(), gr.update() person_idx = parse_person_idx(person_label) joint_idx = joint_name_to_index(joint_name) if person_idx is None or joint_idx < 0: return gr.update(), gr.update(), gr.update() x, y = evt.index w, h = source.size new_poses = move_joint(poses, person_idx, joint_idx, x, y, w, h) return ( new_poses, render_pose_overlay(source, new_poses, person_idx, joint_idx), render_pose_skeleton(new_poses, w, h), ) pose_overlay.select( fn=_on_overlay_click, inputs=[pose_keypoints_state, pose_source_state, active_person_dd, active_joint_dd], outputs=[pose_keypoints_state, pose_overlay, pose_clean], ) def _on_active_change(poses, source, person_label, joint_name): if not poses or source is None: return gr.update() person_idx = parse_person_idx(person_label) or 0 joint_idx = max(joint_name_to_index(joint_name), 0) return render_pose_overlay(source, poses, person_idx, joint_idx) active_person_dd.change( fn=_on_active_change, inputs=[pose_keypoints_state, pose_source_state, active_person_dd, active_joint_dd], outputs=[pose_overlay], ) active_joint_dd.change( fn=_on_active_change, inputs=[pose_keypoints_state, pose_source_state, active_person_dd, active_joint_dd], outputs=[pose_overlay], ) def _on_hide_active(poses, source, person_label, joint_name): person_idx = parse_person_idx(person_label) joint_idx = joint_name_to_index(joint_name) new_poses = hide_joint(poses, person_idx, joint_idx) if source is None: return new_poses, gr.update(), gr.update() w, h = source.size return (new_poses, render_pose_overlay(source, new_poses, person_idx, joint_idx), render_pose_skeleton(new_poses, w, h)) delete_joint_btn.click( fn=_on_hide_active, inputs=[pose_keypoints_state, pose_source_state, active_person_dd, active_joint_dd], outputs=[pose_keypoints_state, pose_overlay, pose_clean], ) def _on_clear_all(poses, source): new_poses = clear_all_joints(poses) if source is None: return new_poses, gr.update(), gr.update() w, h = source.size return (new_poses, render_pose_overlay(source, new_poses, None, None), render_pose_skeleton(new_poses, w, h)) clear_pose_btn.click( fn=_on_clear_all, inputs=[pose_keypoints_state, pose_source_state], outputs=[pose_keypoints_state, pose_overlay, pose_clean], ) send_depth_ref_btn.click( fn=push_pil_to_reference, inputs=[depth_output, ref1, ref2, ref3], outputs=_send_ref_outputs, ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs]) send_depth_base_btn.click( fn=push_pil_to_base, inputs=[depth_output], outputs=[base_image], ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info] ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs]) send_pose_ref_btn.click( fn=push_pil_to_reference, inputs=[pose_clean, ref1, ref2, ref3], outputs=_send_ref_outputs, ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs]) send_pose_base_btn.click( fn=push_pil_to_base, inputs=[pose_clean], outputs=[base_image], ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info] ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs]) if __name__ == "__main__": # Gradio 6.0: theme and css go on launch(), not Blocks() demo.queue().launch(css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)