68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
import importlib
|
|
|
|
import PIL.Image
|
|
import mlx.core as mx
|
|
|
|
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 | None = None,
|
|
controlnet_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 = controlnet_image_path
|
|
self.controlnet_strength = controlnet_strength
|
|
|
|
def save(self, path: str, export_json_metadata: bool = False) -> None:
|
|
from mflux import ImageUtil
|
|
|
|
ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata)
|
|
|
|
def _get_metadata(self) -> dict:
|
|
return {
|
|
"mflux_version": str(GeneratedImage.get_version()),
|
|
"model": str(self.model_config.alias),
|
|
"seed": str(self.seed),
|
|
"steps": str(self.steps),
|
|
"guidance": "None" if self.model_config == ModelConfig.FLUX1_SCHNELL else str(self.guidance),
|
|
"precision": f"{self.precision}",
|
|
"quantization": "None" if self.quantization is None else f"{self.quantization} bit",
|
|
"generation_time": f"{self.generation_time:.2f} seconds",
|
|
"lora_paths": ", ".join(self.lora_paths) if self.lora_paths else "None",
|
|
"lora_scales": ", ".join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else "None",
|
|
"prompt": self.prompt,
|
|
"controlnet_image": "None" if self.controlnet_image is None else self.controlnet_image,
|
|
"controlnet_strength": "None" if self.controlnet_strength is None else f"{self.controlnet_strength:.2f}",
|
|
}
|
|
|
|
@staticmethod
|
|
def get_version():
|
|
try:
|
|
return importlib.metadata.version("mflux")
|
|
except importlib.metadata.PackageNotFoundError:
|
|
return "unknown"
|