65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field, asdict
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import json
|
|
|
|
|
|
class Gender(str, Enum):
|
|
MALE = "male"
|
|
FEMALE = "female"
|
|
OTHER = "other"
|
|
|
|
|
|
@dataclass
|
|
class Appearance:
|
|
hair: str = ""
|
|
eyes: str = ""
|
|
skin: str = ""
|
|
build: str = ""
|
|
features: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class Outfit:
|
|
name: str = "default"
|
|
description: str = ""
|
|
tags: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class Character:
|
|
id: str
|
|
name: str
|
|
gender: Gender
|
|
age: str = ""
|
|
appearance: Appearance = field(default_factory=Appearance)
|
|
outfits: list[Outfit] = field(default_factory=list)
|
|
style_tags: list[str] = field(default_factory=list)
|
|
lora_path: Optional[str] = None
|
|
reference_dir: Optional[str] = None
|
|
|
|
def to_dict(self) -> dict:
|
|
return asdict(self)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict) -> Character:
|
|
data = data.copy()
|
|
data["gender"] = Gender(data["gender"])
|
|
if isinstance(data.get("appearance"), dict):
|
|
data["appearance"] = Appearance(**data["appearance"])
|
|
if isinstance(data.get("outfits"), list):
|
|
data["outfits"] = [Outfit(**o) if isinstance(o, dict) else o for o in data["outfits"]]
|
|
return cls(**data)
|
|
|
|
def get_reference_images(self, registry_root: Path) -> list[Path]:
|
|
if not self.reference_dir:
|
|
ref_dir = registry_root / "references" / self.id
|
|
else:
|
|
ref_dir = Path(self.reference_dir)
|
|
if not ref_dir.exists():
|
|
return []
|
|
return sorted(p for p in ref_dir.iterdir() if p.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"})
|