Also save LoRA metadata

This commit is contained in:
filipstrand 2024-09-06 20:08:53 +02:00
parent 6d05f5b27e
commit 6fb51170de
3 changed files with 20 additions and 1 deletions

View File

@ -30,6 +30,8 @@ class Flux1:
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
):
self.lora_paths = lora_paths
self.lora_scales = lora_scales
self.model_config = model_config
# Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
@ -44,7 +46,12 @@ class Flux1:
self.clip_text_encoder = CLIPEncoder()
# Load the weights from disk, huggingface cache, or download from huggingface
weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path, lora_paths=lora_paths, lora_scales=lora_scales)
weights = WeightHandler(
repo_id=model_config.model_name,
local_path=local_path,
lora_paths=lora_paths,
lora_scales=lora_scales
)
# Set the loaded weights if they are not quantized
if weights.quantization_level is None:
@ -106,6 +113,8 @@ class Flux1:
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict['elapsed'],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
)

View File

@ -24,6 +24,8 @@ class Image:
precision: mx.Dtype,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
):
self.image = image
self.model_config = model_config
@ -34,6 +36,8 @@ class Image:
self.precision = precision
self.quantization = quantization
self.generation_time = generation_time
self.lora_paths = lora_paths
self.lora_scales = lora_scales
def save(self, path: str, export_json_metadata: bool = False) -> None:
file_path = Path(path)
@ -116,5 +120,7 @@ class Image:
'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 '',
'lora_scales': ', '.join([f"{scale:.2f}" for scale in self.lora_scales]) if self.lora_scales else '',
'prompt': self.prompt,
}

View File

@ -16,6 +16,8 @@ class ImageUtil:
prompt: str,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
config: RuntimeConfig,
) -> Image:
normalized = ImageUtil._denormalize(decoded_latents)
@ -31,6 +33,8 @@ class ImageUtil:
precision=config.precision,
quantization=quantization,
generation_time=generation_time,
lora_paths=lora_paths,
lora_scales=lora_scales,
)
@staticmethod