kohya_ss / library /class_source_model.py
Ateras's picture
Upload folder using huggingface_hub
fe6327d
raw
history blame
4.59 kB
import gradio as gr
from .common_gui import (
get_any_file_path,
get_folder_path,
set_pretrained_model_name_or_path_input,
)
folder_symbol = '\U0001f4c2' # πŸ“‚
refresh_symbol = '\U0001f504' # πŸ”„
save_style_symbol = '\U0001f4be' # πŸ’Ύ
document_symbol = '\U0001F4C4' # πŸ“„
class SourceModel:
def __init__(
self,
save_model_as_choices=[
'same as source model',
'ckpt',
'diffusers',
'diffusers_safetensors',
'safetensors',
],
headless=False,
):
self.headless = headless
self.save_model_as_choices = save_model_as_choices
with gr.Tab('Source model'):
# Define the input elements
with gr.Row():
self.model_list = gr.Dropdown(
label='Model Quick Pick',
choices=[
'custom',
# 'stabilityai/stable-diffusion-xl-base-0.9',
# 'stabilityai/stable-diffusion-xl-refiner-0.9',
'stabilityai/stable-diffusion-2-1-base/blob/main/v2-1_512-ema-pruned',
'stabilityai/stable-diffusion-2-1-base',
'stabilityai/stable-diffusion-2-base',
'stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-ema-pruned',
'stabilityai/stable-diffusion-2-1',
'stabilityai/stable-diffusion-2',
'runwayml/stable-diffusion-v1-5',
'CompVis/stable-diffusion-v1-4',
],
value='runwayml/stable-diffusion-v1-5',
)
self.save_model_as = gr.Dropdown(
label='Save trained model as',
choices=save_model_as_choices,
value='safetensors',
)
with gr.Row():
self.pretrained_model_name_or_path = gr.Textbox(
label='Pretrained model name or path',
placeholder='enter the path to custom model or name of pretrained model',
value='runwayml/stable-diffusion-v1-5',
visible=(False and not headless),
)
self.pretrained_model_name_or_path_file = gr.Button(
document_symbol,
elem_id='open_folder_small',
visible=(False and not headless),
)
self.pretrained_model_name_or_path_file.click(
get_any_file_path,
inputs=self.pretrained_model_name_or_path,
outputs=self.pretrained_model_name_or_path,
show_progress=False,
)
self.pretrained_model_name_or_path_folder = gr.Button(
folder_symbol,
elem_id='open_folder_small',
visible=(False and not headless),
)
self.pretrained_model_name_or_path_folder.click(
get_folder_path,
inputs=self.pretrained_model_name_or_path,
outputs=self.pretrained_model_name_or_path,
show_progress=False,
)
self.v2 = gr.Checkbox(label='v2', value=False, visible=False)
self.v_parameterization = gr.Checkbox(
label='v_parameterization', value=False, visible=False
)
self.sdxl_checkbox = gr.Checkbox(
label='SDXL Model', value=False, visible=False
)
self.model_list.change(
set_pretrained_model_name_or_path_input,
inputs=[
self.model_list,
self.pretrained_model_name_or_path,
self.pretrained_model_name_or_path_file,
self.pretrained_model_name_or_path_folder,
self.v2,
self.v_parameterization,
self.sdxl_checkbox,
],
outputs=[
self.model_list,
self.pretrained_model_name_or_path,
self.pretrained_model_name_or_path_file,
self.pretrained_model_name_or_path_folder,
self.v2,
self.v_parameterization,
self.sdxl_checkbox,
],
show_progress=False,
)