import os
import inspect
from html import escape
from typing import cast
import gradio as gr
from modules import errors, sd_models, sd_vae, extras, sd_samplers, ui_symbols, modelstats
from modules.ui_components import ToolButton
from modules.ui_common import create_refresh_button
from modules.call_queue import wrap_gradio_gpu_call
from modules.shared import opts, log
extra_ui = []
def update_model_hashes():
from modules import sd_unet, sd_checkpoint
unets = {}
for k, v in sd_unet.unet_dict.items():
unets[k] = sd_checkpoint.CheckpointInfo(name=k, filename=v, model_type='unet')
yield from sd_models.update_model_hashes(unets, model_type='unet')
yield from sd_models.update_model_hashes(model_type='checkpoint')
def create_models_table(rows: list | None = None):
from modules import sd_detect
if rows is None:
rows = []
rows = sorted(rows, key=lambda row: str(getattr(row, 'model_name', '')).lower())
html = """
<table class="simple-table sortable-table" data-sortable="true" data-default-sort-key="name" data-default-sort-order="asc" data-sort-key="name" data-sort-order="asc">
<thead>
<tr>
<th class="sortable sorted-asc" aria-sort="ascending" data-sort-key="name" data-sort-type="text">Name</th>
<th class="sortable" data-sort-key="family" data-sort-type="text">Family</th>
<th class="sortable" data-sort-key="type" data-sort-type="text">Type</th>
<th class="sortable" data-sort-key="pipeline" data-sort-type="text">Pipeline</th>
<th class="sortable" data-sort-key="size" data-sort-type="number">Size</th>
<th class="sortable" data-sort-key="mtime" data-sort-type="number">MTime</th>
<th class="sortable" data-sort-key="hash" data-sort-type="text">Hash</th>
<th></th>
</tr>
</thead>
<tbody>
{tbody}
</tbody>
</table>
"""
tbody = ''
for row in rows:
try:
f = row.filename
size, mtime = modelstats.stat(f)
size = round(size / 1024 / 1024 / 1024, 3)
if os.path.isfile(f):
typ = os.path.splitext(f)[1][1:]
elif os.path.isdir(f):
typ = 'diffusers'
else:
typ = 'unknown'
guess = 'Stable Diffusion'
guess = sd_detect.guess_by_size(f, guess)
guess = sd_detect.guess_by_name(f, guess)
guess, pipeline = sd_detect.guess_by_diffusers(f, guess)
guess = sd_detect.guess_variant(f, guess)
pipeline = sd_detect.shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline
model_name = escape(str(row.model_name))
pipeline_name = escape(pipeline.__name__ if pipeline else '(unknown)')
typ_name = escape(str(typ))
guess_name = escape(str(guess))
hash_name = escape(str(row.shorthash))
mtime_sort = mtime.timestamp() if hasattr(mtime, 'timestamp') else 0
tbody += f"""
<tr>
<td data-sort-value="{model_name.lower()}">{model_name}</td>
<td data-sort-value="{typ_name.lower()}">{typ_name}</td>
<td data-sort-value="{guess_name.lower()}">{guess_name}</td>
<td data-sort-value="{pipeline_name.lower()}">{pipeline_name}</td>
<td data-sort-value="{size}">{size:.3f} GB</td>
<td data-sort-value="{mtime_sort}">{mtime}</td>
<td data-sort-value="{hash_name.lower()}">{hash_name}</td>
<td style="cursor:pointer;" onclick="deleteFile('{escape(str(row.path))}')">\uf530</td>
</tr>
"""
except Exception as e:
log.error(f'Model list: row={vars(row)} {e}')
return html.format(tbody=tbody)
def create_ui():
log.debug('UI initialize: tab=models')
dummy_component = gr.Label(visible=False)
with gr.Row(elem_id="models_tab"):
with gr.Column(elem_id='models_output_container', scale=1):
models_outcome = gr.HTML(elem_id="models_outcome", value="")
models_file = gr.File(label='', visible=False)
with gr.Column(elem_id='models_input_container', scale=3):
with gr.Tab(label="Active Model", elem_id="models_current_tab"):
def create_modules_table(rows: list):
html = """
<table class="simple-table">
<thead>
<tr><th>Module</th><th>Class</th><th>Device</th><th>Dtype</th><th>Quant</th><th>Params</th><th>Modules</th><th>Config</th></tr>
</thead>
<tbody>
{tbody}
</tbody>
</table>
"""
tbody = ''
for row in rows:
try:
config = str(row.config)
except Exception:
config = '{}'
try:
tbody += f"""
<tr>
<td>{row.name}</td>
<td>{row.cls}</td>
<td>{row.device}</td>
<td>{row.dtype}</td>
<td>{row.quant}</td>
<td>{row.params}</td>
<td>{row.modules}</td>
<td><div class='model-config'>{config}</div></td>
</tr>
"""
except Exception as e:
log.error(f'Model list: row={vars(row)} {e}')
return html.format(tbody=tbody)
def analyze():
model = modelstats.analyze()
if model is None:
return ["Model not loaded", {}]
meta = model.meta
html = create_modules_table(model.modules)
return [html, meta]
with gr.Row():
model_analyze = gr.Button(value="Analyze model", variant='primary')
with gr.Row():
model_desc = gr.HTML(value="", elem_id="model_desc")
with gr.Accordion(label="Save model", open=False):
with gr.Row():
save_name = gr.Textbox(label="Model name", placeholder="Model name to save as")
with gr.Row():
save_path = gr.Textbox(label="Model base path", placeholder="Path to save model to", value=opts.diffusers_dir)
with gr.Row():
save_shard = gr.Textbox(label="Max shard size", placeholder="Maximum shard size", value="5GB")
save_overwrite = gr.Checkbox(label="Overwrite existing", value=False)
with gr.Row():
save_result = gr.HTML(value="", elem_id="model_save_outcome")
with gr.Row():
model_save = gr.Button(value="Save model", variant='primary')
model_save.click(fn=sd_models.save_model, inputs=[save_name, save_path, save_shard, save_overwrite], outputs=[save_result])
with gr.Accordion(label="Metadata", open=False):
model_meta = gr.JSON(label="Metadata", value={}, elem_id="model_meta")
model_analyze.click(fn=analyze, inputs=[], outputs=[model_desc, model_meta])
with gr.Tab(label="Models List", elem_id="models_list_tab"):
with gr.Row():
gr.HTML('<h2>List all locally available models</h2><br>')
with gr.Row():
model_list_btn = gr.Button(value="Refresh list", variant='primary')
model_checkhash_btn = gr.Button(value="Calculate hashes", variant='secondary')
with gr.Row():
model_table = gr.HTML(value=create_models_table(), elem_id="model_list_table", elem_classes="scroll-auto")
model_checkhash_btn.click(fn=update_model_hashes, inputs=[], outputs=[model_table])
model_list_btn.click(fn=lambda: create_models_table(list(sd_models.checkpoints_list.values())), inputs=[], outputs=[model_table])
with gr.Tab(label="Cache List", elem_id="models_cache_tab"):
with gr.Row():
gr.HTML('<h2>List models in Hugging Face cache</h2><br>')
with gr.Row():
model_cache_btn = gr.Button(value="Refresh list", variant='primary')
with gr.Row():
cache_table = gr.HTML(value=create_models_table(), elem_id="model_cache_table", elem_classes="scroll-auto")
model_cache_btn.click(fn=lambda: create_models_table(sd_models.list_hfcache()), inputs=[], outputs=[cache_table])
with gr.Tab(label="Metadata", elem_id="models_metadata_tab"):
from modules.civitai.metadata_civitai import civit_search_metadata, civit_update_metadata
with gr.Row():
gr.HTML('<h2>Fetch model preview metadata</h2><br>')
with gr.Row():
civit_previews_btn = gr.Button(value="Scan missing", variant='primary')
civit_update_btn = gr.Button(value="Update all", variant='primary')
with gr.Row():
civit_metadata = gr.HTML(value='', elem_id="civit_metadata")
civit_previews_btn.click(fn=civit_search_metadata, inputs=[], outputs=[civit_metadata])
civit_update_btn.click(fn=civit_update_metadata, inputs=[], outputs=[civit_metadata])
with gr.Tab(label="Loader", elem_id="models_loader_tab"):
from modules import ui_models_load
ui_models_load.create_ui(models_outcome, models_file)
with gr.Tab(label="Merge", elem_id="models_merge_tab"):
from modules.merging import merge_methods
from modules.merging.merge_utils import BETA_METHODS, TRIPLE_METHODS, interpolate
from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
def sd_model_choices():
return ['None'] + sd_models.checkpoint_titles()
with gr.Row():
gr.HTML('<h3> Merge multiple models<br></h2>')
with gr.Row(equal_height=False):
with gr.Column(variant='compact'):
with gr.Row():
custom_name = gr.Textbox(label="New model name")
with gr.Row():
merge_mode = gr.Dropdown(choices=merge_methods.__all__, value="weighted_sum", label="Interpolation Method")
merge_mode_docs = gr.HTML(value=merge_methods.weighted_sum.__doc__.strip().replace("\n", "<br>"))
with gr.Row():
primary_model_name = gr.Dropdown(sd_model_choices(), label="Primary model", value="None")
create_refresh_button(primary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "checkpoint_A_refresh")
secondary_model_name = gr.Dropdown(sd_model_choices(), label="Secondary model", value="None")
create_refresh_button(secondary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "checkpoint_B_refresh")
tertiary_model_name = gr.Dropdown(sd_model_choices(), label="Tertiary model", value="None", visible=False)
tertiary_refresh = create_refresh_button(tertiary_model_name, sd_models.list_models, lambda: {"choices": sd_model_choices()}, "checkpoint_C_refresh", visible=False)
with gr.Row():
with gr.Tabs() as tabs:
with gr.TabItem(label="Simple Merge", id=0):
with gr.Row():
alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Alpha Ratio', value=0.5)
beta = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Beta Ratio', value=None, visible=False)
with gr.TabItem(label="Preset Block Merge", id=1):
with gr.Row():
sdxl = gr.Checkbox(label="SDXL")
with gr.Row():
alpha_preset = gr.Dropdown(
choices=["None"] + list(BLOCK_WEIGHTS_PRESETS.keys()), value=None,
label="ALPHA Block Weight Preset", multiselect=True, max_choices=2)
alpha_preset_lambda = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Preset Interpolation Ratio', value=None, visible=False)
apply_preset = ToolButton('⇨', visible=True)
with gr.Row():
beta_preset = gr.Dropdown(choices=["None"] + list(BLOCK_WEIGHTS_PRESETS.keys()), value=None, label="BETA Block Weight Preset", multiselect=True, max_choices=2, interactive=True, visible=False)
beta_preset_lambda = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Preset Interpolation Ratio', value=None, interactive=True, visible=False)
beta_apply_preset = ToolButton('⇨', interactive=True, visible=False)
with gr.TabItem(label="Manual Block Merge", id=2):
with gr.Row():
alpha_label = gr.Markdown("# Alpha")
with gr.Row():
alpha_base = gr.Textbox(value=None, label="Base", min_width=70, scale=1)
alpha_in_blocks = gr.Textbox(value=None, label="In Blocks", scale=15)
alpha_mid_block = gr.Textbox(value=None, label="Mid Block", min_width=80, scale=1)
alpha_out_blocks = gr.Textbox(value=None, label="Out Block", scale=15)
with gr.Row():
beta_label = gr.Markdown("# Beta", visible=False)
with gr.Row():
beta_base = gr.Textbox(value=None, label="Base", min_width=70, scale=1, interactive=True, visible=False)
beta_in_blocks = gr.Textbox(value=None, label="In Blocks", interactive=True, scale=15, visible=False)
beta_mid_block = gr.Textbox(value=None, label="Mid Block", min_width=80, interactive=True, scale=1, visible=False)
beta_out_blocks = gr.Textbox(value=None, label="Out Block", interactive=True, scale=15, visible=False)
with gr.Row():
overwrite = gr.Checkbox(label="Overwrite model")
with gr.Row():
save_metadata = gr.Checkbox(value=True, label="Save metadata")
with gr.Row():
weights_clip = gr.Checkbox(label="Weights clip")
prune = gr.Checkbox(label="Prune", value=True, visible=False)
with gr.Row():
re_basin = gr.Checkbox(label="ReBasin")
re_basin_iterations = gr.Slider(minimum=0, maximum=25, step=1, label='Number of ReBasin Iterations', value=None, visible=False)
with gr.Row():
checkpoint_format = gr.Radio(choices=["ckpt", "safetensors"], value="safetensors", visible=False, label="Model format")
with gr.Row():
precision = gr.Radio(choices=["fp16", "fp32"], value="fp16", label="Model precision")
with gr.Row():
device = gr.Radio(choices=["cpu", "shuffle", "gpu"], value="cpu", label="Merge Device")
unload = gr.Checkbox(label="Unload Current Model from VRAM", value=False, visible=False)
with gr.Row():
bake_in_vae = gr.Dropdown(choices=["None"] + list(sd_vae.vae_dict), value="None", interactive=True, label="Replace VAE")
create_refresh_button(bake_in_vae, sd_vae.refresh_vae_list,
lambda: {"choices": ["None"] + list(sd_vae.vae_dict)},
"modelmerger_bake_in_vae_refresh")
with gr.Row():
modelmerger_merge = gr.Button(value="Merge", variant='primary')
def modelmerger(dummy_component,
overwrite,
primary_model_name,
secondary_model_name,
tertiary_model_name,
merge_mode,
alpha,
beta,
alpha_preset,
alpha_preset_lambda,
alpha_base,
alpha_in_blocks,
alpha_mid_block,
alpha_out_blocks,
beta_preset,
beta_preset_lambda,
beta_base,
beta_in_blocks,
beta_mid_block,
beta_out_blocks,
precision,
custom_name,
checkpoint_format,
save_metadata,
weights_clip,
prune,
re_basin,
re_basin_iterations,
device,
unload,
bake_in_vae):
kwargs = {}
for x in inspect.getfullargspec(modelmerger)[0]:
kwargs[x] = locals()[x]
for key in list(kwargs.keys()):
if kwargs[key] in [None, "None", "", 0, []]:
del kwargs[key]
del kwargs['dummy_component']
if kwargs.get("custom_name", None) is None:
log.error('Merge: no output model specified')
return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], "No output model specified"]
elif kwargs.get("primary_model_name", None) is None or kwargs.get("secondary_model_name", None) is None:
log.error('Merge: no models selected')
return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], "No models selected"]
else:
log.debug(f'Merge start: {kwargs}')
try:
results = extras.run_modelmerger(dummy_component, **kwargs)
except Exception as e:
errors.display(e, 'Merge')
sd_models.list_models()
return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Error merging checkpoints: {e}"]
return results
def tertiary(mode):
if mode in TRIPLE_METHODS:
return [gr.update(visible=True) for _ in range(2)]
else:
return [gr.update(visible=False) for _ in range(2)]
def beta_visibility(mode):
if mode in BETA_METHODS:
return [gr.update(visible=True) for _ in range(9)]
else:
return [gr.update(visible=False) for _ in range(9)]
def show_iters(show):
if show:
return gr.Slider.update(value=5, visible=True)
else:
return gr.Slider.update(value=None, visible=False)
def show_help(mode):
try:
doc = getattr(merge_methods, mode).__doc__.strip().replace("\n", "<br>")
except AttributeError:
log.warning(f'Merge mode "{mode}" is missing documentation')
doc = "Error: Documentation missing"
return gr.update(value=doc, visible=True)
def show_unload(device):
if device == "gpu":
return gr.update(visible=True)
else:
return gr.update(visible=False)
def preset_visiblility(x):
if len(x) == 2:
return gr.Slider.update(value=0.5, visible=True)
else:
return gr.Slider.update(value=None, visible=False)
def load_presets(presets, ratio):
for i, p in enumerate(presets):
presets[i] = BLOCK_WEIGHTS_PRESETS[p]
if len(presets) == 2:
preset = interpolate(presets, ratio)
else:
preset = presets[0]
preset = [f'{x:.3f}' if int(x) != x else str(x) for x in preset]
preset = [preset[0], ",".join(preset[1:13]), preset[13], ",".join(preset[14:])]
return [gr.update(value=x) for x in preset] + [gr.update(selected=2)]
def preset_choices(sdxl):
if sdxl:
return [gr.update(choices=["None"] + list(SDXL_BLOCK_WEIGHTS_PRESETS.keys())) for _ in range(2)]
else:
return [gr.update(choices=["None"] + list(BLOCK_WEIGHTS_PRESETS.keys())) for _ in range(2)]
device.change(fn=show_unload, inputs=device, outputs=unload)
merge_mode.change(fn=show_help, inputs=merge_mode, outputs=merge_mode_docs)
sdxl.change(fn=preset_choices, inputs=sdxl, outputs=[alpha_preset, beta_preset])
alpha_preset.change(fn=preset_visiblility, inputs=alpha_preset, outputs=alpha_preset_lambda)
beta_preset.change(fn=preset_visiblility, inputs=alpha_preset, outputs=beta_preset_lambda)
merge_mode.input(fn=tertiary, inputs=merge_mode, outputs=[tertiary_model_name, tertiary_refresh])
merge_mode.input(fn=beta_visibility, inputs=merge_mode, outputs=[beta, alpha_label, beta_label, beta_apply_preset, beta_preset, beta_base, beta_in_blocks, beta_mid_block, beta_out_blocks])
re_basin.change(fn=show_iters, inputs=re_basin, outputs=re_basin_iterations)
apply_preset.click(fn=load_presets, inputs=[alpha_preset, alpha_preset_lambda], outputs=[alpha_base, alpha_in_blocks, alpha_mid_block, alpha_out_blocks, cast("gr.components.Component", tabs)])
beta_apply_preset.click(fn=load_presets, inputs=[beta_preset, beta_preset_lambda], outputs=[beta_base, beta_in_blocks, beta_mid_block, beta_out_blocks, cast("gr.components.Component", tabs)])
modelmerger_merge.click(
fn=wrap_gradio_gpu_call(modelmerger, extra_outputs=lambda: [gr.update() for _ in range(4)], name='Models'),
_js='modelmerger',
inputs=[
dummy_component,
overwrite,
primary_model_name,
secondary_model_name,
tertiary_model_name,
merge_mode,
alpha,
beta,
alpha_preset,
alpha_preset_lambda,
alpha_base,
alpha_in_blocks,
alpha_mid_block,
alpha_out_blocks,
beta_preset,
beta_preset_lambda,
beta_base,
beta_in_blocks,
beta_mid_block,
beta_out_blocks,
precision,
custom_name,
checkpoint_format,
save_metadata,
weights_clip,
prune,
re_basin,
re_basin_iterations,
device,
unload,
bake_in_vae,
],
outputs=[
primary_model_name,
secondary_model_name,
tertiary_model_name,
dummy_component,
models_outcome,
]
)
with gr.Tab(label="Replace", elem_id="models_replace_tab"):
with gr.Row():
gr.HTML('<h3> Replace model components<br></h2>')
with gr.Row():
with gr.Column(scale=3):
model_type = gr.Dropdown(label="Base model type", choices=['sd15', 'sdxl', 'sd21', 'sd35', 'flux.1'], value='sdxl', interactive=False)
with gr.Column(scale=5):
with gr.Row():
model_name = gr.Dropdown(sd_models.checkpoint_titles(), label="Input model")
create_refresh_button(model_name, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_titles()}, "checkpoint_Z_refresh")
with gr.Column(scale=5):
custom_name = gr.Textbox(label="Output model", placeholder="Output model path")
with gr.Row():
with gr.Column(scale=3):
gr.HTML('Model components<br><span style="color: var(--body-text-color-subdued)">Specify the components to include<br>Paths can be relative or absolute</span><br>')
with gr.Column(scale=5):
comp_unet = gr.Textbox(placeholder="UNet model", show_label=False)
comp_vae = gr.Textbox(placeholder="VAE model", show_label=False)
with gr.Column(scale=5):
comp_te1 = gr.Textbox(placeholder="Text encoder 1", show_label=False)
comp_te2 = gr.Textbox(placeholder="Text encoder 2", show_label=False)
with gr.Row():
with gr.Column(scale=3):
gr.HTML('Model settings<br>')
with gr.Column(scale=10):
with gr.Row():
precision = gr.Dropdown(label="Model precision", choices=["fp32", "fp16", "bf16"], value="fp16")
comp_scheduler = gr.Dropdown(label="Sampler", choices=[s.name for s in sd_samplers.samplers if s.constructor is not None])
comp_prediction = gr.Dropdown(label="Prediction type", choices=["epsilon", "v"], value="epsilon")
with gr.Row():
with gr.Column(scale=3):
gr.HTML('Merge LoRA<br>')
with gr.Column(scale=9):
comp_lora = gr.Textbox(label="Comma separated list with optional strength per LoRA", placeholder="LoRA models")
with gr.Column(scale=1):
comp_fuse = gr.Number(label="Fuse strength", value=1.0)
with gr.Row():
gr.HTML('<br>')
with gr.Row():
with gr.Column(scale=2):
gr.HTML('Model metadata<br>')
with gr.Column(scale=5):
meta_author = gr.Textbox(placeholder="Author name", show_label=False)
meta_version = gr.Textbox(placeholder="Model version", show_label=False)
meta_license = gr.Textbox(placeholder="Model license", show_label=False)
with gr.Column(scale=5):
meta_desc = gr.Textbox(placeholder="Model description", lines=3, show_label=False)
meta_hint = gr.Textbox(placeholder="Model hint", lines=3, show_label=False)
with gr.Column(scale=3):
meta_thumbnail = gr.Image(label="Thumbnail", type='pil')
with gr.Row():
gr.HTML('Note: Save is optional as you can merge in-memory and use newly created model immediately')
with gr.Row():
create_diffusers = gr.Checkbox(label="Save diffusers", value=True)
create_safetensors = gr.Checkbox(label="Save safetensors", value=True)
debug = gr.Checkbox(label="Debug info", value=False)
model_modules_btn = gr.Button(value="Merge Modules", variant='primary')
model_modules_btn.click(
fn=extras.run_model_modules,
inputs=[
model_type, model_name, custom_name,
comp_unet, comp_vae, comp_te1, comp_te2,
precision, comp_scheduler, comp_prediction,
comp_lora, comp_fuse,
meta_author, meta_version, meta_license, meta_desc, meta_hint, meta_thumbnail,
create_diffusers, create_safetensors, debug,
],
outputs=[models_outcome]
)
with gr.Tab(label="CivitAI", elem_id="models_civitai_tab") as civitai_tab:
from modules.civitai.search_civitai import search_civitai, create_model_cards, base_models
sort_fallback = ['', 'Most Downloaded', 'Highest Rated', 'Most Liked', 'Most Discussed',
'Most Collected', 'Most Images', 'Newest', 'Oldest']
type_fallback = ['', 'Checkpoint', 'TextualInversion', 'Hypernetwork', 'AestheticGradient',
'LORA', 'LoCon', 'DoRA', 'Controlnet', 'Upscaler', 'MotionModule',
'VAE', 'Poses', 'Wildcards', 'Workflows', 'Detection', 'Other']
def civitai_search(civit_search_text, civit_search_tag, civit_nsfw, civit_type,
civit_base, civit_token, civit_sort, civit_period):
if civit_search_text and civit_search_tag:
civit_search_tag = ''
results = search_civitai(query=civit_search_text, tag=civit_search_tag, nsfw=civit_nsfw,
types=civit_type, base=civit_base, token=civit_token,
sort=civit_sort, period=civit_period)
html = create_model_cards(results)
return html
def civitai_update_token(token):
log.debug('CivitAI update token')
opts.civitai_token = token
opts.save()
def civitai_download(model_urls, model_names, model_types, model_bases,
model_ids, version_ids, model_path, civit_token, model_output):
from modules.civitai.download_civitai import download_civit_model
for model_url, model_name, model_type, model_base, model_id, version_id in zip(
model_urls, model_names, model_types, model_bases, model_ids, version_ids, strict=False):
msg = f"<h4>Initiating download</h4><div>{model_name} | {model_type} | <a href='{model_url}'>{model_url}</a></div><br>"
yield msg + model_output
download_civit_model(model_url, model_name, model_path, model_type, civit_token,
base_model=model_base, model_id=int(model_id or 0), version_id=int(version_id or 0))
yield model_output
def civitai_toggle_subfolder(enabled, template):
opts.data['civitai_save_subfolder_enabled'] = enabled
if enabled and not template:
template = '{{BASEMODEL}}'
opts.data['civitai_save_subfolder'] = template
opts.save()
return gr.update(value=template, interactive=enabled)
with gr.Row():
gr.HTML('<h2>Search & Download</h2>')
with gr.Row(elem_id='civitai_search_row'):
civit_search_text = gr.Textbox(label='', placeholder='keyword, model id, or civitai url', elem_id="civit_search_text")
civit_search_tag = gr.Textbox(label='', placeholder='tag', elem_id="civit_search_tag")
civit_search_text_btn = ToolButton(value=ui_symbols.search, interactive=True, elem_id="civit_text_search")
with gr.Accordion(label='Options', open=False, elem_id="civitai_search_options"):
civit_download_btn = gr.Button(value="Download model", variant='primary', elem_id="civitai_download_btn", visible=False)
with gr.Row():
civit_type = gr.Dropdown(choices=type_fallback, label='CivitAI model type', value='', elem_id='civit_type')
civit_base = gr.Dropdown(choices=base_models, label='CivitAI base model', value='')
with gr.Row():
civit_sort = gr.Dropdown(choices=sort_fallback, label='CivitAI sort', value='', elem_id='civit_sort')
civit_period = gr.Dropdown(
choices=['', 'AllTime', 'Year', 'Month', 'Week', 'Day'],
label='CivitAI time period', value='', elem_id='civit_period',
)
with gr.Row():
civit_nsfw = gr.Checkbox(label='CivitAI NSFW allowed', value=True)
with gr.Row():
civit_token = gr.Textbox(opts.civitai_token, label='CivitAI token', placeholder='optional access token for private or gated models', elem_id="civitai_token")
with gr.Row():
civit_folder = gr.Textbox(label='Download folder', placeholder='optional folder for downloads')
with gr.Row():
civit_subfolder_enabled = gr.Checkbox(
label='Sort downloads into subfolders',
value=getattr(opts, 'civitai_save_subfolder_enabled', False),
elem_id='civit_subfolder_enabled',
)
civit_subfolder_template = gr.Textbox(
label='Subfolder template',
value=getattr(opts, 'civitai_save_subfolder', '') if getattr(opts, 'civitai_save_subfolder_enabled', False) else '',
placeholder='e.g. {{BASEMODEL}} or {{CREATOR}}/{{BASEMODEL}}',
interactive=getattr(opts, 'civitai_save_subfolder_enabled', False),
elem_id='civit_subfolder_template',
)
with gr.Row():
civitai_models_output = gr.HTML('', elem_id="civitai_models_output")
_dummy = gr.Label(visible=False)
civit_inputs = [civit_search_text, civit_search_tag, civit_nsfw, civit_type,
civit_base, civit_token, civit_sort, civit_period]
civit_search_text_btn.click(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output])
civit_search_text.submit(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output])
civit_search_tag.submit(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output])
civit_token.change(fn=civitai_update_token, inputs=[civit_token], outputs=[])
civit_subfolder_enabled.change(
fn=civitai_toggle_subfolder,
inputs=[civit_subfolder_enabled, civit_subfolder_template],
outputs=[civit_subfolder_template],
)
civit_subfolder_template.change(
fn=lambda v: (setattr(opts, 'civitai_save_subfolder', v), opts.save()),
inputs=[civit_subfolder_template], outputs=[],
)
civit_download_btn.click(
fn=civitai_download,
_js="downloadCivitModel",
inputs=[_dummy, _dummy, _dummy, _dummy, _dummy, _dummy, civit_folder, civit_token, civitai_models_output],
outputs=[civitai_models_output],
show_progress='full',
)
_civitai_loaded = False
def civitai_on_tab_enter():
nonlocal _civitai_loaded
if _civitai_loaded:
return [gr.update(), gr.update(), gr.update(), gr.update(), gr.update()]
_civitai_loaded = True
from modules.civitai.client_civitai import client
options = client.discover_options()
type_choices = [''] + (options.get('types', []) or type_fallback[1:])
sort_choices = [''] + (options.get('sort', []) or sort_fallback[1:])
base_choices = [''] + (options.get('base_models', []) or base_models[1:])
results = search_civitai(query='', sort='Most Downloaded', period='AllTime', limit=20)
html = create_model_cards(results)
return [
gr.update(choices=type_choices),
gr.update(choices=sort_choices, value='Most Downloaded'),
gr.update(choices=base_choices),
gr.update(value='AllTime'),
html,
]
civitai_tab.select(
fn=civitai_on_tab_enter,
inputs=[],
outputs=[civit_type, civit_sort, civit_base, civit_period, civitai_models_output],
)
with gr.Tab(label="Huggingface", elem_id="models_huggingface_tab"):
from modules.models_hf import hf_search, hf_select, hf_download_model, hf_update_token
with gr.Column(scale=6):
with gr.Row():
gr.HTML('<h3> Download model from huggingface<br></h2>')
with gr.Row():
hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models')
hf_search_btn = ToolButton(value=ui_symbols.search, interactive=True, elem_id="hf_text_search")
with gr.Row():
hf_selected = gr.Textbox('', label='Select model', placeholder='select model from search results or enter model name manually')
with gr.Accordion(label='Advanced', open=False, elem_id="hf_search_options"):
with gr.Row():
hf_token = gr.Textbox(opts.huggingface_token, label='Huggingface token', placeholder='optional access token for private or gated models', elem_id="hf_token")
with gr.Row():
hf_variant = gr.Textbox('', label='Specify model variant', placeholder='')
hf_revision = gr.Textbox('', label='Specify model revision', placeholder='')
with gr.Row():
hf_mirror = gr.Textbox('', label='Huggingface mirror', placeholder='optional mirror site for downloads')
hf_custom_pipeline = gr.Textbox('', label='Custom pipeline', placeholder='optional pipeline for downloads')
with gr.Column(scale=1):
gr.HTML('<br>')
hf_download_model_btn = gr.Button(value="Download model", variant='primary')
with gr.Row():
hf_headers = ['Name', 'Pipeline', 'Tags', 'Downloads', 'Updated', 'URL']
hf_types = ['str', 'str', 'str', 'number', 'date', 'markdown']
hf_results = gr.DataFrame(None, label='Search results', show_label=True, interactive=False, wrap=True, headers=hf_headers, datatype=hf_types)
hf_search_text.submit(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results])
hf_search_btn.click(fn=hf_search, inputs=[hf_search_text], outputs=[hf_results])
hf_results.select(fn=hf_select, inputs=[hf_results], outputs=[hf_selected])
hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_custom_pipeline], outputs=[models_outcome])
hf_token.change(fn=hf_update_token, inputs=[hf_token], outputs=[])
from modules.lora.lora_extract import create_ui as lora_extract_ui
lora_extract_ui()
for ui in extra_ui:
if callable(ui):
ui()