File size: 10,858 Bytes
b6ff6dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 |
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"<lora:{lora_name}:{preferred_weight}>, {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"<lora:{lora_file.stem}:1>,"
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)
|