KohyaSS/library/convert_model_gui.py

248 lines
9.1 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import gradio as gr
from easygui import msgbox
import subprocess
import os
import shutil
from .common_gui import get_folder_path, get_file_path
folder_symbol = '\U0001f4c2' # 📂
refresh_symbol = '\U0001f504' # 🔄
save_style_symbol = '\U0001f4be' # 💾
document_symbol = '\U0001F4C4' # 📄
PYTHON = 'python3' if os.name == 'posix' else './venv/Scripts/python.exe'
def convert_model(
source_model_input,
source_model_type,
target_model_folder_input,
target_model_name_input,
target_model_type,
target_save_precision_type,
):
# Check for caption_text_input
if source_model_type == '':
msgbox('Invalid source model type')
return
# Check if source model exist
if os.path.isfile(source_model_input):
print('The provided source model is a file')
elif os.path.isdir(source_model_input):
print('The provided model is a folder')
else:
msgbox('The provided source model is neither a file nor a folder')
return
# Check if source model exist
if os.path.isdir(target_model_folder_input):
print('The provided model folder exist')
else:
msgbox('The provided target folder does not exist')
return
run_cmd = f'{PYTHON} "tools/convert_diffusers20_original_sd.py"'
v1_models = [
'runwayml/stable-diffusion-v1-5',
'CompVis/stable-diffusion-v1-4',
]
# check if v1 models
if str(source_model_type) in v1_models:
print('SD v1 model specified. Setting --v1 parameter')
run_cmd += ' --v1'
else:
print('SD v2 model specified. Setting --v2 parameter')
run_cmd += ' --v2'
if not target_save_precision_type == 'unspecified':
run_cmd += f' --{target_save_precision_type}'
if (
target_model_type == 'diffuser'
or target_model_type == 'diffuser_safetensors'
):
run_cmd += f' --reference_model="{source_model_type}"'
if target_model_type == 'diffuser_safetensors':
run_cmd += ' --use_safetensors'
run_cmd += f' "{source_model_input}"'
if (
target_model_type == 'diffuser'
or target_model_type == 'diffuser_safetensors'
):
target_model_path = os.path.join(
target_model_folder_input, target_model_name_input
)
run_cmd += f' "{target_model_path}"'
else:
target_model_path = os.path.join(
target_model_folder_input,
f'{target_model_name_input}.{target_model_type}',
)
run_cmd += f' "{target_model_path}"'
print(run_cmd)
# Run the command
if os.name == 'posix':
os.system(run_cmd)
else:
subprocess.run(run_cmd)
if (
not target_model_type == 'diffuser'
or target_model_type == 'diffuser_safetensors'
):
v2_models = [
'stabilityai/stable-diffusion-2-1-base',
'stabilityai/stable-diffusion-2-base',
]
v_parameterization = [
'stabilityai/stable-diffusion-2-1',
'stabilityai/stable-diffusion-2',
]
if str(source_model_type) in v2_models:
inference_file = os.path.join(
target_model_folder_input, f'{target_model_name_input}.yaml'
)
print(f'Saving v2-inference.yaml as {inference_file}')
shutil.copy(
f'./v2_inference/v2-inference.yaml',
f'{inference_file}',
)
if str(source_model_type) in v_parameterization:
inference_file = os.path.join(
target_model_folder_input, f'{target_model_name_input}.yaml'
)
print(f'Saving v2-inference-v.yaml as {inference_file}')
shutil.copy(
f'./v2_inference/v2-inference-v.yaml',
f'{inference_file}',
)
# parser = argparse.ArgumentParser()
# parser.add_argument("--v1", action='store_true',
# help='load v1.x model (v1 or v2 is required to load checkpoint) / 1.xのモデルを読み込む')
# parser.add_argument("--v2", action='store_true',
# help='load v2.0 model (v1 or v2 is required to load checkpoint) / 2.0のモデルを読み込む')
# parser.add_argument("--fp16", action='store_true',
# help='load as fp16 (Diffusers only) and save as fp16 (checkpoint only) / fp16形式で読み込みDiffusers形式のみ対応、保存するcheckpointのみ対応')
# parser.add_argument("--bf16", action='store_true', help='save as bf16 (checkpoint only) / bf16形式で保存するcheckpointのみ対応')
# parser.add_argument("--float", action='store_true',
# help='save as float (checkpoint only) / float(float32)形式で保存するcheckpointのみ対応')
# parser.add_argument("--epoch", type=int, default=0, help='epoch to write to checkpoint / checkpointに記録するepoch数の値')
# parser.add_argument("--global_step", type=int, default=0,
# help='global_step to write to checkpoint / checkpointに記録するglobal_stepの値')
# parser.add_argument("--reference_model", type=str, default=None,
# help="reference model for schduler/tokenizer, required in saving Diffusers, copy schduler/tokenizer from this / scheduler/tokenizerのコピー元のDiffusersモデル、Diffusers形式で保存するときに必要")
# parser.add_argument("model_to_load", type=str, default=None,
# help="model to load: checkpoint file or Diffusers model's directory / 読み込むモデル、checkpointかDiffusers形式モデルのディレクトリ")
# parser.add_argument("model_to_save", type=str, default=None,
# help="model to save: checkpoint (with extension) or Diffusers model's directory (without extension) / 変換後のモデル、拡張子がある場合はcheckpoint、ない場合はDiffusesモデルとして保存")
###
# Gradio UI
###
def gradio_convert_model_tab():
with gr.Tab('Convert model'):
gr.Markdown(
'This utility can be used to convert from one stable diffusion model format to another.'
)
with gr.Row():
source_model_input = gr.Textbox(
label='Source model',
placeholder='path to source model folder of file to convert...',
interactive=True,
)
button_source_model_dir = gr.Button(
folder_symbol, elem_id='open_folder_small'
)
button_source_model_dir.click(
get_folder_path,
outputs=source_model_input,
show_progress=False,
)
button_source_model_file = gr.Button(
document_symbol, elem_id='open_folder_small'
)
button_source_model_file.click(
get_file_path,
inputs=[source_model_input],
outputs=source_model_input,
show_progress=False,
)
source_model_type = gr.Dropdown(
label='Source model type',
choices=[
'stabilityai/stable-diffusion-2-1-base',
'stabilityai/stable-diffusion-2-base',
'stabilityai/stable-diffusion-2-1',
'stabilityai/stable-diffusion-2',
'runwayml/stable-diffusion-v1-5',
'CompVis/stable-diffusion-v1-4',
],
)
with gr.Row():
target_model_folder_input = gr.Textbox(
label='Target model folder',
placeholder='path to target model folder of file name to create...',
interactive=True,
)
button_target_model_folder = gr.Button(
folder_symbol, elem_id='open_folder_small'
)
button_target_model_folder.click(
get_folder_path,
outputs=target_model_folder_input,
show_progress=False,
)
target_model_name_input = gr.Textbox(
label='Target model name',
placeholder='target model name...',
interactive=True,
)
target_model_type = gr.Dropdown(
label='Target model type',
choices=[
'diffuser',
'diffuser_safetensors',
'ckpt',
'safetensors',
],
)
target_save_precision_type = gr.Dropdown(
label='Target model precision',
choices=['unspecified', 'fp16', 'bf16', 'float'],
value='unspecified',
)
convert_button = gr.Button('Convert model')
convert_button.click(
convert_model,
inputs=[
source_model_input,
source_model_type,
target_model_folder_input,
target_model_name_input,
target_model_type,
target_save_precision_type,
],
show_progress=False,
)