import os
import sys
# Environment variables
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
os.environ['GRADIO_ANALYTICS_ENABLED'] = '0'
# No need to adjust sys.path if using proper module imports
import subprocess
import gradio as gr
from PIL import Image
import torch
import uuid
import shutil
import json
import yaml
from slugify import slugify
from transformers import AutoProcessor, AutoModelForCausalLM
from gradio_logsview import LogsView, LogsViewRunner
from huggingface_hub import hf_hub_download, HfApi
from fluxgym_main.library import flux_train_utils, huggingface_util
from argparse import Namespace
from fluxgym_main import train_network
import toml
import re
MAX_IMAGES = 150
with open('models.yaml', 'r') as file:
models = yaml.safe_load(file)
def readme(base_model, lora_name, instance_prompt, sample_prompts):
# model license
model_config = models[base_model]
model_file = model_config["file"]
base_model_name = model_config["base"]
license = None
license_name = None
license_link = None
license_items = []
if "license" in model_config:
license = model_config["license"]
license_items.append(f"license: {license}")
if "license_name" in model_config:
license_name = model_config["license_name"]
license_items.append(f"license_name: {license_name}")
if "license_link" in model_config:
license_link = model_config["license_link"]
license_items.append(f"license_link: {license_link}")
license_str = "\n".join(license_items)
print(f"license_items={license_items}")
print(f"license_str = {license_str}")
# tags
tags = [ "text-to-image", "flux", "lora", "diffusers", "template:sd-lora", "fluxgym" ]
# widgets
widgets = []
sample_image_paths = []
output_name = slugify(lora_name)
samples_dir = resolve_path_without_quotes(f"outputs/{output_name}/sample")
try:
for filename in os.listdir(samples_dir):
# Filename Schema: [name]_[steps]_[index]_[timestamp].png
match = re.search(r"_(\d+)_(\d+)_(\d+)\.png$", filename)
if match:
steps, index, timestamp = int(match.group(1)), int(match.group(2)), int(match.group(3))
sample_image_paths.append((steps, index, f"sample/{filename}"))
# Sort by numeric index
sample_image_paths.sort(key=lambda x: x[0], reverse=True)
final_sample_image_paths = sample_image_paths[:len(sample_prompts)]
final_sample_image_paths.sort(key=lambda x: x[1])
for i, prompt in enumerate(sample_prompts):
_, _, image_path = final_sample_image_paths[i]
widgets.append(
{
"text": prompt,
"output": {
"url": image_path
},
}
)
except:
print(f"no samples")
dtype = "torch.bfloat16"
# Construct the README content
readme_content = f"""---
tags:
{yaml.dump(tags, indent=4).strip()}
{"widget:" if os.path.isdir(samples_dir) else ""}
{yaml.dump(widgets, indent=4).strip() if widgets else ""}
base_model: {base_model_name}
{"instance_prompt: " + instance_prompt if instance_prompt else ""}
{license_str}
---
# {lora_name}
A Flux LoRA trained on a local computer with [Fluxgym](https://github.com/cocktailpeanut/fluxgym)
<Gallery />
## Trigger words
{"You should use `" + instance_prompt + "` to trigger the image generation." if instance_prompt else "No trigger words defined."}
## Download model and use it with ComfyUI, AUTOMATIC1111, SD.Next, Invoke AI, Forge, etc.
Weights for this model are available in Safetensors format.
"""
return readme_content
def account_hf():
try:
with open("HF_TOKEN", "r") as file:
token = file.read()
api = HfApi(token=token)
try:
account = api.whoami()
return { "token": token, "account": account['name'] }
except:
return None
except:
return None
"""
hf_logout.click(fn=logout_hf, outputs=[hf_token, hf_login, hf_logout, repo_owner])
"""
def logout_hf():
os.remove("HF_TOKEN")
global current_account
current_account = account_hf()
print(f"current_account={current_account}")
return gr.update(value=""), gr.update(visible=True), gr.update(visible=False), gr.update(value="", visible=False)
"""
hf_login.click(fn=login_hf, inputs=[hf_token], outputs=[hf_token, hf_login, hf_logout, repo_owner])
"""
def login_hf(hf_token):
api = HfApi(token=hf_token)
try:
account = api.whoami()
if account != None:
if "name" in account:
with open("HF_TOKEN", "w") as file:
file.write(hf_token)
global current_account
current_account = account_hf()
return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(value=current_account["account"], visible=True)
return gr.update(), gr.update(), gr.update(), gr.update()
except:
print(f"incorrect hf_token")
return gr.update(), gr.update(), gr.update(), gr.update()
def upload_hf(base_model, lora_rows, repo_owner, repo_name, repo_visibility, hf_token):
src = lora_rows
repo_id = f"{repo_owner}/{repo_name}"
gr.Info(f"Uploading to Huggingface. Please Stand by...", duration=None)
args = Namespace(
huggingface_repo_id=repo_id,
huggingface_repo_type="model",
huggingface_repo_visibility=repo_visibility,
huggingface_path_in_repo="",
huggingface_token=hf_token,
async_upload=False
)
print(f"upload_hf args={args}")
huggingface_util.upload(args=args, src=src)
gr.Info(f"[Upload Complete] https://huggingface.co/{repo_id}", duration=None)
def load_captioning(uploaded_files, concept_sentence):
uploaded_images = [file for file in uploaded_files if not file.endswith('.txt')]
txt_files = [file for file in uploaded_files if file.endswith('.txt')]
txt_files_dict = {os.path.splitext(os.path.basename(txt_file))[0]: txt_file for txt_file in txt_files}
updates = []
if len(uploaded_images) <= 1:
raise gr.Error(
"Please upload at least 2 images to train your model (the ideal number with default settings is between 4-30)"
)
elif len(uploaded_images) > MAX_IMAGES:
raise gr.Error(f"For now, only {MAX_IMAGES} or less images are allowed for training")
# Update for the captioning_area
# for _ in range(3):
updates.append(gr.update(visible=True))
# Update visibility and image for each captioning row and image
for i in range(1, MAX_IMAGES + 1):
# Determine if the current row and image should be visible
visible = i <= len(uploaded_images)
# Update visibility of the captioning row
updates.append(gr.update(visible=visible))
# Update for image component - display image if available, otherwise hide
image_value = uploaded_images[i - 1] if visible else None
updates.append(gr.update(value=image_value, visible=visible))
corresponding_caption = False
if(image_value):
base_name = os.path.splitext(os.path.basename(image_value))[0]
if base_name in txt_files_dict:
with open(txt_files_dict[base_name], 'r') as file:
corresponding_caption = file.read()
# Update value of captioning area
text_value = corresponding_caption if visible and corresponding_caption else concept_sentence if visible and concept_sentence else None
updates.append(gr.update(value=text_value, visible=visible))
# Update for the sample caption area
updates.append(gr.update(visible=True))
updates.append(gr.update(visible=True))
return updates
def hide_captioning():
return gr.update(visible=False), gr.update(visible=False)
def resize_image(image_path, output_path, size):
with Image.open(image_path) as img:
width, height = img.size
if width < height:
new_width = size
new_height = int((size/width) * height)
else:
new_height = size
new_width = int((size/height) * width)
print(f"resize {image_path} : {new_width}x{new_height}")
img_resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
img_resized.save(output_path)
def create_dataset(destination_folder, size, *inputs):
print("Creating dataset")
images = inputs[0]
if not os.path.exists(destination_folder):
os.makedirs(destination_folder)
for index, image in enumerate(images):
# copy the images to the datasets folder
new_image_path = shutil.copy(image, destination_folder)
# if it's a caption text file skip the next bit
ext = os.path.splitext(new_image_path)[-1].lower()
if ext == '.txt':
continue
# resize the images
resize_image(new_image_path, new_image_path, size)
# copy the captions
original_caption = inputs[index + 1]
image_file_name = os.path.basename(new_image_path)
caption_file_name = os.path.splitext(image_file_name)[0] + ".txt"
caption_path = resolve_path_without_quotes(os.path.join(destination_folder, caption_file_name))
print(f"image_path={new_image_path}, caption_path = {caption_path}, original_caption={original_caption}")
# if caption_path exists, do not write
if os.path.exists(caption_path):
print(f"{caption_path} already exists. use the existing .txt file")
else:
print(f"{caption_path} create a .txt caption file")
with open(caption_path, 'w') as file:
file.write(original_caption)
print(f"destination_folder {destination_folder}")
return destination_folder
def run_captioning(images, concept_sentence, *captions):
print(f"run_captioning")
print(f"concept sentence {concept_sentence}")
print(f"captions {captions}")
#Load internally to not consume resources for training
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"device={device}")
torch_dtype = torch.float16
model = AutoModelForCausalLM.from_pretrained(
"multimodalart/Florence-2-large-no-flash-attn", torch_dtype=torch_dtype, trust_remote_code=True
).to(device)
processor = AutoProcessor.from_pretrained("multimodalart/Florence-2-large-no-flash-attn", trust_remote_code=True)
captions = list(captions)
for i, image_path in enumerate(images):
print(captions[i])
if isinstance(image_path, str): # If image is a file path
image = Image.open(image_path).convert("RGB")
prompt = "<DETAILED_CAPTION>"
inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
print(f"inputs {inputs}")
generated_ids = model.generate(
input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3
)
print(f"generated_ids {generated_ids}")
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
print(f"generated_text: {generated_text}")
parsed_answer = processor.post_process_generation(
generated_text, task=prompt, image_size=(image.width, image.height)
)
print(f"parsed_answer = {parsed_answer}")
caption_text = parsed_answer["<DETAILED_CAPTION>"].replace("The image shows ", "")
print(f"caption_text = {caption_text}, concept_sentence={concept_sentence}")
if concept_sentence:
caption_text = f"{concept_sentence} {caption_text}"
captions[i] = caption_text
yield captions
model.to("cpu")
del model
del processor
if torch.cuda.is_available():
torch.cuda.empty_cache()
def recursive_update(d, u):
for k, v in u.items():
if isinstance(v, dict) and v:
d[k] = recursive_update(d.get(k, {}), v)
else:
d[k] = v
return d
def download(base_model):
model = models[base_model]
model_file = model["file"]
repo = model["repo"]
# download unet
if base_model == "flux-dev" or base_model == "flux-schnell":
unet_folder = "models/unet"
else:
unet_folder = f"models/unet/{repo}"
unet_path = os.path.join(unet_folder, model_file)
if not os.path.exists(unet_path):
os.makedirs(unet_folder, exist_ok=True)
gr.Info(f"Downloading base model: {base_model}. Please wait. (You can check the terminal for the download progress)", duration=None)
print(f"download {base_model}")
hf_hub_download(repo_id=repo, local_dir=unet_folder, filename=model_file)
# download vae
vae_folder = "models/vae"
vae_path = os.path.join(vae_folder, "ae.sft")
if not os.path.exists(vae_path):
os.makedirs(vae_folder, exist_ok=True)
gr.Info(f"Downloading vae")
print(f"downloading ae.sft...")
hf_hub_download(repo_id="cocktailpeanut/xulf-dev", local_dir=vae_folder, filename="ae.sft")
# download clip
clip_folder = "models/clip"
clip_l_path = os.path.join(clip_folder, "clip_l.safetensors")
if not os.path.exists(clip_l_path):
os.makedirs(clip_folder, exist_ok=True)
gr.Info(f"Downloading clip...")
print(f"download clip_l.safetensors")
hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", local_dir=clip_folder, filename="clip_l.safetensors")
# download t5xxl
t5xxl_path = os.path.join(clip_folder, "t5xxl_fp16.safetensors")
if not os.path.exists(t5xxl_path):
print(f"download t5xxl_fp16.safetensors")
gr.Info(f"Downloading t5xxl...")
hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", local_dir=clip_folder, filename="t5xxl_fp16.safetensors")
def resolve_path(p):
current_dir = os.path.dirname(os.path.abspath(__file__))
norm_path = os.path.normpath(os.path.join(current_dir, p))
return f"\"{norm_path}\""
def resolve_path_without_quotes(p):
current_dir = os.path.dirname(os.path.abspath(__file__))
norm_path = os.path.normpath(os.path.join(current_dir, p))
return norm_path
def gen_sh(
base_model,
output_name,
resolution,
seed,
workers,
learning_rate,
network_dim,
max_train_epochs,
save_every_n_epochs,
timestep_sampling,
guidance_scale,
vram,
sample_prompts,
sample_every_n_steps,
*advanced_components
):
print(f"gen_sh: network_dim:{network_dim}, max_train_epochs={max_train_epochs}, save_every_n_epochs={save_every_n_epochs}, timestep_sampling={timestep_sampling}, guidance_scale={guidance_scale}, vram={vram}, sample_prompts={sample_prompts}, sample_every_n_steps={sample_every_n_steps}")
output_dir = resolve_path(f"outputs/{output_name}")
sample_prompts_path = resolve_path(f"outputs/{output_name}/sample_prompts.txt")
line_break = "\\"
file_type = "sh"
if sys.platform == "win32":
line_break = "^"
file_type = "bat"
############# Sample args ########################
sample = ""
if len(sample_prompts) > 0 and sample_every_n_steps > 0:
sample = f"""--sample_prompts={sample_prompts_path} --sample_every_n_steps="{sample_every_n_steps}" {line_break}"""
############# Optimizer args ########################
# if vram == "8G":
# optimizer = f"""--optimizer_type adafactor {line_break}
# --optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" {line_break}
# --split_mode {line_break}
# --network_args "train_blocks=single" {line_break}
# --lr_scheduler constant_with_warmup {line_break}
# --max_grad_norm 0.0 {line_break}"""
if vram == "16G":
# 16G VRAM
optimizer = f"""--optimizer_type adafactor {line_break}
--optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" {line_break}
--lr_scheduler constant_with_warmup {line_break}
--max_grad_norm 0.0 {line_break}"""
elif vram == "12G":
# 12G VRAM
optimizer = f"""--optimizer_type adafactor {line_break}
--optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" {line_break}
--split_mode {line_break}
--network_args "train_blocks=single" {line_break}
--lr_scheduler constant_with_warmup {line_break}
--max_grad_norm 0.0 {line_break}"""
else:
# 20G+ VRAM
optimizer = f"--optimizer_type adamw8bit {line_break}"
#######################################################
model_config = models[base_model]
model_file = model_config["file"]
repo = model_config["repo"]
if base_model == "flux-dev" or base_model == "flux-schnell":
model_folder = "models/unet"
else:
model_folder = f"models/unet/{repo}"
model_path = os.path.join(model_folder, model_file)
pretrained_model_path = resolve_path(model_path)
clip_path = resolve_path("models/clip/clip_l.safetensors")
t5_path = resolve_path("models/clip/t5xxl_fp16.safetensors")
ae_path = resolve_path("models/vae/ae.sft")
sh = f"""accelerate launch {line_break}
--mixed_precision bf16 {line_break}
--num_cpu_threads_per_process 1 {line_break}
sd-scripts/flux_train_network.py {line_break}
--pretrained_model_name_or_path {pretrained_model_path} {line_break}
--clip_l {clip_path} {line_break}
--t5xxl {t5_path} {line_break}
--ae {ae_path} {line_break}
--cache_latents_to_disk {line_break}
--save_model_as safetensors {line_break}
--sdpa --persistent_data_loader_workers {line_break}
--max_data_loader_n_workers {workers} {line_break}
--seed {seed} {line_break}
--gradient_checkpointing {line_break}
--mixed_precision bf16 {line_break}
--save_precision bf16 {line_break}
--network_module networks.lora_flux {line_break}
--network_dim {network_dim} {line_break}
{optimizer}{sample}
--learning_rate {learning_rate} {line_break}
--cache_text_encoder_outputs {line_break}
--cache_text_encoder_outputs_to_disk {line_break}
--fp8_base {line_break}
--highvram {line_break}
--max_train_epochs {max_train_epochs} {line_break}
--save_every_n_epochs {save_every_n_epochs} {line_break}
--dataset_config {resolve_path(f"outputs/{output_name}/dataset.toml")} {line_break}
--output_dir {output_dir} {line_break}
--output_name {output_name} {line_break}
--timestep_sampling {timestep_sampling} {line_break}
--discrete_flow_shift 3.1582 {line_break}
--model_prediction_type raw {line_break}
--guidance_scale {guidance_scale} {line_break}
--loss_type l2 {line_break}"""
############# Advanced args ########################
global advanced_component_ids
global original_advanced_component_values
# check dirty
print(f"original_advanced_component_values = {original_advanced_component_values}")
advanced_flags = []
for i, current_value in enumerate(advanced_components):
# print(f"compare {advanced_component_ids[i]}: old={original_advanced_component_values[i]}, new={current_value}")
if original_advanced_component_values[i] != current_value:
# dirty
if current_value == True:
# Boolean
advanced_flags.append(advanced_component_ids[i])
else:
# string
advanced_flags.append(f"{advanced_component_ids[i]} {current_value}")
if len(advanced_flags) > 0:
advanced_flags_str = f" {line_break}\n ".join(advanced_flags)
sh = sh + "\n " + advanced_flags_str
return sh
def gen_toml(
dataset_folder,
resolution,
class_tokens,
num_repeats
):
toml = f"""[general]
shuffle_caption = false
caption_extension = '.txt'
keep_tokens = 1
[[datasets]]
resolution = {resolution}
batch_size = 1
keep_tokens = 1
[[datasets.subsets]]
image_dir = '{resolve_path_without_quotes(dataset_folder)}'
class_tokens = '{class_tokens}'
num_repeats = {num_repeats}"""
return toml
def update_total_steps(max_train_epochs, num_repeats, images):
try:
num_images = len(images)
total_steps = max_train_epochs * num_images * num_repeats
print(f"max_train_epochs={max_train_epochs} num_images={num_images}, num_repeats={num_repeats}, total_steps={total_steps}")
return gr.update(value = total_steps)
except:
print("")
def set_repo(lora_rows):
selected_name = os.path.basename(lora_rows)
return gr.update(value=selected_name)
def get_loras():
try:
outputs_path = resolve_path_without_quotes(f"outputs")
files = os.listdir(outputs_path)
folders = [os.path.join(outputs_path, item) for item in files if os.path.isdir(os.path.join(outputs_path, item)) and item != "sample"]
folders.sort(key=lambda file: os.path.getctime(file), reverse=True)
return folders
except Exception as e:
return []
def get_samples(lora_name):
output_name = slugify(lora_name)
try:
samples_path = resolve_path_without_quotes(f"outputs/{output_name}/sample")
files = [os.path.join(samples_path, file) for file in os.listdir(samples_path)]
files.sort(key=lambda file: os.path.getctime(file), reverse=True)
return files
except:
return []
def start_training(
base_model,
lora_name,
train_script,
train_config,
sample_prompts,
):
# write custom script and toml
if not os.path.exists("models"):
os.makedirs("models", exist_ok=True)
if not os.path.exists("outputs"):
os.makedirs("outputs", exist_ok=True)
output_name = slugify(lora_name)
output_dir = resolve_path_without_quotes(f"outputs/{output_name}")
if not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
download(base_model)
file_type = "sh"
if sys.platform == "win32":
file_type = "bat"
sh_filename = f"train.{file_type}"
sh_filepath = resolve_path_without_quotes(f"outputs/{output_name}/{sh_filename}")
with open(sh_filepath, 'w', encoding="utf-8") as file:
file.write(train_script)
gr.Info(f"Generated train script at {sh_filename}")
dataset_path = resolve_path_without_quotes(f"outputs/{output_name}/dataset.toml")
with open(dataset_path, 'w', encoding="utf-8") as file:
file.write(train_config)
gr.Info(f"Generated dataset.toml")
sample_prompts_path = resolve_path_without_quotes(f"outputs/{output_name}/sample_prompts.txt")
with open(sample_prompts_path, 'w', encoding='utf-8') as file:
file.write(sample_prompts)
gr.Info(f"Generated sample_prompts.txt")
# Train
if sys.platform == "win32":
command = sh_filepath
else:
command = f"bash \"{sh_filepath}\""
# Use Popen to run the command and capture output in real-time
env = os.environ.copy()
env['PYTHONIOENCODING'] = 'utf-8'
env['LOG_LEVEL'] = 'DEBUG'
runner = LogsViewRunner()
cwd = os.path.dirname(os.path.abspath(__file__))
gr.Info(f"Started training")
yield from runner.run_command([command], cwd=cwd)
yield runner.log(f"Runner: {runner}")
# Generate Readme
config = toml.loads(train_config)
concept_sentence = config['datasets'][0]['subsets'][0]['class_tokens']
print(f"concept_sentence={concept_sentence}")
print(f"lora_name {lora_name}, concept_sentence={concept_sentence}, output_name={output_name}")
sample_prompts_path = resolve_path_without_quotes(f"outputs/{output_name}/sample_prompts.txt")
with open(sample_prompts_path, "r", encoding="utf-8") as f:
lines = f.readlines()
sample_prompts = [line.strip() for line in lines if len(line.strip()) > 0 and line[0] != "#"]
md = readme(base_model, lora_name, concept_sentence, sample_prompts)
readme_path = resolve_path_without_quotes(f"outputs/{output_name}/README.md")
with open(readme_path, "w", encoding="utf-8") as f:
f.write(md)
gr.Info(f"Training Complete. Check the outputs folder for the LoRA files.", duration=None)
def update(
base_model,
lora_name,
resolution,
seed,
workers,
class_tokens,
learning_rate,
network_dim,
max_train_epochs,
save_every_n_epochs,
timestep_sampling,
guidance_scale,
vram,
num_repeats,
sample_prompts,
sample_every_n_steps,
*advanced_components,
):
output_name = slugify(lora_name)
dataset_folder = str(f"datasets/{output_name}")
sh = gen_sh(
base_model,
output_name,
resolution,
seed,
workers,
learning_rate,
network_dim,
max_train_epochs,
save_every_n_epochs,
timestep_sampling,
guidance_scale,
vram,
sample_prompts,
sample_every_n_steps,
*advanced_components,
)
toml = gen_toml(
dataset_folder,
resolution,
class_tokens,
num_repeats
)
return gr.update(value=sh), gr.update(value=toml), dataset_folder
"""
demo.load(fn=loaded, js=js, outputs=[hf_token, hf_login, hf_logout, hf_account])
"""
def loaded():
global current_account
current_account = account_hf()
print(f"current_account={current_account}")
if current_account != None:
return gr.update(value=current_account["token"]), gr.update(visible=False), gr.update(visible=True), gr.update(value=current_account["account"], visible=True)
else:
return gr.update(value=""), gr.update(visible=True), gr.update(visible=False), gr.update(value="", visible=False)
def update_sample(concept_sentence):
return gr.update(value=concept_sentence)
def refresh_publish_tab():
loras = get_loras()
return gr.Dropdown(label="Trained LoRAs", choices=loras)
def init_advanced():
# if basic_args
basic_args = {
'pretrained_model_name_or_path',
'clip_l',
't5xxl',
'ae',
'cache_latents_to_disk',
'save_model_as',
'sdpa',
'persistent_data_loader_workers',
'max_data_loader_n_workers',
'seed',
'gradient_checkpointing',
'mixed_precision',
'save_precision',
'network_module',
'network_dim',
'learning_rate',
'cache_text_encoder_outputs',
'cache_text_encoder_outputs_to_disk',
'fp8_base',
'highvram',
'max_train_epochs',
'save_every_n_epochs',
'dataset_config',
'output_dir',
'output_name',
'timestep_sampling',
'discrete_flow_shift',
'model_prediction_type',
'guidance_scale',
'loss_type',
'optimizer_type',
'optimizer_args',
'lr_scheduler',
'sample_prompts',
'sample_every_n_steps',
'max_grad_norm',
'split_mode',
'network_args'
}
# generate a UI config
# if not in basic_args, create a simple form
parser = train_network.setup_parser()
flux_train_utils.add_flux_train_arguments(parser)
args_info = {}
for action in parser._actions:
if action.dest != 'help': # Skip the default help argument
# if the dest is included in basic_args
args_info[action.dest] = {
"action": action.option_strings, # Option strings like '--use_8bit_adam'
"type": action.type, # Type of the argument
"help": action.help, # Help message
"default": action.default, # Default value, if any
"required": action.required # Whether the argument is required
}
temp = []
for key in args_info:
temp.append({ 'key': key, 'action': args_info[key] })
temp.sort(key=lambda x: x['key'])
advanced_component_ids = []
advanced_components = []
for item in temp:
key = item['key']
action = item['action']
if key in basic_args:
print("")
else:
action_type = str(action['type'])
component = None
with gr.Column(min_width=300):
if action_type == "None":
# radio
component = gr.Checkbox()
# elif action_type == "<class 'str'>":
# component = gr.Textbox()
# elif action_type == "<class 'int'>":
# component = gr.Number(precision=0)
# elif action_type == "<class 'float'>":
# component = gr.Number()
# elif "int_or_float" in action_type:
# component = gr.Number()
else:
component = gr.Textbox(value="")
if component != None:
component.interactive = True
component.elem_id = action['action'][0]
component.label = component.elem_id
component.elem_classes = ["advanced"]
if action['help'] != None:
component.info = action['help']
advanced_components.append(component)
advanced_component_ids.append(component.elem_id)
return advanced_components, advanced_component_ids
theme = gr.themes.Monochrome(
text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"),
font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui", "sans-serif"],
)
css = """
@keyframes rotate {
0% {
transform: rotate(0deg);
}
100% {
transform: rotate(360deg);
}
}
#advanced_options .advanced:nth-child(even) { background: rgba(0,0,100,0.04) !important; }
h1{font-family: georgia; font-style: italic; font-weight: bold; font-size: 30px; letter-spacing: -1px;}
h3{margin-top: 0}
.tabitem{border: 0px}
.group_padding{}
nav{position: fixed; top: 0; left: 0; right: 0; z-index: 1000; text-align: center; padding: 10px; box-sizing: border-box; display: flex; align-items: center; backdrop-filter: blur(10px); }
nav button { background: none; color: firebrick; font-weight: bold; border: 2px solid firebrick; padding: 5px 10px; border-radius: 5px; font-size: 14px; }
nav img { height: 40px; width: 40px; border-radius: 40px; }
nav img.rotate { animation: rotate 2s linear infinite; }
.flexible { flex-grow: 1; }
.tast-details { margin: 10px 0 !important; }
.toast-wrap { bottom: var(--size-4) !important; top: auto !important; border: none !important; backdrop-filter: blur(10px); }
.toast-title, .toast-text, .toast-icon, .toast-close { color: black !important; font-size: 14px; }
.toast-body { border: none !important; }
#terminal { box-shadow: none !important; margin-bottom: 25px; background: rgba(0,0,0,0.03); }
#terminal .generating { border: none !important; }
#terminal label { position: absolute !important; }
.tabs { margin-top: 50px; }
.hidden { display: none !important; }
.codemirror-wrapper .cm-line { font-size: 12px !important; }
label { font-weight: bold !important; }
#start_training.clicked { background: silver; color: black; }
"""
js = """
function() {
let autoscroll = document.querySelector("#autoscroll")
if (window.iidxx) {
window.clearInterval(window.iidxx);
}
window.iidxx = window.setInterval(function() {
let text=document.querySelector(".codemirror-wrapper .cm-line").innerText.trim()
let img = document.querySelector("#logo")
if (text.length > 0) {
autoscroll.classList.remove("hidden")
if (autoscroll.classList.contains("on")) {
autoscroll.textContent = "Autoscroll ON"
window.scrollTo(0, document.body.scrollHeight, { behavior: "smooth" });
img.classList.add("rotate")
} else {
autoscroll.textContent = "Autoscroll OFF"
img.classList.remove("rotate")
}
}
}, 500);
console.log("autoscroll", autoscroll)
autoscroll.addEventListener("click", (e) => {
autoscroll.classList.toggle("on")
})
function debounce(fn, delay) {
let timeoutId;
return function(...args) {
clearTimeout(timeoutId);
timeoutId = setTimeout(() => fn(...args), delay);
};
}
function handleClick() {
console.log("refresh")
document.querySelector("#refresh").click();
}
const debouncedClick = debounce(handleClick, 1000);
document.addEventListener("input", debouncedClick);
document.querySelector("#start_training").addEventListener("click", (e) => {
e.target.classList.add("clicked")
e.target.innerHTML = "Training..."
})
}
"""
current_account = account_hf()
print(f"current_account={current_account}")
with gr.Blocks(elem_id="app", theme=theme, css=css, fill_width=True) as demo:
with gr.Tabs() as tabs:
with gr.TabItem("Gym"):
output_components = []
with gr.Row():
gr.HTML("""<nav>
<img id='logo' src='/file=icon.png' width='80' height='80'>
<div class='flexible'></div>
<button id='autoscroll' class='on hidden'></button>
</nav>
""")
with gr.Row(elem_id='container'):
with gr.Column():
gr.Markdown(
"""# Step 1. LoRA Info
<p style="margin-top:0">Configure your LoRA train settings.</p>
""", elem_classes="group_padding")
lora_name = gr.Textbox(
label="The name of your LoRA",
info="This has to be a unique name",
placeholder="e.g.: Persian Miniature Painting style, Cat Toy",
)
concept_sentence = gr.Textbox(
elem_id="--concept_sentence",
label="Trigger word/sentence",
info="Trigger word or sentence to be used",
placeholder="uncommon word like p3rs0n or trtcrd, or sentence like 'in the style of CNSTLL'",
interactive=True,
)
model_names = list(models.keys())
print(f"model_names={model_names}")
base_model = gr.Dropdown(label="Base model (edit the models.yaml file to add more to this list)", choices=model_names, value=model_names[0])
vram = gr.Radio(["20G", "16G", "12G" ], value="20G", label="VRAM", interactive=True)
num_repeats = gr.Number(value=10, precision=0, label="Repeat trains per image", interactive=True)
max_train_epochs = gr.Number(label="Max Train Epochs", value=16, interactive=True)
total_steps = gr.Number(0, interactive=False, label="Expected training steps")
sample_prompts = gr.Textbox("", lines=5, label="Sample Image Prompts (Separate with new lines)", interactive=True)
sample_every_n_steps = gr.Number(0, precision=0, label="Sample Image Every N Steps", interactive=True)
resolution = gr.Number(value=512, precision=0, label="Resize dataset images")
with gr.Column():
gr.Markdown(
"""# Step 2. Dataset
<p style="margin-top:0">Make sure the captions include the trigger word.</p>
""", elem_classes="group_padding")
with gr.Group():
images = gr.File(
file_types=["image", ".txt"],
label="Upload your images",
#info="If you want, you can also manually upload caption files that match the image names (example: img0.png => img0.txt)",
file_count="multiple",
interactive=True,
visible=True,
scale=1,
)
with gr.Group(visible=False) as captioning_area:
do_captioning = gr.Button("Add AI captions with Florence-2")
output_components.append(captioning_area)
#output_components = [captioning_area]
caption_list = []
for i in range(1, MAX_IMAGES + 1):
locals()[f"captioning_row_{i}"] = gr.Row(visible=False)
with locals()[f"captioning_row_{i}"]:
locals()[f"image_{i}"] = gr.Image(
type="filepath",
width=111,
height=111,
min_width=111,
interactive=False,
scale=2,
show_label=False,
show_share_button=False,
show_download_button=False,
)
locals()[f"caption_{i}"] = gr.Textbox(
label=f"Caption {i}", scale=15, interactive=True
)
output_components.append(locals()[f"captioning_row_{i}"])
output_components.append(locals()[f"image_{i}"])
output_components.append(locals()[f"caption_{i}"])
caption_list.append(locals()[f"caption_{i}"])
with gr.Column():
gr.Markdown(
"""# Step 3. Train
<p style="margin-top:0">Press start to start training.</p>
""", elem_classes="group_padding")
refresh = gr.Button("Refresh", elem_id="refresh", visible=False)
start = gr.Button("Start training", visible=False, elem_id="start_training")
output_components.append(start)
train_script = gr.Textbox(label="Train script", max_lines=100, interactive=True)
train_config = gr.Textbox(label="Train config", max_lines=100, interactive=True)
with gr.Accordion("Advanced options", elem_id='advanced_options', open=False):
with gr.Row():
with gr.Column(min_width=300):
seed = gr.Number(label="--seed", info="Seed", value=42, interactive=True)
with gr.Column(min_width=300):
workers = gr.Number(label="--max_data_loader_n_workers", info="Number of Workers", value=2, interactive=True)
with gr.Column(min_width=300):
learning_rate = gr.Textbox(label="--learning_rate", info="Learning Rate", value="8e-4", interactive=True)
with gr.Column(min_width=300):
save_every_n_epochs = gr.Number(label="--save_every_n_epochs", info="Save every N epochs", value=4, interactive=True)
with gr.Column(min_width=300):
guidance_scale = gr.Number(label="--guidance_scale", info="Guidance Scale", value=1.0, interactive=True)
with gr.Column(min_width=300):
timestep_sampling = gr.Textbox(label="--timestep_sampling", info="Timestep Sampling", value="shift", interactive=True)
with gr.Column(min_width=300):
network_dim = gr.Number(label="--network_dim", info="LoRA Rank", value=4, minimum=4, maximum=128, step=4, interactive=True)
advanced_components, advanced_component_ids = init_advanced()
with gr.Row():
terminal = LogsView(label="Train log", elem_id="terminal")
with gr.Row():
gallery = gr.Gallery(get_samples, inputs=[lora_name], label="Samples", every=10, columns=6)
with gr.TabItem("Publish") as publish_tab:
hf_token = gr.Textbox(label="Huggingface Token")
hf_login = gr.Button("Login")
hf_logout = gr.Button("Logout")
with gr.Row() as row:
gr.Markdown("**LoRA**")
gr.Markdown("**Upload**")
loras = get_loras()
with gr.Row():
lora_rows = refresh_publish_tab()
with gr.Column():
with gr.Row():
repo_owner = gr.Textbox(label="Account", interactive=False)
repo_name = gr.Textbox(label="Repository Name")
repo_visibility = gr.Textbox(label="Repository Visibility ('public' or 'private')", value="public")
upload_button = gr.Button("Upload to HuggingFace")
upload_button.click(
fn=upload_hf,
inputs=[
base_model,
lora_rows,
repo_owner,
repo_name,
repo_visibility,
hf_token,
]
)
hf_login.click(fn=login_hf, inputs=[hf_token], outputs=[hf_token, hf_login, hf_logout, repo_owner])
hf_logout.click(fn=logout_hf, outputs=[hf_token, hf_login, hf_logout, repo_owner])
publish_tab.select(refresh_publish_tab, outputs=lora_rows)
lora_rows.select(fn=set_repo, inputs=[lora_rows], outputs=[repo_name])
dataset_folder = gr.State()
listeners = [
base_model,
lora_name,
resolution,
seed,
workers,
concept_sentence,
learning_rate,
network_dim,
max_train_epochs,
save_every_n_epochs,
timestep_sampling,
guidance_scale,
vram,
num_repeats,
sample_prompts,
sample_every_n_steps,
*advanced_components
]
advanced_component_ids = [x.elem_id for x in advanced_components]
original_advanced_component_values = [comp.value for comp in advanced_components]
images.upload(
load_captioning,
inputs=[images, concept_sentence],
outputs=output_components
)
images.delete(
load_captioning,
inputs=[images, concept_sentence],
outputs=output_components
)
images.clear(
hide_captioning,
outputs=[captioning_area, start]
)
max_train_epochs.change(
fn=update_total_steps,
inputs=[max_train_epochs, num_repeats, images],
outputs=[total_steps]
)
num_repeats.change(
fn=update_total_steps,
inputs=[max_train_epochs, num_repeats, images],
outputs=[total_steps]
)
images.upload(
fn=update_total_steps,
inputs=[max_train_epochs, num_repeats, images],
outputs=[total_steps]
)
images.delete(
fn=update_total_steps,
inputs=[max_train_epochs, num_repeats, images],
outputs=[total_steps]
)
images.clear(
fn=update_total_steps,
inputs=[max_train_epochs, num_repeats, images],
outputs=[total_steps]
)
concept_sentence.change(fn=update_sample, inputs=[concept_sentence], outputs=sample_prompts)
start.click(fn=create_dataset, inputs=[dataset_folder, resolution, images] + caption_list, outputs=dataset_folder).then(
fn=start_training,
inputs=[
base_model,
lora_name,
train_script,
train_config,
sample_prompts,
],
outputs=terminal,
)
do_captioning.click(fn=run_captioning, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)
demo.load(fn=loaded, js=js, outputs=[hf_token, hf_login, hf_logout, repo_owner])
refresh.click(update, inputs=listeners, outputs=[train_script, train_config, dataset_folder])
if __name__ == "__main__":
cwd = os.path.dirname(os.path.abspath(__file__))
demo.launch(debug=True, show_error=True, allowed_paths=[cwd])