107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
import importlib
|
|
import pathlib
|
|
import typing as t
|
|
|
|
import mlx.core as mx
|
|
import PIL.Image
|
|
import toml
|
|
|
|
from mflux.config.model_config import ModelConfig
|
|
|
|
|
|
class GeneratedImage:
|
|
def __init__(
|
|
self,
|
|
image: PIL.Image.Image,
|
|
model_config: ModelConfig,
|
|
seed: int,
|
|
prompt: str,
|
|
steps: int,
|
|
guidance: float | None,
|
|
precision: mx.Dtype,
|
|
quantization: int,
|
|
generation_time: float,
|
|
lora_paths: list[str],
|
|
lora_scales: list[float],
|
|
controlnet_image_path: str | pathlib.Path | None = None,
|
|
controlnet_strength: float | None = None,
|
|
init_image_path: str | pathlib.Path | None = None,
|
|
init_image_strength: float | None = None,
|
|
):
|
|
self.image = image
|
|
self.model_config = model_config
|
|
self.seed = seed
|
|
self.prompt = prompt
|
|
self.steps = steps
|
|
self.guidance = guidance
|
|
self.precision = precision
|
|
self.quantization = quantization
|
|
self.generation_time = generation_time
|
|
self.lora_paths = lora_paths
|
|
self.lora_scales = lora_scales
|
|
self.controlnet_image_path = controlnet_image_path
|
|
self.controlnet_strength = controlnet_strength
|
|
self.init_image_path = init_image_path
|
|
self.init_image_strength = init_image_strength
|
|
|
|
def save(
|
|
self,
|
|
path: t.Union[str, pathlib.Path],
|
|
export_json_metadata: bool = False,
|
|
overwrite: bool = False,
|
|
) -> None:
|
|
from mflux import ImageUtil
|
|
|
|
ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata, overwrite)
|
|
|
|
def _get_metadata(self) -> dict:
|
|
"""Generate metadata for reference as well as input data for
|
|
command line --config-from-metadata arg in future generations.
|
|
"""
|
|
return {
|
|
# mflux_version is used by future metadata readers
|
|
# to determine supportability of metadata-derived workflows
|
|
"mflux_version": GeneratedImage.get_version(),
|
|
"model": self.model_config.model_name,
|
|
"base_model": str(self.model_config.base_model),
|
|
"seed": self.seed,
|
|
"steps": self.steps,
|
|
"guidance": self.guidance if self.model_config.supports_guidance else None,
|
|
"precision": str(self.precision),
|
|
"quantize": self.quantization,
|
|
"generation_time_seconds": round(self.generation_time, 2),
|
|
"lora_paths": [str(p) for p in self.lora_paths] if self.lora_paths else None,
|
|
"lora_scales": [round(scale, 2) for scale in self.lora_scales] if self.lora_scales else None,
|
|
"init_image_path": str(self.init_image_path) if self.init_image_path else None,
|
|
"init_image_strength": self.init_image_strength if self.init_image_path else None,
|
|
"controlnet_image_path": str(self.controlnet_image_path) if self.controlnet_image_path else None,
|
|
"controlnet_strength": round(self.controlnet_strength, 2) if self.controlnet_strength else None,
|
|
"prompt": self.prompt,
|
|
}
|
|
|
|
@staticmethod
|
|
def get_version() -> str:
|
|
version = GeneratedImage._get_version_from_toml()
|
|
if version:
|
|
return version
|
|
|
|
# Fallback to installed package version
|
|
try:
|
|
return str(importlib.metadata.version("mflux"))
|
|
except importlib.metadata.PackageNotFoundError:
|
|
return "unknown"
|
|
|
|
@staticmethod
|
|
def _get_version_from_toml() -> str | None:
|
|
# Search for pyproject.toml by traversing up from the current working directory
|
|
current_dir = pathlib.Path(__file__).resolve().parent
|
|
for parent in current_dir.parents:
|
|
pyproject_path = parent / "pyproject.toml"
|
|
if pyproject_path.exists():
|
|
try:
|
|
pyproject_data = toml.load(pyproject_path)
|
|
return pyproject_data.get("project", {}).get("version")
|
|
except (toml.TomlDecodeError, KeyError, TypeError):
|
|
return None
|
|
return None
|