Also save LoRA metadata
This commit is contained in:
parent
6d05f5b27e
commit
6fb51170de
@ -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,
|
||||
)
|
||||
|
||||
|
||||
@ -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,
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user