| import os
|
| from collections import defaultdict
|
| from dataclasses import dataclass
|
|
|
| import gradio as gr
|
|
|
|
|
| MODEL_FILE_STATUS_MISSING = 0
|
| MODEL_FILE_STATUS_PARTIAL = 1
|
| MODEL_FILE_STATUS_EXPECTED = 2
|
| MODEL_STATUS_PREFIXES = {
|
| MODEL_FILE_STATUS_MISSING: "\u2B1B",
|
| MODEL_FILE_STATUS_EXPECTED: "\U0001F7E6",
|
| MODEL_FILE_STATUS_PARTIAL: "\U0001F7E8",
|
| }
|
|
|
|
|
| @dataclass
|
| class DropdownDeps:
|
| transformer_types: list
|
| displayed_model_types: list
|
| transformer_type: str
|
| three_levels_hierarchy: bool
|
| families_infos: dict
|
| server_config: dict
|
| transformer_quantization: str
|
| transformer_dtype_policy: str
|
| text_encoder_quantization: str
|
| get_model_def: callable
|
| get_model_recursive_prop: callable
|
| get_model_filename: callable
|
| get_local_model_filename: callable
|
| get_lora_dir: callable
|
| get_parent_model_type: callable
|
| get_base_model_type: callable
|
| get_model_family: callable
|
| get_model_name: callable
|
| get_transformer_dtype: callable
|
|
|
|
|
| def compact_name(family_name, model_name):
|
| if model_name.startswith(family_name):
|
| return model_name[len(family_name):].strip()
|
| return model_name
|
|
|
|
|
| def decorate_model_dropdown_label(label, status):
|
| if not isinstance(label, str):
|
| return label
|
| prefix = MODEL_STATUS_PREFIXES.get(status, "")
|
| return f"{prefix} {label}" if len(prefix) > 0 else label
|
|
|
|
|
| def decorate_dropdown_choices_with_status(choices, status_map):
|
| decorated = []
|
| for choice in choices:
|
| if not isinstance(choice, tuple) or len(choice) < 2:
|
| decorated.append(choice)
|
| continue
|
| label, value = choice[0], choice[1]
|
| status = status_map.get(value, MODEL_FILE_STATUS_MISSING)
|
| decorated.append((decorate_model_dropdown_label(label, status), value, *choice[2:]))
|
| return decorated
|
|
|
|
|
| def get_dropdown_model_types(deps):
|
| dropdown_types = list(deps.transformer_types) if len(deps.transformer_types) > 0 else list(deps.displayed_model_types)
|
| if deps.transformer_type not in dropdown_types:
|
| dropdown_types.append(deps.transformer_type)
|
| return list(dict.fromkeys(dropdown_types))
|
|
|
|
|
| def get_family_dropdown_model_types(deps, current_model_family, dropdown_types=None):
|
| dropdown_types = get_dropdown_model_types(deps) if dropdown_types is None else dropdown_types
|
| if current_model_family is None:
|
| return dropdown_types
|
| return [model_type for model_type in dropdown_types if deps.get_model_family(model_type, for_ui=True) == current_model_family]
|
|
|
|
|
| def _get_module_files_for_status(deps, model_type, quantization, dtype_policy):
|
| transformer_dtype = deps.get_transformer_dtype(model_type, dtype_policy)
|
| modules = deps.get_model_recursive_prop(model_type, "modules", return_list=True)
|
| modules = [deps.get_model_recursive_prop(module, "modules", sub_prop_name="_list", return_list=True) if isinstance(module, str) else module for module in modules]
|
| module_files = []
|
| for module_type in modules:
|
| if isinstance(module_type, dict):
|
| URLs1 = module_type.get("URLs", None)
|
| if URLs1 is None:
|
| return None
|
| module_files.append(deps.get_model_filename(model_type, quantization, transformer_dtype, URLs=URLs1))
|
| URLs2 = module_type.get("URLs2", None)
|
| if URLs2 is None:
|
| return None
|
| module_files.append(deps.get_model_filename(model_type, quantization, transformer_dtype, URLs=URLs2))
|
| else:
|
| module_files.append(deps.get_model_filename(model_type, quantization, transformer_dtype, module_type=module_type))
|
| return module_files
|
|
|
|
|
| def _get_status_quantization_and_dtype(deps):
|
| quantization = deps.server_config.get("transformer_quantization", deps.transformer_quantization)
|
| dtype_policy = deps.server_config.get("transformer_dtype_policy", deps.transformer_dtype_policy)
|
| return quantization, dtype_policy
|
|
|
|
|
| def _append_expected_file_entry(entries, seen, filename, extra_paths=None):
|
| if not isinstance(filename, str) or len(filename) == 0:
|
| return
|
| if extra_paths is None:
|
| extra_list = []
|
| elif isinstance(extra_paths, list):
|
| extra_list = [path for path in extra_paths if isinstance(path, str) and len(path) > 0]
|
| else:
|
| extra_list = [extra_paths] if isinstance(extra_paths, str) and len(extra_paths) > 0 else []
|
| key = (filename.casefold(), tuple(path.casefold() for path in extra_list))
|
| if key in seen:
|
| return
|
| seen.add(key)
|
| entries.append({"filename": filename, "extra_paths": extra_list if len(extra_list) > 0 else None})
|
|
|
|
|
| def _append_expected_local_path_entry(entries, seen, local_path):
|
| if not isinstance(local_path, str) or len(local_path) == 0:
|
| return
|
| path_key = local_path.casefold()
|
| if path_key in seen:
|
| return
|
| seen.add(path_key)
|
| entries.append({"path": local_path})
|
|
|
|
|
| def get_expected_core_file_entries_for_status(deps, model_type):
|
| model_def = deps.get_model_def(model_type)
|
| if model_def is None:
|
| return []
|
| quantization, dtype_policy = _get_status_quantization_and_dtype(deps)
|
| entries = []
|
| seen = set()
|
|
|
| expected_filename = deps.get_model_filename(model_type, quantization=quantization, dtype_policy=dtype_policy)
|
| _append_expected_file_entry(entries, seen, expected_filename)
|
| if isinstance(model_def, dict) and "URLs2" in model_def:
|
| expected_filename2 = deps.get_model_filename(model_type, quantization=quantization, dtype_policy=dtype_policy, submodel_no=2)
|
| _append_expected_file_entry(entries, seen, expected_filename2)
|
|
|
| module_files = _get_module_files_for_status(deps, model_type, quantization, dtype_policy)
|
| if isinstance(module_files, list):
|
| for filename in module_files:
|
| _append_expected_file_entry(entries, seen, filename)
|
|
|
| text_encoder_URLs = deps.get_model_recursive_prop(model_type, "text_encoder_URLs", return_list=True)
|
| if text_encoder_URLs is not None:
|
| text_encoder_filename = deps.get_model_filename(model_type=model_type, quantization=deps.text_encoder_quantization, dtype_policy=dtype_policy, URLs=text_encoder_URLs)
|
| text_encoder_folder = model_def.get("text_encoder_folder", None)
|
| _append_expected_file_entry(entries, seen, text_encoder_filename, extra_paths=text_encoder_folder)
|
| return entries
|
|
|
|
|
| def get_missing_core_file_entries_for_status(deps, model_type):
|
| missing_entries = []
|
| for entry in get_expected_core_file_entries_for_status(deps, model_type):
|
| filename = entry.get("filename", "")
|
| extra_paths = entry.get("extra_paths", None)
|
| if deps.get_local_model_filename(filename, extra_paths=extra_paths) is None:
|
| missing_entries.append(entry)
|
| return missing_entries
|
|
|
|
|
| def get_expected_secondary_file_entries_for_status(deps, model_type):
|
| model_def = deps.get_model_def(model_type)
|
| if model_def is None:
|
| return []
|
| entries = []
|
| seen = set()
|
|
|
| preload_urls = deps.get_model_recursive_prop(model_type, "preload_URLs", return_list=True)
|
| if preload_urls is None:
|
| preload_urls = []
|
| if not isinstance(preload_urls, list):
|
| preload_urls = [preload_urls]
|
| for url in preload_urls:
|
| if isinstance(url, str) and len(url) > 0:
|
| _append_expected_file_entry(entries, seen, url)
|
|
|
| vae_urls = model_def.get("VAE_URLs", [])
|
| if vae_urls is None:
|
| vae_urls = []
|
| if not isinstance(vae_urls, list):
|
| vae_urls = [vae_urls]
|
| for url in vae_urls:
|
| if isinstance(url, str) and len(url) > 0:
|
| _append_expected_file_entry(entries, seen, url)
|
|
|
| model_loras = deps.get_model_recursive_prop(model_type, "loras", return_list=True)
|
| if model_loras is None:
|
| model_loras = []
|
| if not isinstance(model_loras, list):
|
| model_loras = [model_loras]
|
| lora_dir = deps.get_lora_dir(model_type)
|
| for url in model_loras:
|
| if not isinstance(url, str) or len(url) == 0:
|
| continue
|
| basename = os.path.basename(url)
|
| if len(basename) == 0:
|
| continue
|
| _append_expected_local_path_entry(entries, seen, os.path.join(lora_dir, basename))
|
|
|
| return entries
|
|
|
|
|
| def has_secondary_model_files_for_status(deps, model_type, quantization, dtype_policy):
|
| model_def = deps.get_model_def(model_type)
|
| if model_def is None:
|
| return True
|
|
|
| text_encoder_URLs = deps.get_model_recursive_prop(model_type, "text_encoder_URLs", return_list=True)
|
| if text_encoder_URLs is not None:
|
| text_encoder_filename = deps.get_model_filename(model_type=model_type, quantization=deps.text_encoder_quantization, dtype_policy=dtype_policy, URLs=text_encoder_URLs)
|
| if isinstance(text_encoder_filename, str) and len(text_encoder_filename) > 0:
|
| text_encoder_folder = model_def.get("text_encoder_folder", None)
|
| if deps.get_local_model_filename(text_encoder_filename, extra_paths=text_encoder_folder) is None:
|
| return False
|
|
|
| for prop, recursive in (("preload_URLs", True), ("VAE_URLs", False)):
|
| if recursive:
|
| urls = deps.get_model_recursive_prop(model_type, prop, return_list=True)
|
| else:
|
| urls = model_def.get(prop, [])
|
| if urls is None:
|
| continue
|
| if not isinstance(urls, list):
|
| urls = [urls]
|
| for url in urls:
|
| if not isinstance(url, str) or len(url) == 0:
|
| continue
|
| if deps.get_local_model_filename(url) is None:
|
| return False
|
|
|
| model_loras = deps.get_model_recursive_prop(model_type, "loras", return_list=True)
|
| if model_loras is None:
|
| model_loras = []
|
| if not isinstance(model_loras, list):
|
| model_loras = [model_loras]
|
| lora_dir = deps.get_lora_dir(model_type)
|
| for url in model_loras:
|
| if not isinstance(url, str) or len(url) == 0:
|
| continue
|
| if not os.path.isfile(os.path.join(lora_dir, os.path.basename(url))):
|
| return False
|
|
|
| module_files = _get_module_files_for_status(deps, model_type, quantization, dtype_policy)
|
| if module_files is None:
|
| return False
|
| for filename in module_files:
|
| if not isinstance(filename, str) or len(filename) == 0:
|
| continue
|
| if deps.get_local_model_filename(filename) is None:
|
| return False
|
| return True
|
|
|
|
|
| def get_model_download_status(deps, model_type):
|
| quantization, dtype_policy = _get_status_quantization_and_dtype(deps)
|
| model_def = deps.get_model_def(model_type)
|
| expected_filenames = []
|
| expected_filename = deps.get_model_filename(model_type, quantization=quantization, dtype_policy=dtype_policy)
|
| if isinstance(expected_filename, str) and len(expected_filename) > 0:
|
| expected_filenames.append(expected_filename)
|
| if isinstance(model_def, dict) and "URLs2" in model_def:
|
| expected_filename2 = deps.get_model_filename(model_type, quantization=quantization, dtype_policy=dtype_policy, submodel_no=2)
|
| if isinstance(expected_filename2, str) and len(expected_filename2) > 0:
|
| expected_filenames.append(expected_filename2)
|
|
|
| expected_exists = []
|
| for filename in expected_filenames:
|
| expected_exists.append(deps.get_local_model_filename(filename) is not None)
|
|
|
| if len(expected_exists) > 0 and all(expected_exists):
|
| if not has_secondary_model_files_for_status(deps, model_type, quantization, dtype_policy):
|
| return MODEL_FILE_STATUS_PARTIAL
|
| return MODEL_FILE_STATUS_EXPECTED
|
|
|
| if any(expected_exists):
|
| return MODEL_FILE_STATUS_PARTIAL
|
|
|
| candidate_urls = []
|
| for prop in ("URLs", "URLs2"):
|
| urls = deps.get_model_recursive_prop(model_type, prop, return_list=True)
|
| if not isinstance(urls, list):
|
| urls = [urls] if urls else []
|
| candidate_urls += urls
|
|
|
| checked_candidates = set()
|
| expected_set = {name.casefold() for name in expected_filenames if isinstance(name, str) and len(name) > 0}
|
| for candidate in candidate_urls:
|
| if not isinstance(candidate, str) or len(candidate) == 0:
|
| continue
|
| candidate_key = candidate.casefold()
|
| if candidate_key in checked_candidates:
|
| continue
|
| checked_candidates.add(candidate_key)
|
| if candidate_key in expected_set:
|
| continue
|
| if deps.get_local_model_filename(candidate) is not None:
|
| return MODEL_FILE_STATUS_PARTIAL
|
| return MODEL_FILE_STATUS_MISSING
|
|
|
|
|
| def get_model_download_status_maps(deps, dropdown_types=None):
|
| direct_status_map = {}
|
| dropdown_types = get_dropdown_model_types(deps) if dropdown_types is None else dropdown_types
|
| parent_to_children = defaultdict(list)
|
|
|
| for model_type in dropdown_types:
|
| if deps.get_model_def(model_type) is None:
|
| continue
|
| status = get_model_download_status(deps, model_type)
|
| direct_status_map[model_type] = status
|
| parent_model_type = deps.get_parent_model_type(model_type)
|
| if parent_model_type is not None:
|
| parent_to_children[parent_model_type].append(model_type)
|
|
|
| aggregated_parent_status_map = dict(direct_status_map)
|
| for parent_model_type, children in parent_to_children.items():
|
| child_statuses = [direct_status_map.get(child, MODEL_FILE_STATUS_MISSING) for child in children]
|
| if len(child_statuses) == 0:
|
| continue
|
| parent_status = MODEL_FILE_STATUS_MISSING
|
| if any(status == MODEL_FILE_STATUS_EXPECTED for status in child_statuses):
|
| parent_status = MODEL_FILE_STATUS_EXPECTED
|
| elif any(status == MODEL_FILE_STATUS_PARTIAL for status in child_statuses):
|
| parent_status = MODEL_FILE_STATUS_PARTIAL
|
| aggregated_parent_status_map[parent_model_type] = max(aggregated_parent_status_map.get(parent_model_type, MODEL_FILE_STATUS_MISSING), parent_status)
|
| return direct_status_map, aggregated_parent_status_map
|
|
|
|
|
| def get_model_download_status_map(deps, dropdown_types=None):
|
| return get_model_download_status_maps(deps, dropdown_types)[1]
|
|
|
|
|
| def create_models_hierarchy(rows):
|
| """
|
| rows: list of (model_name, model_id, parent_model_id)
|
| returns:
|
| parents_list: list[(parent_header, parent_id)]
|
| children_dict: dict[parent_id] -> list[(child_display_name, child_id)]
|
| """
|
| toks = lambda s: [t for t in s.split() if t]
|
| norm = lambda s: " ".join(s.split()).casefold()
|
|
|
| groups, parents, order = defaultdict(list), {}, []
|
| for name, mid, pmid in rows:
|
| groups[pmid].append((name, mid))
|
| if mid == pmid and pmid not in parents:
|
| parents[pmid] = name
|
| order.append(pmid)
|
|
|
| parents_list, children_dict = [], {}
|
|
|
| for pid in order:
|
| p_name = parents[pid]
|
| p_tok = toks(p_name)
|
| p_low = [w.casefold() for w in p_tok]
|
| n = len(p_low)
|
| p_last = p_low[-1]
|
| p_set = set(p_low)
|
|
|
| kids = []
|
| for name, mid in groups.get(pid, []):
|
| ot = toks(name)
|
| lt = [w.casefold() for w in ot]
|
| st = set(lt)
|
| kids.append((name, mid, ot, lt, st))
|
|
|
| outliers = {mid for _, mid, _, _, st in kids if mid != pid and p_set.isdisjoint(st)}
|
|
|
| prefix_non = []
|
| for name, mid, ot, lt, st in kids:
|
| if mid == pid or (mid not in outliers and lt and lt[0] == p_low[0]):
|
| prefix_non.append((ot, lt))
|
|
|
| def lcp_len(a, b):
|
| i = 0
|
| m = min(len(a), len(b))
|
| while i < m and a[i] == b[i]:
|
| i += 1
|
| return i
|
|
|
| L = n if len(prefix_non) <= 1 else min(lcp_len(lt, p_low) for _, lt in prefix_non)
|
| if L == 0 and len(prefix_non) > 1:
|
| L = n
|
|
|
| shares_last = any(mid != pid and mid not in outliers and lt and lt[-1] == p_last for _, mid, _, lt, _ in kids)
|
| header_tokens_disp = p_tok[:L] + ([p_tok[-1]] if shares_last and L < n else [])
|
| header = " ".join(header_tokens_disp)
|
| header_has_last = (L == n) or (shares_last and L < n)
|
|
|
| prefix_low = p_low[:L]
|
|
|
| def startswith_prefix(lt):
|
| if L == 0 or len(lt) < L:
|
| return False
|
| for i in range(L):
|
| if lt[i] != prefix_low[i]:
|
| return False
|
| return True
|
|
|
| def base_rem(ot, lt):
|
| return ot[L:] if startswith_prefix(lt) else ot[:]
|
|
|
| def trim_rem(rem, lt):
|
| out = rem[:]
|
| if header_has_last and lt and lt[-1] == p_last and out and out[-1].casefold() == p_last:
|
| out = out[:-1]
|
| return out
|
|
|
| kid_infos = []
|
| for name, mid, ot, lt, _ in kids:
|
| rem_core = base_rem(ot, lt) if mid not in outliers else ot[:]
|
| kid_infos.append({
|
| "name": name,
|
| "mid": mid,
|
| "ot": ot,
|
| "lt": lt,
|
| "outlier": mid in outliers,
|
| "rem_core": rem_core,
|
| "rem_trim": trim_rem(rem_core, lt) if mid not in outliers else ot[:],
|
| "rem_set": {w.casefold() for w in rem_core} if mid not in outliers else set(),
|
| "rem_trim_set": {w.casefold() for w in (trim_rem(rem_core, lt) if mid not in outliers else ot[:])} if mid not in outliers else set(),
|
| })
|
|
|
| default_info = next(info for info in kid_infos if info["mid"] == pid)
|
| other_words = set()
|
| for info in kid_infos:
|
| if info["mid"] != pid:
|
| other_words |= info["rem_set"]
|
| default_shares = bool(default_info["rem_set"] & other_words)
|
|
|
| def disp(info):
|
| if info["outlier"]:
|
| return info["name"]
|
| if info["mid"] == pid:
|
| if not default_shares:
|
| return "Default"
|
| rem = info["rem_trim"]
|
| else:
|
| rem = info["rem_trim"]
|
| s = " ".join(rem).strip()
|
| return s if s else "Default"
|
|
|
| entries = [(disp(default_info), pid)]
|
| for info in kid_infos:
|
| if info["mid"] == pid:
|
| continue
|
| entries.append((disp(info), info["mid"]))
|
|
|
| p_full = norm(p_name)
|
| full_by_mid = {mid: name for name, mid, *_ in kids}
|
| num = 2
|
| numbered = [entries[0]]
|
| for dname, mid in entries[1:]:
|
| if dname == "Default" and norm(full_by_mid[mid]) == p_full:
|
| numbered.append((f"Default #{num}", mid))
|
| num += 1
|
| else:
|
| numbered.append((dname, mid))
|
|
|
| parents_list.append((header, pid))
|
| children_dict[pid] = numbered
|
|
|
| for pid in groups.keys():
|
| if pid in parents:
|
| continue
|
| first_name = groups[pid][0][0]
|
| parents_list.append((first_name, pid))
|
| children_dict[pid] = [(name, mid) for name, mid in groups[pid]]
|
|
|
| parents_list = sorted(parents_list, key=lambda c: c[0])
|
| return parents_list, children_dict
|
|
|
|
|
| def get_sorted_dropdown(deps, dropdown_types, current_model_family, current_model_type, three_levels=True):
|
| models_families = [deps.get_model_family(t, for_ui=True) for t in dropdown_types]
|
| families = {}
|
| for family in models_families:
|
| if family not in families:
|
| families[family] = 1
|
|
|
| families_orders = [deps.families_infos[family][0] for family in families]
|
| families_labels = [deps.families_infos[family][1] for family in families]
|
| sorted_familes = [info[1:] for info in sorted(zip(families_orders, families_labels, families), key=lambda c: c[0])]
|
| if current_model_family is None:
|
| dropdown_choices = [(deps.families_infos[family][0], deps.get_model_name(model_type), model_type) for model_type, family in zip(dropdown_types, models_families)]
|
| else:
|
| dropdown_choices = [(deps.families_infos[family][0], compact_name(deps.families_infos[family][1], deps.get_model_name(model_type)), model_type) for model_type, family in zip(dropdown_types, models_families) if family == current_model_family]
|
| dropdown_choices = sorted(dropdown_choices, key=lambda c: (c[0], c[1]))
|
| if three_levels:
|
| dropdown_choices = [(*model[1:], deps.get_parent_model_type(model[2])) for model in dropdown_choices]
|
| sorted_choices, finetunes_dict = create_models_hierarchy(dropdown_choices)
|
| return sorted_familes, sorted_choices, finetunes_dict[deps.get_parent_model_type(current_model_type)]
|
| dropdown_types_list = list({deps.get_base_model_type(model[2]) for model in dropdown_choices})
|
| dropdown_choices = [model[1:] for model in dropdown_choices]
|
| return sorted_familes, dropdown_types_list, dropdown_choices
|
|
|
|
|
| def generate_dropdown_model_list(deps, current_model_type):
|
| dropdown_types = list(deps.transformer_types) if len(deps.transformer_types) > 0 else list(deps.displayed_model_types)
|
| if current_model_type not in dropdown_types:
|
| dropdown_types.append(current_model_type)
|
| current_model_family = deps.get_model_family(current_model_type, for_ui=True)
|
| sorted_familes, sorted_models, sorted_finetunes = get_sorted_dropdown(deps, dropdown_types, current_model_family, current_model_type, three_levels=deps.three_levels_hierarchy)
|
| status_model_types = get_family_dropdown_model_types(deps, current_model_family, dropdown_types)
|
| if current_model_type not in status_model_types:
|
| status_model_types.append(current_model_type)
|
| direct_status_map, aggregated_parent_status_map = get_model_download_status_maps(deps, status_model_types)
|
| sorted_models = decorate_dropdown_choices_with_status(sorted_models, aggregated_parent_status_map)
|
| sorted_finetunes = decorate_dropdown_choices_with_status(sorted_finetunes, direct_status_map)
|
|
|
| dropdown_families = gr.Dropdown(choices=sorted_familes, value=current_model_family, show_label=False, scale=2 if deps.three_levels_hierarchy else 1, elem_id="family_list", min_width=50)
|
| dropdown_models = gr.Dropdown(choices=sorted_models, value=deps.get_parent_model_type(current_model_type) if deps.three_levels_hierarchy else deps.get_base_model_type(current_model_type), show_label=False, scale=3 if len(sorted_finetunes) > 1 else 7, elem_id="model_base_types_list", visible=deps.three_levels_hierarchy)
|
| dropdown_finetunes = gr.Dropdown(choices=sorted_finetunes, value=current_model_type, show_label=False, scale=4, visible=len(sorted_finetunes) > 1 or not deps.three_levels_hierarchy, elem_id="model_list")
|
| return dropdown_families, dropdown_models, dropdown_finetunes
|
|
|
|
|
| def change_model_family(deps, state, current_model_family):
|
| dropdown_types = list(deps.transformer_types) if len(deps.transformer_types) > 0 else list(deps.displayed_model_types)
|
| current_family_name = deps.families_infos[current_model_family][1]
|
| models_families = [deps.get_model_family(t, for_ui=True) for t in dropdown_types]
|
| dropdown_choices = [(compact_name(current_family_name, deps.get_model_name(model_type)), model_type) for model_type, family in zip(dropdown_types, models_families) if family == current_model_family]
|
| dropdown_choices = sorted(dropdown_choices, key=lambda c: c[0])
|
| family_dropdown_types = [choice[1] for choice in dropdown_choices]
|
| direct_status_map, aggregated_parent_status_map = get_model_download_status_maps(deps, family_dropdown_types)
|
| last_model_per_family = state.get("last_model_per_family", {})
|
| model_type = last_model_per_family.get(current_model_family, "")
|
| if len(model_type) == "" or model_type not in [choice[1] for choice in dropdown_choices]:
|
| model_type = dropdown_choices[0][1]
|
|
|
| if deps.three_levels_hierarchy:
|
| parent_model_type = deps.get_parent_model_type(model_type)
|
| dropdown_choices = [(*tup, deps.get_parent_model_type(tup[1])) for tup in dropdown_choices]
|
| dropdown_base_types_choices, finetunes_dict = create_models_hierarchy(dropdown_choices)
|
| dropdown_choices = decorate_dropdown_choices_with_status(finetunes_dict[parent_model_type], direct_status_map)
|
| dropdown_base_types_choices = decorate_dropdown_choices_with_status(dropdown_base_types_choices, aggregated_parent_status_map)
|
| model_finetunes_visible = len(dropdown_choices) > 1
|
| else:
|
| parent_model_type = deps.get_base_model_type(model_type)
|
| model_finetunes_visible = True
|
| dropdown_base_types_choices = list({deps.get_base_model_type(model[1]) for model in dropdown_choices})
|
| dropdown_choices = decorate_dropdown_choices_with_status(dropdown_choices, direct_status_map)
|
|
|
| return gr.Dropdown(choices=dropdown_base_types_choices, value=parent_model_type, scale=3 if model_finetunes_visible else 7), gr.Dropdown(choices=dropdown_choices, value=model_type, visible=model_finetunes_visible)
|
|
|
|
|
| def change_model_base_types(deps, state, current_model_family, model_base_type_choice):
|
| if not deps.three_levels_hierarchy:
|
| return gr.update()
|
| dropdown_types = list(deps.transformer_types) if len(deps.transformer_types) > 0 else list(deps.displayed_model_types)
|
| current_family_name = deps.families_infos[current_model_family][1]
|
| dropdown_choices = [(compact_name(current_family_name, deps.get_model_name(model_type)), model_type, model_base_type_choice) for model_type in dropdown_types if deps.get_parent_model_type(model_type) == model_base_type_choice and deps.get_model_family(model_type, for_ui=True) == current_model_family]
|
| dropdown_choices = sorted(dropdown_choices, key=lambda c: c[0])
|
| _, finetunes_dict = create_models_hierarchy(dropdown_choices)
|
| base_dropdown_types = [choice[1] for choice in dropdown_choices]
|
| direct_status_map, _ = get_model_download_status_maps(deps, base_dropdown_types)
|
| dropdown_choices = decorate_dropdown_choices_with_status(finetunes_dict[model_base_type_choice], direct_status_map)
|
| model_finetunes_visible = len(dropdown_choices) > 1
|
| last_model_per_type = state.get("last_model_per_type", {})
|
| model_type = last_model_per_type.get(model_base_type_choice, "")
|
| if len(model_type) == "" or model_type not in [choice[1] for choice in dropdown_choices]:
|
| model_type = dropdown_choices[0][1]
|
| return gr.update(scale=3 if model_finetunes_visible else 7), gr.Dropdown(choices=dropdown_choices, value=model_type, visible=model_finetunes_visible)
|
|
|
