From 6fb51170dec55fa3e3a98c17d45230da7d5e0f94 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Fri, 6 Sep 2024 20:08:53 +0200
Subject: [PATCH] Also save LoRA metadata

---
 src/flux_1/flux.py                       | 11 ++++++++++-
 src/flux_1/post_processing/image.py      |  6 ++++++
 src/flux_1/post_processing/image_util.py |  4 ++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py
index c07e6f1..50511f2 100644
--- a/src/flux_1/flux.py
+++ b/src/flux_1/flux.py
@@ -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,
         )
 
diff --git a/src/flux_1/post_processing/image.py b/src/flux_1/post_processing/image.py
index 674fbfc..542bfc7 100644
--- a/src/flux_1/post_processing/image.py
+++ b/src/flux_1/post_processing/image.py
@@ -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,
         }
diff --git a/src/flux_1/post_processing/image_util.py b/src/flux_1/post_processing/image_util.py
index e83bf6e..b5edd08 100644
--- a/src/flux_1/post_processing/image_util.py
+++ b/src/flux_1/post_processing/image_util.py
@@ -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