ui_extra_networks_checkpoints_user_metadata.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. import gradio as gr
  2. from modules import ui_extra_networks_user_metadata, sd_vae
  3. from modules.ui_common import create_refresh_button
  4. class CheckpointUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataEditor):
  5. def __init__(self, ui, tabname, page):
  6. super().__init__(ui, tabname, page)
  7. self.select_vae = None
  8. def save_user_metadata(self, name, desc, notes, vae):
  9. user_metadata = self.get_user_metadata(name)
  10. user_metadata["description"] = desc
  11. user_metadata["notes"] = notes
  12. user_metadata["vae"] = vae
  13. self.write_user_metadata(name, user_metadata)
  14. def put_values_into_components(self, name):
  15. user_metadata = self.get_user_metadata(name)
  16. values = super().put_values_into_components(name)
  17. return [
  18. *values[0:5],
  19. user_metadata.get('vae', ''),
  20. ]
  21. def create_editor(self):
  22. self.create_default_editor_elems()
  23. with gr.Row():
  24. 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")
  25. 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")
  26. self.edit_notes = gr.TextArea(label='Notes', lines=4)
  27. self.create_default_buttons()
  28. viewed_components = [
  29. self.edit_name,
  30. self.edit_description,
  31. self.html_filedata,
  32. self.html_preview,
  33. self.edit_notes,
  34. self.select_vae,
  35. ]
  36. self.button_edit\
  37. .click(fn=self.put_values_into_components, inputs=[self.edit_name_input], outputs=viewed_components)\
  38. .then(fn=lambda: gr.update(visible=True), inputs=[], outputs=[self.box])
  39. edited_components = [
  40. self.edit_description,
  41. self.edit_notes,
  42. self.select_vae,
  43. ]
  44. self.setup_save_handler(self.button_save, self.save_user_metadata, edited_components)