Skip to content

Commit

Permalink
Merge pull request #2450 from bmaltais/dev
Browse files Browse the repository at this point in the history
v24.1.0
  • Loading branch information
bmaltais committed May 5, 2024
2 parents 415d7b2 + d26dad2 commit 95f8883
Show file tree
Hide file tree
Showing 39 changed files with 1,060 additions and 943 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/typos.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,4 @@ jobs:
- uses: actions/checkout@v4

- name: typos-action
uses: crate-ci/typos@v1.19.0
uses: crate-ci/typos@v1.21.0
2 changes: 1 addition & 1 deletion .release
Original file line number Diff line number Diff line change
@@ -1 +1 @@
v24.0.9
v24.1.0
238 changes: 16 additions & 222 deletions README.md

Large diffs are not rendered by default.

4 changes: 1 addition & 3 deletions kohya_gui.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,13 +78,11 @@ def UI(**kwargs):
reg_data_dir_input=reg_data_dir_input,
output_dir_input=output_dir_input,
logging_dir_input=logging_dir_input,
enable_copy_info_button=True,
headless=headless,
config=config,
use_shell_flag=use_shell_flag,
)
with gr.Tab("LoRA"):
_ = LoRATools(headless=headless, use_shell_flag=use_shell_flag)
_ = LoRATools(headless=headless)
with gr.Tab("About"):
gr.Markdown(f"kohya_ss GUI release {release}")
with gr.Tab("README"):
Expand Down
74 changes: 39 additions & 35 deletions kohya_gui/basic_caption_gui.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import gradio as gr
from easygui import msgbox
import subprocess
from .common_gui import (
get_folder_path,
Expand Down Expand Up @@ -28,7 +27,6 @@ def caption_images(
postfix: str,
find_text: str,
replace_text: str,
use_shell: bool = False,
):
"""
Captions images in a given directory with a given caption text.
Expand All @@ -46,72 +44,78 @@ def caption_images(
Returns:
None
"""
# Check if images_dir is provided
# Check if images_dir and caption_ext are provided
missing_parameters = []
if not images_dir:
msgbox(
"Image folder is missing. Please provide the directory containing the images to caption."
)
return

# Check if caption_ext is provided
missing_parameters.append("image directory")
if not caption_ext:
msgbox("Please provide an extension for the caption files.")
missing_parameters.append("caption file extension")

if missing_parameters:
log.info(
"The following parameter(s) are missing: {}. "
"Please provide these to proceed with captioning the images.".format(", ".join(missing_parameters))
)
return

# Log the captioning process
if caption_text:
log.info(f"Captioning files in {images_dir} with {caption_text}...")

# Build the command to run caption.py
run_cmd = rf'"{PYTHON}" "{scriptdir}/tools/caption.py"'
run_cmd += f' --caption_text="{caption_text}"'
run_cmd = [
rf"{PYTHON}",
rf"{scriptdir}/tools/caption.py",
"--caption_text",
caption_text,
]

# Add optional flags to the command
if overwrite:
run_cmd += f" --overwrite"
run_cmd.append("--overwrite")
if caption_ext:
run_cmd += f' --caption_file_ext="{caption_ext}"'
run_cmd.append("--caption_file_ext")
run_cmd.append(caption_ext)

run_cmd += f' "{images_dir}"'
run_cmd.append(rf"{images_dir}")

# Log the command
log.info(run_cmd)
# Reconstruct the safe command string for display
command_to_run = " ".join(run_cmd)
log.info(f"Executing command: {command_to_run}")

# Set the environment variable for the Python path
env = os.environ.copy()
env["PYTHONPATH"] = (
f"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
)
env["TF_ENABLE_ONEDNN_OPTS"] = "0"

log.info(f"Executing command: {run_cmd} with shell={use_shell}")

# Run the command in the sd-scripts folder context
subprocess.run(run_cmd, env=env, shell=use_shell)

subprocess.run(run_cmd, env=env, shell=False)

# Check if overwrite option is enabled
if overwrite:
# Add prefix and postfix to caption files
if prefix or postfix:
# Add prefix and postfix to caption files or find and replace text in caption files
if prefix or postfix or find_text:
# Add prefix and/or postfix to caption files
add_pre_postfix(
folder=images_dir,
caption_file_ext=caption_ext,
prefix=prefix,
postfix=postfix,
)
# Find and replace text in caption files
if find_text:
find_replace(
folder_path=images_dir,
caption_file_ext=caption_ext,
search_text=find_text,
replace_text=replace_text,
)
# Replace specified text in caption files if find and replace text is provided
if find_text and replace_text:
find_replace(
folder_path=images_dir,
caption_file_ext=caption_ext,
search_text=find_text,
replace_text=replace_text,
)
else:
# Show a message if modification is not possible without overwrite option enabled
if prefix or postfix:
msgbox(
log.info(
'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected.'
)

Expand All @@ -120,7 +124,7 @@ def caption_images(


# Gradio UI
def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None, use_shell: bool = False):
def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None):
"""
Creates a Gradio tab for basic image captioning.
Expand Down Expand Up @@ -200,6 +204,7 @@ def list_images_dirs(path):
choices=[".cap", ".caption", ".txt"],
value=".txt",
interactive=True,
allow_custom_value=True,
)
# Checkbox to overwrite existing captions
overwrite = gr.Checkbox(
Expand Down Expand Up @@ -258,7 +263,6 @@ def list_images_dirs(path):
postfix,
find_text,
replace_text,
gr.Checkbox(value=use_shell, visible=False),
],
show_progress=False,
)
Expand Down
41 changes: 20 additions & 21 deletions kohya_gui/blip_caption_gui.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import gradio as gr
from easygui import msgbox
import subprocess
import os
import sys
Expand All @@ -23,7 +22,6 @@ def caption_images(
beam_search: bool,
prefix: str = "",
postfix: str = "",
use_shell: bool = False,
) -> None:
"""
Automatically generates captions for images in the specified directory using the BLIP model.
Expand All @@ -47,59 +45,60 @@ def caption_images(
"""
# Check if the image folder is provided
if not train_data_dir:
msgbox("Image folder is missing...")
log.info("Image folder is missing...")
return

# Check if the caption file extension is provided
if not caption_file_ext:
msgbox("Please provide an extension for the caption files.")
log.info("Please provide an extension for the caption files.")
return

log.info(f"Captioning files in {train_data_dir}...")

# Construct the command to run make_captions.py
run_cmd = [fr'"{PYTHON}"', fr'"{scriptdir}/sd-scripts/finetune/make_captions.py"']
run_cmd = [rf"{PYTHON}", rf"{scriptdir}/sd-scripts/finetune/make_captions.py"]

# Add required arguments
run_cmd.append('--batch_size')
run_cmd.append("--batch_size")
run_cmd.append(str(batch_size))
run_cmd.append('--num_beams')
run_cmd.append("--num_beams")
run_cmd.append(str(num_beams))
run_cmd.append('--top_p')
run_cmd.append("--top_p")
run_cmd.append(str(top_p))
run_cmd.append('--max_length')
run_cmd.append("--max_length")
run_cmd.append(str(max_length))
run_cmd.append('--min_length')
run_cmd.append("--min_length")
run_cmd.append(str(min_length))

# Add optional flags to the command
if beam_search:
run_cmd.append("--beam_search")
if caption_file_ext:
run_cmd.append('--caption_extension')
run_cmd.append("--caption_extension")
run_cmd.append(caption_file_ext)

# Add the directory containing the training data
run_cmd.append(fr'"{train_data_dir}"')
run_cmd.append(rf"{train_data_dir}")

# Add URL for caption model weights
run_cmd.append('--caption_weights')
run_cmd.append("https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth")
run_cmd.append("--caption_weights")
run_cmd.append(
rf"https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth"
)

# Set up the environment
env = os.environ.copy()
env["PYTHONPATH"] = (
f"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
)
env["TF_ENABLE_ONEDNN_OPTS"] = "0"

# Reconstruct the safe command string for display
command_to_run = " ".join(run_cmd)
log.info(f"Executing command: {command_to_run} with shell={use_shell}")

# Run the command in the sd-scripts folder context
subprocess.run(command_to_run, env=env, shell=use_shell, cwd=f"{scriptdir}/sd-scripts")
log.info(f"Executing command: {command_to_run}")

# Run the command in the sd-scripts folder context
subprocess.run(run_cmd, env=env, shell=False, cwd=rf"{scriptdir}/sd-scripts")

# Add prefix and postfix
add_pre_postfix(
Expand All @@ -117,7 +116,7 @@ def caption_images(
###


def gradio_blip_caption_gui_tab(headless=False, default_train_dir=None, use_shell: bool = False):
def gradio_blip_caption_gui_tab(headless=False, default_train_dir=None):
from .common_gui import create_refresh_button

default_train_dir = (
Expand Down Expand Up @@ -167,6 +166,7 @@ def list_train_dirs(path):
choices=[".cap", ".caption", ".txt"],
value=".txt",
interactive=True,
allow_custom_value=True,
)

prefix = gr.Textbox(
Expand Down Expand Up @@ -207,7 +207,6 @@ def list_train_dirs(path):
beam_search,
prefix,
postfix,
gr.Checkbox(value=use_shell, visible=False),
],
show_progress=False,
)
Expand Down
6 changes: 3 additions & 3 deletions kohya_gui/class_command_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def __init__(self, headless: bool = False):
"Stop training", visible=self.process is not None or headless, variant="stop"
)

def execute_command(self, run_cmd: str, use_shell: bool = False, **kwargs):
def execute_command(self, run_cmd: str, **kwargs):
"""
Execute a command if no other command is currently running.
Expand All @@ -44,10 +44,10 @@ def execute_command(self, run_cmd: str, use_shell: bool = False, **kwargs):

# Reconstruct the safe command string for display
command_to_run = " ".join(run_cmd)
log.info(f"Executing command: {command_to_run} with shell={use_shell}")
log.info(f"Executing command: {command_to_run}")

# Execute the command securely
self.process = subprocess.Popen(command_to_run, **kwargs, shell=use_shell)
self.process = subprocess.Popen(run_cmd, **kwargs)
log.info("Command executed.")

def kill_command(self):
Expand Down
15 changes: 7 additions & 8 deletions kohya_gui/class_lora_tab.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,14 @@ class LoRATools:
def __init__(
self,
headless: bool = False,
use_shell_flag: bool = False,
):
gr.Markdown("This section provide various LoRA tools...")
gradio_extract_dylora_tab(headless=headless, use_shell=use_shell_flag)
gradio_convert_lcm_tab(headless=headless, use_shell=use_shell_flag)
gradio_extract_lora_tab(headless=headless, use_shell=use_shell_flag)
gradio_extract_lycoris_locon_tab(headless=headless, use_shell=use_shell_flag)
gradio_merge_lora_tab = GradioMergeLoRaTab(use_shell=use_shell_flag)
gradio_merge_lycoris_tab(headless=headless, use_shell=use_shell_flag)
gradio_svd_merge_lora_tab(headless=headless, use_shell=use_shell_flag)
gradio_extract_dylora_tab(headless=headless)
gradio_convert_lcm_tab(headless=headless)
gradio_extract_lora_tab(headless=headless)
gradio_extract_lycoris_locon_tab(headless=headless)
gradio_merge_lora_tab = GradioMergeLoRaTab()
gradio_merge_lycoris_tab(headless=headless)
gradio_svd_merge_lora_tab(headless=headless)
gradio_resize_lora_tab(headless=headless)
gradio_verify_lora_tab(headless=headless)
Loading

0 comments on commit 95f8883

Please sign in to comment.