v18.1: Model conversion utility

.gitignore

@@ -3,3 +3,5 @@ venv1
 mytraining.ps
 __pycache__
 .vscode
+*.egg-info
+build

README.md

@@ -130,6 +130,8 @@ Drop by the discord server for support: https://discord.com/channels/10415185624
 
 ## Change history
 
+* 12/18 (v18.1) update:
+    - Add Stable Diffusion model conversion utility. Make sure to run `pip upgrade -U -r requirements.txt` after updating to this release as this introduce new pip requirements.
 * 12/17 (v18) update:
     - Save model as option added to train_db_fixed.py
     - Save model as option added to GUI

dreambooth_gui.py

@@ -10,14 +10,15 @@ import os
 import subprocess
 import pathlib
 import shutil
-from dreambooth_gui.dreambooth_folder_creation import gradio_dreambooth_folder_creation_tab
-from dreambooth_gui.caption_gui import gradio_caption_gui_tab
-from dreambooth_gui.dataset_balancing import gradio_dataset_balancing_tab
-from dreambooth_gui.common_gui import (
+from library.dreambooth_folder_creation_gui import gradio_dreambooth_folder_creation_tab
+from library.caption_gui import gradio_caption_gui_tab
+from library.dataset_balancing_gui import gradio_dataset_balancing_tab
+from library.common_gui import (
     get_folder_path,
     remove_doublequote,
     get_file_path,
 )
+from library.convert_model_gui import gradio_convert_model_tab
 from easygui import filesavebox, msgbox
 
 folder_symbol = '\U0001f4c2'  # 📂
@@ -699,6 +700,7 @@ with interface:
         # Captionning tab
         gradio_caption_gui_tab()
         gradio_dataset_balancing_tab()
+        gradio_convert_model_tab()
         # with gr.Tab('Model conversion'):
         #     convert_to_safetensors_input = gr.Checkbox(
         #         label='Convert to SafeTensors', value=True

library/convert_model_gui.py

@@ -0,0 +1,191 @@
+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' # 📄
+
+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'.\\venv\Scripts\python.exe "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":
+        run_cmd += f' --reference_model="{source_model_type}"'
+        
+    run_cmd += f' "{source_model_input}"'
+    
+    if target_model_type == "diffuser":
+        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
+    subprocess.run(run_cmd)
+    
+    if not target_model_type == "diffuser":
+        
+        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
+            )
+            
+            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
+            )
+            
+            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
+            )
+            
+            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',
+                    'ckpt',
+                    'safetensors',
+                ],)
+            target_save_precision_type = gr.Dropdown(label="Target model precison", 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
+            ],
+        )

library/dataset_balancing_gui.py

@@ -102,7 +102,6 @@ def gradio_dataset_balancing_tab():
 
             total_repeats_number = gr.Number(
                 value=1000,
-                min=1,
                 interactive=True,
                 label='Training steps per concept per epoch',
             )

requirements.txt

@@ -12,3 +12,4 @@ safetensors==0.2.6
 gradio
 altair
 easygui
+.

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup, find_packages
+ 
+setup(name = "library", packages = find_packages())

tools/convert_diffusers20_original_sd.py

@@ -5,9 +5,7 @@ import argparse
 import os
 import torch
 from diffusers import StableDiffusionPipeline
-
-import model_util
-
+from library import model_util as model_util
 
 def convert(args):
   # 引数を確認する

train_db_fixed.py

@@ -43,7 +43,7 @@ import cv2
 from einops import rearrange
 from torch import einsum
 
-import model_util
+import library.model_util as model_util
 
 # Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
 TOKENIZER_PATH = "openai/clip-vit-large-patch14"