ui_extra_networks_checkpoints.py 1.2 KB

12345678910111213141516171819202122232425262728293031
  1. import html
  2. import json
  3. import os
  4. from modules import shared, ui_extra_networks, sd_models
  5. class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
  6. def __init__(self):
  7. super().__init__('Checkpoints')
  8. def refresh(self):
  9. shared.refresh_checkpoints()
  10. def list_items(self):
  11. checkpoint: sd_models.CheckpointInfo
  12. for name, checkpoint in sd_models.checkpoints_list.items():
  13. path, ext = os.path.splitext(checkpoint.filename)
  14. yield {
  15. "name": checkpoint.name_for_extra,
  16. "filename": path,
  17. "preview": self.find_preview(path),
  18. "description": self.find_description(path),
  19. "search_term": self.search_terms_from_path(checkpoint.filename) + " " + (checkpoint.sha256 or ""),
  20. "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
  21. "local_preview": f"{path}.{shared.opts.samples_format}",
  22. }
  23. def allowed_directories_for_previews(self):
  24. return [v for v in [shared.cmd_opts.ckpt_dir, sd_models.model_path] if v is not None]