File size: 5,252 Bytes
bb7f1f4 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import os
import gradio as gr
import scripts.mo.ui_styled_html as styled
from scripts.mo.environment import env, logger
from scripts.mo.ui_navigation import generate_ui_token
from scripts.mo.utils import find_preview_file, find_info_file
def _on_id_change(record_id):
logger.info('_on_id_change record_id: %s', record_id)
if not record_id:
return [
gr.HTML.update(value='Record is missing.'),
gr.Button.update(visible=False),
gr.Button.update(visible=False),
gr.Button.update(visible=False),
gr.Button.update(visible=False)
]
if os.path.isfile(record_id):
return [
gr.HTML.update(
value=styled.alert_danger(f'Are you sure you what to remove "{os.path.basename(record_id)}"?')),
gr.Button.update(visible=True),
gr.Button.update(visible=False),
gr.Button.update(visible=True),
gr.Button.update(visible=False)
]
else:
record = env.storage.get_record_by_id(record_id)
if record is None:
return [
gr.HTML.update(value='Record was not found in database.'),
gr.Button.update(visible=False),
gr.Button.update(visible=False),
gr.Button.update(visible=False),
gr.Button.update(visible=False)
]
elif os.path.exists(record.location):
return [
gr.HTML.update(value=styled.alert_danger(f'Are you sure you what to remove "{record.name}"?')),
gr.Button.update(visible=True),
gr.Button.update(visible=True),
gr.Button.update(visible=True),
gr.Button.update(visible=True)
]
else:
return [
gr.HTML.update(value=styled.alert_danger(f'Are you sure you what to remove "{record.name}"?')),
gr.Button.update(visible=True),
gr.Button.update(visible=True),
gr.Button.update(visible=False),
gr.Button.update(visible=False)
]
def _on_remove_record_button_click(record_id):
env.storage.remove_record(record_id)
logger.info('removed record: %s', record_id)
return generate_ui_token()
def _on_remove_files_button_click(record_id):
if os.path.isfile(record_id):
logger.info('removed local model file: %s', record_id)
os.remove(record_id)
preview_path = find_preview_file(record_id)
if preview_path and os.path.exists(preview_path):
logger.info('removed local preview files: %s', preview_path)
os.remove(preview_path)
else:
record = env.storage.get_record_by_id(record_id)
if record.location and os.path.exists(record.location):
logger.info('removed model file: %s', record.location)
os.remove(record.location)
preview_path = find_preview_file(record.location)
if preview_path and os.path.exists(preview_path):
logger.info('removed preview file: %s', preview_path)
os.remove(preview_path)
info_file = find_info_file(record.location)
if info_file is not None and os.path.isfile(info_file):
logger.info('removed info file: %s', info_file)
os.remove(info_file)
return generate_ui_token()
def _on_remove_both_button(record_id):
_on_remove_files_button_click(record_id)
_on_remove_record_button_click(record_id)
return generate_ui_token()
def remove_ui_block():
with gr.Blocks():
remove_id_box = gr.Textbox(label='remove_id_box', elem_classes='mo-alert-warning', visible=False)
remove_back_box = gr.Textbox(label='remove_back_box', elem_classes='mo-alert-warning', visible=False)
gr.Markdown('## Record removal')
html_widget = gr.HTML()
with gr.Row():
gr.Markdown()
cancel_button = gr.Button('Cancel')
remove_record_button = gr.Button('Remove Record', visible=False, elem_id='mo_button_remove')
remove_files_button = gr.Button('Remove Files', visible=False, elem_id='mo_button_remove')
remove_both_button = gr.Button('Remove Record and Files', visible=False, elem_id='mo_button_remove')
gr.Markdown()
remove_record_button.click(_on_remove_record_button_click,
inputs=remove_id_box,
outputs=remove_back_box)
remove_files_button.click(_on_remove_files_button_click,
inputs=remove_id_box,
outputs=remove_back_box)
remove_both_button.click(_on_remove_both_button,
inputs=remove_id_box,
outputs=remove_back_box)
cancel_button.click(fn=None, _js='navigateBack')
remove_id_box.change(_on_id_change, inputs=remove_id_box,
outputs=[html_widget, cancel_button, remove_record_button, remove_files_button,
remove_both_button])
remove_back_box.change(fn=None, _js='navigateHome')
return remove_id_box
|