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