254 lines
9.4 KiB
Python
254 lines
9.4 KiB
Python
import os
|
||
import shutil
|
||
import uuid
|
||
from dataclasses import dataclass, fields
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from pythonapp.Libs.ConfigDataClass import Config
|
||
|
||
|
||
@dataclass
|
||
class PackageInfo(Config):
|
||
"""Информация о пакете"""
|
||
uuid: str = ""
|
||
name: str = ""
|
||
description: str = ""
|
||
release_date: str = ""
|
||
package_type: str = "" # unet, vae, text encoder
|
||
lineage: str = "" # sd 1.5, sdxl, flux.1
|
||
size_bytes: int = 0
|
||
version: str = ""
|
||
quantization: str = "" # fp8, bf16
|
||
dependencies: List[str] = None
|
||
resources: List[str] = None
|
||
|
||
def __post_init__(self):
|
||
if self.dependencies is None:
|
||
self.dependencies = []
|
||
if self.resources is None:
|
||
self.resources = []
|
||
super().__post_init__()
|
||
|
||
|
||
class ModelPackage:
|
||
def __init__(self, package_path: str, file_paths: List[str] = None, package_info: PackageInfo = None, load=False):
|
||
self.path = Path(package_path)
|
||
|
||
# Создаем директорию если она не существует
|
||
self.path.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Путь к файлу package.json
|
||
self.package_file = self.path / "package.json"
|
||
|
||
# Загружаем существующую информацию из файла
|
||
self.info = PackageInfo(filename=str(self.package_file))
|
||
|
||
self.files = []
|
||
|
||
if load:
|
||
self.files = self.get_files(package_path)
|
||
return
|
||
|
||
# Если package_info передан и не пустой, обновляем информацию
|
||
if package_info is not None:
|
||
# Обновляем только те поля, которые не определены
|
||
for field in fields(package_info):
|
||
field_name = field.name
|
||
field_value = getattr(package_info, field_name)
|
||
if field_value is not None and field_value != "" and field_name != "filename":
|
||
current_value = getattr(self.info, field_name)
|
||
if current_value is None or current_value == "" or current_value == 0 or len(current_value) == 0:
|
||
setattr(self.info, field_name, field_value)
|
||
|
||
# Генерируем UUID если он не определен
|
||
if not self.info.uuid:
|
||
self.info.uuid = str(uuid.uuid4())
|
||
|
||
# Сохраняем информацию о пакете
|
||
self.info.save()
|
||
|
||
# Создаем директорию files
|
||
self.files_path = self.path / "files"
|
||
self.files_path.mkdir(exist_ok=True)
|
||
|
||
|
||
if file_paths:
|
||
# Перемещаем файлы в директорию files
|
||
for file_path in file_paths:
|
||
src = Path(file_path)
|
||
if src.exists():
|
||
dst = self.files_path / src.name
|
||
if src.is_dir():
|
||
shutil.copytree(src, dst)
|
||
else:
|
||
shutil.copy2(src, dst)
|
||
|
||
def __str__(self): return self.name
|
||
def __repr__(self): return self.name
|
||
|
||
@classmethod
|
||
def interactive(cls, package_path: str, pkg_uuid = None):
|
||
"""Интерактивное создание пакета через консоль"""
|
||
print("Введите информацию о пакете:")
|
||
if os.path.exists(str(Path(package_path) / "package.json")): raise RuntimeError("package exists!")
|
||
package_info = PackageInfo(str(Path(package_path) / "package.json"))
|
||
package_info.name = input("Название: ").strip()
|
||
package_info.description = input("Описание: ").strip()
|
||
package_info.release_date = input("Дата выхода (YYYY-MM-DD): ").strip()
|
||
package_info.package_type = input("Тип пакета (unet, vae, text encoder): ").strip()
|
||
package_info.lineage = input("Линейка (sd 1.5, sdxl, flux.1): ").strip()
|
||
|
||
# Размер в байтах
|
||
while True:
|
||
try:
|
||
size = int(input("Размер в байтах: ").strip())
|
||
package_info.size_bytes = size
|
||
break
|
||
except ValueError:
|
||
print("Введите корректное число")
|
||
|
||
package_info.version = input("Версия: ").strip()
|
||
package_info.quantization = input("Квантование (fp8, bf16): ").strip()
|
||
|
||
# Ввод зависимостей
|
||
print("Зависимости (введите по одной, пустая строка для завершения):")
|
||
dependencies = []
|
||
while True:
|
||
dep = input().strip()
|
||
if not dep:
|
||
break
|
||
dependencies.append(dep)
|
||
package_info.dependencies = dependencies
|
||
|
||
# Ввод ресурсов
|
||
print("Ресурсы (введите по одному, пустая строка для завершения):")
|
||
resources = []
|
||
while True:
|
||
resource = input().strip()
|
||
if not resource:
|
||
break
|
||
resources.append(resource)
|
||
package_info.resources = resources
|
||
|
||
# Генерируем UUID случайным образом (не запрашиваем у пользователя)
|
||
package_info.uuid = pkg_uuid
|
||
if not package_info.uuid:
|
||
package_info.uuid = str(uuid.uuid4())
|
||
print(f"Сгенерирован UUID: {package_info.uuid}")
|
||
|
||
# Ввод путей к файлам
|
||
print("Пути к файлам и директориям (введите по одному, пустая строка для завершения):")
|
||
file_paths = []
|
||
while True:
|
||
file_path = input().strip()
|
||
if not file_path:
|
||
break
|
||
file_paths.append(file_path)
|
||
|
||
# Создаем пакет
|
||
return cls(package_path, file_paths, package_info)
|
||
|
||
@property
|
||
def uuid(self) -> str:
|
||
"""Возвращает UUID пакета"""
|
||
return self.info.uuid
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.name
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.description
|
||
|
||
@property
|
||
def release_date(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.release_date
|
||
|
||
@property
|
||
def package_type(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.package_type
|
||
|
||
@property
|
||
def lineage(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.lineage
|
||
|
||
@property
|
||
def size_bytes(self) -> int:
|
||
"""Возвращает название пакета"""
|
||
return self.info.size_bytes
|
||
|
||
@property
|
||
def version(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.version
|
||
|
||
@property
|
||
def quantization(self) -> str:
|
||
"""Возвращает название пакета"""
|
||
return self.info.quantization
|
||
|
||
|
||
@property
|
||
def dependencies(self) -> List[str]:
|
||
"""Возвращает список зависимостей пакета"""
|
||
return self.info.dependencies.copy() # Возвращаем копию для защиты от изменений
|
||
|
||
@property
|
||
def provides(self) -> List[str]:
|
||
"""Возвращает список ресурсов, предоставляемых пакетом (включая имя пакета)"""
|
||
provides_list = self.info.resources.copy() # Возвращаем копию
|
||
if self.info.name: # Добавляем имя пакета, если оно есть
|
||
provides_list.append(self.info.name)
|
||
return provides_list
|
||
|
||
@classmethod
|
||
def load(cls, package_path: str): return cls(package_path, load=True)
|
||
|
||
|
||
@staticmethod
|
||
def get_files(package_path):
|
||
"""
|
||
Получает путь к директории, ищет в ней поддиректорию 'files',
|
||
и возвращает список относительных путей ко всем файлам внутри 'files'.
|
||
|
||
:param package_path: Путь к базовой директории
|
||
:return: Список относительных путей внутри директории 'files'
|
||
"""
|
||
files_dir = os.path.join(package_path, 'files')
|
||
|
||
# Проверяем, существует ли директория 'files'
|
||
if not os.path.exists(files_dir):
|
||
return []
|
||
|
||
if not os.path.isdir(files_dir):
|
||
return []
|
||
|
||
relative_paths = []
|
||
|
||
# Рекурсивно обходим директорию 'files'
|
||
for root, dirs, files in os.walk(files_dir):
|
||
# Вычисляем относительный путь от 'files' к текущей директории
|
||
rel_root = os.path.relpath(root, files_dir)
|
||
|
||
# Добавляем файлы
|
||
for file in files:
|
||
if rel_root == '.':
|
||
relative_paths.append(file)
|
||
else:
|
||
relative_paths.append(os.path.join(rel_root, file))
|
||
|
||
return relative_paths
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
p = ModelPackage('/tmp/pkg')
|
||
pass |