File size: 3,641 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 | import os.path
from enum import Enum
class ModelType(Enum):
CHECKPOINT = 'Checkpoint'
VAE = 'VAE'
LORA = 'Lora'
HYPER_NETWORK = 'Hyper Network'
LYCORIS = 'LyCORIS'
EMBEDDING = 'Embedding'
OTHER = 'Other'
@staticmethod
def by_value(value: str):
for model_type in ModelType:
if model_type.value == value.strip():
return model_type
raise ValueError(f'There is no model type defined with name "{value}"')
class ModelSort(Enum):
TIME_ADDED_ASC = 'Time Added'
TIME_ADDED_DESC = 'Time Added Reversed'
NAME_ASC = 'Name'
NAME_DESC = 'Name Reversed'
@staticmethod
def by_value(value: str):
for model_sort in ModelSort:
if model_sort.value == value.strip():
return model_sort
raise ValueError(f'There is no model sort defined with name "{value}"')
class Record:
def __init__(self, id_,
name: str,
model_type: ModelType,
download_url: str = '',
url: str = '',
download_path: str = '',
download_filename: str = '',
preview_url: str = '',
description: str = '',
positive_prompts: str = '',
negative_prompts: str = '',
sha256_hash: str = '',
md5_hash: str = '',
location: str = '',
created_at: float = 0,
groups=None,
subdir: str = '',
weight: float = 1):
if groups is None:
groups = []
self.id_ = id_
self.name = name
self.model_type = model_type
self.url = url
self.download_url = download_url
self.download_path = download_path
self.download_filename = download_filename
self.preview_url = preview_url
self.description = description
self.positive_prompts = positive_prompts
self.negative_prompts = negative_prompts
self.sha256_hash = sha256_hash
self.md5_hash = md5_hash
self.location = location
self.created_at = created_at
self.groups = groups
self.subdir = subdir
self.weight = weight;
def is_file_exists(self) -> bool:
return bool(self.location) and os.path.isfile(self.location)
def is_downloadable(self) -> bool:
return bool(self.download_url)
def is_download_possible(self) -> bool:
return self.is_downloadable() and not self.is_file_exists()
def is_local_file_record(self):
return self.id_ is None and self.is_file_exists()
def __str__(self):
return f'id="{self.id_}", ' \
f'name="{self.name}", ' \
f'model_type="{self.model_type}", ' \
f'url="{self.url}", ' \
f'download_url="{self.download_url}", ' \
f'download_path="{self.download_path}", ' \
f'download_filename="{self.download_filename}", ' \
f'preview_url="{self.preview_url}", ' \
f'description="{self.description}", ' \
f'positive_prompts="{self.positive_prompts}", ' \
f'negative_prompts="{self.negative_prompts}", ' \
f'sha256_hash="{self.sha256_hash}", ' \
f'md5_hash="{self.md5_hash}", ' \
f'location="{self.location}", ' \
f'created_at="{self.created_at}", ' \
f'groups="{self.groups}", ' \
f'subdir="{self.subdir}",' \
f'weight="{self.weight}".'
|