import gradio as gr from modules import ui_extra_networks_user_metadata, sd_vae from modules.ui_common import create_refresh_button class CheckpointUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataEditor): def __init__(self, ui, tabname, page): super().__init__(ui, tabname, page) self.select_vae = None def save_user_metadata(self, name, desc, notes, vae): user_metadata = self.get_user_metadata(name) user_metadata["description"] = desc user_metadata["notes"] = notes user_metadata["vae"] = vae self.write_user_metadata(name, user_metadata) def put_values_into_components(self, name): user_metadata = self.get_user_metadata(name) values = super().put_values_into_components(name) return [ *values[0:5], user_metadata.get('vae', ''), ] def create_editor(self): self.create_default_editor_elems() with gr.Row(): 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") 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") self.edit_notes = gr.TextArea(label='Notes', lines=4) self.create_default_buttons() viewed_components = [ self.edit_name, self.edit_description, self.html_filedata, self.html_preview, self.edit_notes, self.select_vae, ] self.button_edit\ .click(fn=self.put_values_into_components, inputs=[self.edit_name_input], outputs=viewed_components)\ .then(fn=lambda: gr.update(visible=True), inputs=[], outputs=[self.box]) edited_components = [ self.edit_description, self.edit_notes, self.select_vae, ] self.setup_save_handler(self.button_save, self.save_user_metadata, edited_components)