Upload loraiterate.py

loraiterate.py

@@ -0,0 +1,275 @@
+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)