File size: 3,495 Bytes
5039ac3
 
ee7e1de
 
5039ac3
 
 
 
 
 
 
 
 
ee7e1de
 
 
 
 
5039ac3
 
 
ee7e1de
 
 
 
 
5039ac3
 
 
2c89886
5039ac3
2c89886
ee7e1de
5039ac3
950e7a7
5039ac3
 
2c89886
 
5039ac3
 
2c89886
950e7a7
 
 
2c89886
950e7a7
 
 
 
 
 
 
 
 
2c89886
 
 
 
950e7a7
5039ac3
c12573a
950e7a7
5039ac3
950e7a7
 
5039ac3
950e7a7
 
5039ac3
950e7a7
 
5039ac3
950e7a7
 
5039ac3
950e7a7
 
5039ac3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
##~ Display Download Results Widgets | by: ANXETY ~##

from directory_setup import *

import re
import os
import json
import time
import ipywidgets as widgets
from ipywidgets import widgets, Layout, Label, Button, VBox, HBox
from IPython.display import display, HTML, Javascript, clear_output


# Setup Env
env = os.getenv('ENV_NAME')
root_path = os.getenv('ROOT_PATH')
webui_path = os.getenv('WEBUI_PATH')
free_plan = os.getenv('FREE_PLAN')


# ==================== CSS ====================
# Main CSS
css_file_path = f"{root_path}/CSS/dl_display_result.css"
with open(css_file_path , "r") as f:
    CSS = f.read()
display(HTML(f"<style>{CSS}</style>"))


# ==================== WIDGETS ====================
# Constants
HR = widgets.HTML('<hr>')
HEADER_DL = 'DOWNLOAD RESULTS'
VERSION = 'v0.33'

## Functions
def output_container_generator(header, items):
    header_widget = widgets.HTML(value=f'<div class="header_outputs_dl">{header} ➤</div>')
    content_widgets = [widgets.HTML(value=f'<div class="items_dl">{item}</div>') for item in items]
    container_widget = widgets.VBox([header_widget, *content_widgets]).add_class("outputs_dl")
    return container_widget

def get_files_list(directory, extensions):
    return [file for file in os.listdir(directory) if file.endswith(extensions)]

def get_folders_list(directory):
    return [folder for folder in os.listdir(extensions_dir) if os.path.isdir(os.path.join(extensions_dir, folder))]

def get_controlnets_list(directory, filter_pattern):
    filter_name = re.compile(filter_pattern)
    return [
        filter_name.match(file).group(1) if filter_name.match(file) else file
        for file in os.listdir(directory)
        if not file.endswith(('.txt', '.yaml'))
    ]

## Widgets
header_widget = widgets.HTML(value=f'''

<div><span class="header_dl">{HEADER_DL}</span> <span style="color: grey; opacity: 0.25;">| {VERSION}</span></div>

''')

# Models
models_list = get_files_list(models_dir, ('.safetensors', '.ckpt'))
models_widget = output_container_generator('Models', models_list)
# Vaes
vaes_list = get_files_list(vaes_dir, '.safetensors')
vaes_widget = output_container_generator('VAEs', vaes_list)
# Embeddings
embeddings_list = get_files_list(embeddings_dir, ('.safetensors', '.pt'))
embeddings_widget = output_container_generator('Embeddings', embeddings_list)
# LoRAs
loras_list = get_files_list(loras_dir, '.safetensors')
loras_widget = output_container_generator('LoRAs', loras_list)
# Extensions
extensions_list = get_folders_list(extensions_dir)
extensions_widget = output_container_generator('Extensions', extensions_list).add_class("extension") # for fix height
# ControlNet
controlnets_list = get_controlnets_list(control_dir, r'^[^_]*_[^_]*_[^_]*_(.*)_fp16\.safetensors')
controlnets_widget = output_container_generator('ControlNets', controlnets_list)

## Sorting and Output
widgets_dict = {
    models_widget: models_list,
    vaes_widget: vaes_list,
    embeddings_widget: embeddings_list,
    loras_widget: loras_list,
    extensions_widget: extensions_list,
    controlnets_widget: controlnets_list
}
outputs_widgets_list = [widget for widget, widget_list in widgets_dict.items() if widget_list]
result_output_widget = widgets.HBox(outputs_widgets_list).add_class("result_output_dl")

container_widget = widgets.VBox([header_widget, HR, result_output_widget, HR]).add_class("container_dl")
display(container_widget)