File size: 4,585 Bytes
11c2c17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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,
            )