vaiola/modules/civit/datamodel.py

from dataclasses import dataclass
from typing import Optional

from modules.shared.DataClassDatabase import DataClassDatabase


@dataclass
class ModelVersionFileHashes(DataClassDatabase):
    # primitives
    AutoV1: Optional[str] = None
    AutoV2: Optional[str] = None
    AutoV3: Optional[str] = None
    CRC32: Optional[str] = None
    SHA256: Optional[str] = None
    BLAKE3: Optional[str] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._table_name = 'model_versions_files_hashes'
        self._standalone_entity = False

@dataclass
class ModelVersionFileMetadata(DataClassDatabase):
    # primitives
    format: Optional[str] = None
    size: Optional[str] = None
    fp: Optional[str] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._table_name = 'model_versions_files_metadata'
        self._standalone_entity = False

@dataclass
class ModelVersionFile(DataClassDatabase):
    # primary key
    id: Optional[int] = None
    # primitives
    sizeKB: Optional[float] = None
    name: Optional[str] = None
    type: Optional[str] = None
    pickleScanResult: Optional[str] = None
    pickleScanMessage: Optional[str] = None
    virusScanResult: Optional[str] = None
    virusScanMessage: Optional[str] = None
    scannedAt: Optional[str] = None
    downloadUrl: Optional[str] = None
    primary: Optional[bool] = None
    # child entities
    metadata: Optional[ModelVersionFileMetadata] = None
    hashes: Optional[ModelVersionFileHashes] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {
            'metadata': ModelVersionFileMetadata,
            'hashes': ModelVersionFileHashes,
        }
        self._key_field = 'id'
        self._table_name = 'model_versions_files'
        self._standalone_entity = False

@dataclass
class ModelVersionStats(DataClassDatabase):
    downloadCount: Optional[int] = None
    ratingCount: Optional[int] = None
    rating: Optional[float] = None
    thumbsUpCount: Optional[int] = None
    thumbsDownCount: Optional[int] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._table_name = 'model_versions_stats'
        self._standalone_entity = False

@dataclass
class ModelVersionImage(DataClassDatabase):
    # primary key
    id: Optional[int] = None
    # primitives
    url: Optional[str] = None
    nsfwLevel: Optional[int] = None
    width: Optional[int] = None
    height: Optional[int] = None
    hash: Optional[str] = None
    type: Optional[str] = None
    minor: Optional[bool] = None
    poi: Optional[bool] = None
    hasMeta: Optional[bool] = None
    hasPositivePrompt: Optional[bool] = None
    onSite: Optional[int] = None
    remixOfId: Optional[int] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._key_field = 'id'
        self._table_name = 'model_versions_images'
        self._standalone_entity = False

@dataclass
class ModelVersion(DataClassDatabase):
    # primary key
    id: Optional[int] = None
    # primitives
    index: Optional[int] = None
    name: Optional[str] = None
    baseModel: Optional[str] = None
    baseModelType: Optional[str] = None
    publishedAt: Optional[str] = None
    availability: Optional[str] = None
    nsfwLevel: Optional[int] = None
    description: Optional[str] = None
    supportsGeneration: Optional[bool] = None
    downloadUrl: Optional[str] = None
    # list of primitives
    trainedWords: Optional[list] = None
    # child entities
    stats: Optional[ModelVersionStats] = None
    files: Optional[list[ModelVersionFile]] = None
    images: Optional[list[ModelVersionImage]] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {
            'stats': ModelVersionStats,
            'files': ModelVersionFile,
            'images': ModelVersionImage,
        }
        self._key_field = 'id'
        self._table_name = 'model_versions'
        self._standalone_entity = False

@dataclass
class ModelStats(DataClassDatabase):
    # primitives
    downloadCount: Optional[int] = None
    favoriteCount: Optional[int] = None
    thumbsUpCount: Optional[int] = None
    thumbsDownCount: Optional[int] = None
    commentCount: Optional[int] = None
    ratingCount: Optional[int] = None
    rating: Optional[int] = None
    tippedAmountCount: Optional[int] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._table_name = 'model_stats'
        self._standalone_entity = False

@dataclass
class ModelCreator(DataClassDatabase):
    # primitives
    username: Optional[str] = None
    image: Optional[str] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._table_name = 'model_creators'
        self._standalone_entity = False

@dataclass
class Model(DataClassDatabase):
    # primary key
    id: Optional[int] = None
    # primitives
    name: Optional[str] = None
    description: Optional[str] = None
    allowNoCredit: Optional[bool] = None
    allowCommercialUse: Optional[list] = None
    allowDerivatives: Optional[bool] = None
    allowDifferentLicense: Optional[bool] = None
    type: Optional[str] = None
    minor: Optional[bool] = None
    sfwOnly: Optional[bool] = None
    poi: Optional[bool] = None
    nsfw: Optional[bool] = None
    nsfwLevel: Optional[int] = None
    availability: Optional[str] = None
    cosmetic: Optional[str] = None
    supportsGeneration: Optional[bool] = None
    mode: Optional[str] = None
    # list of primitives
    tags: Optional[list] = None
    # child entities
    stats: Optional[ModelStats] = None
    creator: Optional[ModelCreator] = None
    modelVersions: Optional[list[ModelVersion]] = None


    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {
            'stats': ModelStats,
            'creator': ModelCreator,
            'modelVersions': ModelVersion,
        }
        self._key_field = 'id'
        self._table_name = 'models'
        self._standalone_entity = True

@dataclass
class Tag(DataClassDatabase):
    name: Optional[str] = None
    link: Optional[str] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._key_field = 'name'
        self._table_name = 'tags'
        self._standalone_entity = True

@dataclass
class Creator(DataClassDatabase):
    # primary key
    username: Optional[str] = None
    # primitives
    modelCount: Optional[int] = None
    link: Optional[str] = None
    image: Optional[str] = None

    def __post_init__(self):
        super().__post_init__()
        self._forwarding = {}
        self._key_field = 'username'
        self._table_name = 'creators'
        self._standalone_entity = True