import os import shutil import uuid from dataclasses import dataclass, fields from pathlib import Path from typing import List from modules.shared.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 tags: List[str] = None def __post_init__(self): if self.dependencies is None: self.dependencies = [] if self.resources is None: self.resources = [] if self.tags is None: self.tags = [] 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(str(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 print("Теги (введите по одному, пустая строка для завершения):") tags = [] while True: tag = input().strip() if not tag: break tags.append(resource) package_info.tags = tags # Генерируем 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) provides_list.extend(self.info.tags) 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