vaiola/modelspace/ModelPackage.py

import os.path
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):
        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))

        # Если 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)

    @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 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

if __name__ == "__main__":
    p = ModelPackage('/tmp/pkg')
    pass