import os import json import copy import random import math import gradio as gr from PIL import Image, ImageDraw, ImageFont from modules import sd_samplers, errors, scripts, images, sd_models from modules.paths_internal import roboto_ttf_file from modules.processing import Processed, process_images from modules.shared import state, cmd_opts, opts from pathlib import Path lora_dir = Path(cmd_opts.lora_dir).resolve() def allowed_path(path): return Path(path).resolve().is_relative_to(lora_dir) def get_base_path(is_use_custom_path, custom_path): return lora_dir.joinpath(custom_path) if is_use_custom_path else lora_dir def is_directory_contain_lora(path): try: if allowed_path(path): safetensor_files = [f for f in os.listdir( path) if f.endswith('.safetensors')] return len(safetensor_files) > 0 except FileNotFoundError: pass except Exception as e: print(e) return False def get_directories(base_path, include_root=True): directories = ["/"] if include_root else [] try: if allowed_path(base_path): for entry in os.listdir(base_path): full_path = os.path.join(base_path, entry) if os.path.isdir(full_path): if is_directory_contain_lora(full_path): directories.append(entry) nested_directories = get_directories( full_path, include_root=False) directories.extend([os.path.join(entry, d) for d in nested_directories]) except FileNotFoundError: pass except Exception as e: print(e) return directories def read_json_file(file_path): with open(file_path, 'r') as file: return json.load(file) def get_lora_name(lora_path): if opts.lora_preferred_name == "Filename": lora_name = lora_path.stem else: metadata = sd_models.read_metadata_from_safetensors(lora_path) lora_name = metadata.get('ss_output_name', lora_path.stem) return lora_name def get_lora_prompt(lora_path, json_path): with open(json_path, 'r', encoding='utf-8') as file: data = json.load(file) preferred_weight = data.get("preferred weight", 1) activation_text = data.get("activation text", "") try: if float(preferred_weight) == 0: preferred_weight = 1 except: preferred_weight = 1 lora_name = get_lora_name(lora_path) return f", {activation_text}," def image_grid_with_text(imgs, texts, rows=None, cols=None, font_path=None, font_size=20, text_color="#FFFFFF", stroke_color="#000000", stroke_width=2, add_text=True): if rows is None: rows = round(math.sqrt(len(imgs))) cols = math.ceil(len(imgs) / rows) if cols is None else cols w, h = imgs[0].size grid = Image.new('RGB', (cols * w, rows * h), 'black') for i, img in enumerate(imgs): grid.paste(img, (i % cols * w, i // cols * h)) if add_text: draw = ImageDraw.Draw(grid) try: font = ImageFont.truetype(font_path, font_size) if font_path and os.path.exists( font_path) else ImageFont.truetype(roboto_ttf_file, font_size) except: font = ImageFont.truetype(roboto_ttf_file, font_size) for i, text in enumerate(texts): x = (i % cols) * w y = (i // cols) * h for dx, dy in [(j, k) for j in range(-stroke_width, stroke_width+1) for k in range(-stroke_width, stroke_width+1)]: draw.text((x+5+dx, y+5+dy), text, font=font, fill=stroke_color) draw.text((x+5, y+5), text, font=font, fill=text_color) return grid class Script(scripts.Script): def title(self): return "Apply on every Lora" def ui(self, is_img2img): def build_lora_tree(base_path): tree = {"__root__": {"name": base_path.name, "children": {}}} for root, dirs, files in os.walk(base_path): rel_path = os.path.relpath(root, base_path) current_node = tree["__root__"] if rel_path != ".": for part in rel_path.split(os.sep): current_node = current_node["children"].setdefault( part, {"name": part, "children": {}, "loras": []}) loras = [f[:-12] for f in files if f.endswith(".safetensors")] current_node["loras"] = loras return tree["__root__"] def update_tree(is_use_custom, custom_path): base_path = get_base_path(is_use_custom, custom_path) return gr.Tree.update(value=build_lora_tree(base_path)) with gr.Column(): base_dir_checkbox = gr.Checkbox( label="Use Custom Lora path", value=False) base_dir_textbox = gr.Textbox( label="Lora directory", visible=False) with gr.Row(): lora_dir_dropdown = gr.Dropdown( label="LORA Directory", choices=["/"] + get_directories(lora_dir), value="/", interactive=True ) refresh_btn = gr.Button("🔄", variant="tool") lora_checkboxes = gr.CheckboxGroup( label="Select LoRAs", interactive=True ) def update_directory(current_dir): base_path = lora_dir.joinpath(current_dir.lstrip('/')) loras = [] if allowed_path(base_path): for root, _, files in os.walk(base_path): for file in files: if file.endswith(('.safetensors', '.pt')): rel_path = os.path.relpath(root, lora_dir) loras.append( f"{rel_path}/{file}" if rel_path != '.' else file) return gr.CheckboxGroup.update(choices=loras) def scan_loras(current_dir): return update_directory(current_dir) lora_dir_dropdown.change( fn=scan_loras, inputs=[lora_dir_dropdown], outputs=lora_checkboxes ) refresh_btn.click( fn=lambda: scan_loras(lora_dir_dropdown.value), outputs=lora_checkboxes ) prompt_lines = gr.Textbox(label="Prompts (one per line)", lines=5) lora_tags_position_radio = gr.Radio( ["Prepend", "Append"], value="Prepend", label="LoRA Tags Position") checkbox_save_grid = gr.Checkbox( label="Save grid image", value=True) font_path = gr.Textbox(label="Custom Font Path") with gr.Row(): use_random_seed = gr.Checkbox( label="Random seed", value=True) use_fixed_seed = gr.Checkbox(label="Fixed seed", value=False) file_upload = gr.File( label="Load prompts from file", file_types=[".txt"], type='binary') def load_prompt_file(file, current_prompts): if file is None: return None, current_prompts, gr.update() lines = [x.strip() for x in file.decode( 'utf8', errors='ignore').split("\n")] return None, "\n".join(lines), gr.update(lines=max(7, len(lines))) file_upload.change( fn=load_prompt_file, inputs=[file_upload, prompt_lines], outputs=[file_upload, prompt_lines, prompt_lines], show_progress=False ) base_dir_checkbox.change( fn=lambda is_use, path: get_base_path(is_use, path), inputs=[base_dir_checkbox, base_dir_textbox], outputs=lora_dir_dropdown ) return [base_dir_checkbox, base_dir_textbox, lora_checkboxes, prompt_lines, lora_tags_position_radio, checkbox_save_grid, font_path] def run(self, p, is_use_custom_path, custom_path, lora_checkboxes, prompt_lines, lora_tags_position, is_save_grid, font_path): selected_loras = [ str(lora_dir.joinpath(lora)) for lora in lora_checkboxes if lora.endswith(('.safetensors', '.pt')) ] if not selected_loras or not prompt_lines: return Processed(p, [], p.seed, "No LoRAs or prompts selected") prompts = [line.strip() for line in prompt_lines.splitlines() if line.strip()] combinations = [(lora, prompt) for lora in selected_loras for prompt in prompts] state.job_count = len(combinations) result_images = [] all_prompts = [] infotexts = [] grid_texts = [] for lora_path, prompt in combinations: if state.interrupted: break current_p = copy.copy(p) lora_file = Path(lora_path) json_file = lora_file.with_suffix('.json') try: lora_tags = get_lora_prompt( lora_file, json_file) if json_file.exists() else f"," except Exception as e: print(f"Error loading Lora {lora_file}: {str(e)}") continue final_prompt = f"{lora_tags} {prompt}" if lora_tags_position == "Prepend" else f"{prompt} {lora_tags}" current_p.prompt = final_prompt proc = process_images(current_p) result_images.extend(proc.images) all_prompts.extend(proc.all_prompts) infotexts.extend(proc.infotexts) grid_texts.extend( [f"{lora_file.stem}\n{prompt}"] * len(proc.images)) if is_save_grid and len(result_images) > 1: rows = round(math.sqrt(len(result_images))) grid_image = image_grid_with_text( result_images, grid_texts, rows=rows, font_path=font_path, text_color="#FFFFFF", stroke_color="#000000", stroke_width=2 ) images.save_image(grid_image, p.outpath_grids, "grid", grid=True, p=p) result_images.insert(0, grid_image) return Processed(p, result_images, p.seed, "", all_prompts=all_prompts, infotexts=infotexts)