stable-diffusion-webui

1
import gradio as gr
2

3
from modules import ui_extra_networks_user_metadata, sd_vae, shared
4
from modules.ui_common import create_refresh_button
5

6

7
class CheckpointUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataEditor):
8
    def __init__(self, ui, tabname, page):
9
        super().__init__(ui, tabname, page)
10

11
        self.select_vae = None
12

13
    def save_user_metadata(self, name, desc, notes, vae):
14
        user_metadata = self.get_user_metadata(name)
15
        user_metadata["description"] = desc
16
        user_metadata["notes"] = notes
17
        user_metadata["vae"] = vae
18

19
        self.write_user_metadata(name, user_metadata)
20

21
    def update_vae(self, name):
22
        if name == shared.sd_model.sd_checkpoint_info.name_for_extra:
23
            sd_vae.reload_vae_weights()
24

25
    def put_values_into_components(self, name):
26
        user_metadata = self.get_user_metadata(name)
27
        values = super().put_values_into_components(name)
28

29
        return [
30
            *values[0:5],
31
            user_metadata.get('vae', ''),
32
        ]
33

34
    def create_editor(self):
35
        self.create_default_editor_elems()
36

37
        with gr.Row():
38
            self.select_vae = gr.Dropdown(choices=["Automatic", "None"] + list(sd_vae.vae_dict), value="None", label="Preferred VAE", elem_id="checpoint_edit_user_metadata_preferred_vae")
39
            create_refresh_button(self.select_vae, sd_vae.refresh_vae_list, lambda: {"choices": ["Automatic", "None"] + list(sd_vae.vae_dict)}, "checpoint_edit_user_metadata_refresh_preferred_vae")
40

41
        self.edit_notes = gr.TextArea(label='Notes', lines=4)
42

43
        self.create_default_buttons()
44

45
        viewed_components = [
46
            self.edit_name,
47
            self.edit_description,
48
            self.html_filedata,
49
            self.html_preview,
50
            self.edit_notes,
51
            self.select_vae,
52
        ]
53

54
        self.button_edit\
55
            .click(fn=self.put_values_into_components, inputs=[self.edit_name_input], outputs=viewed_components)\
56
            .then(fn=lambda: gr.update(visible=True), inputs=[], outputs=[self.box])
57

58
        edited_components = [
59
            self.edit_description,
60
            self.edit_notes,
61
            self.select_vae,
62
        ]
63

64
        self.setup_save_handler(self.button_save, self.save_user_metadata, edited_components)
65
        self.button_save.click(fn=self.update_vae, inputs=[self.edit_name_input])
66

67