Apply lora with --apply-lora <path-to-safetensors>

.gitignore

@@ -11,3 +11,4 @@
 *.png
 *.jpg
 *.pyc
+*.safetensors

main.py

@@ -23,6 +23,8 @@ def main():
     parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
     parser.add_argument('--quantize',  "-q", type=int, choices=[4, 8], default=None, help='Quantize the model (4 or 8, Default is None)')
     parser.add_argument('--path', type=str, default=None, help='Local path for loading a model from disk')
+    parser.add_argument('--apply-lora', type=str, default=None, help='Local safetensors for applying LORA from disk')
+
 
     args = parser.parse_args()
 
@@ -34,7 +36,8 @@ def main():
     flux = Flux1(
         model_config=ModelConfig.from_alias(args.model),
         quantize_full_weights=args.quantize,
-        local_path=args.path
+        local_path=args.path,
+        lora_path=args.apply_lora
     )
 
     image = flux.generate_image(

src/flux_1/flux.py

@@ -19,6 +19,9 @@ from flux_1.tokenizer.clip_tokenizer import TokenizerCLIP
 from flux_1.tokenizer.t5_tokenizer import TokenizerT5
 from flux_1.tokenizer.tokenizer_handler import TokenizerHandler
 from flux_1.weights.weight_handler import WeightHandler
+import safetensors
+from safetensors import safe_open
+
 
 
 class Flux1:
@@ -28,15 +31,16 @@ class Flux1:
             model_config: ModelConfig,
             quantize_full_weights: int | None = None,
             local_path: str | None = None,
+            lora_path: str | None = None
     ):
         self.model_config = model_config
         self.quantize_full_weights = quantize_full_weights
+        self.lora_path = lora_path
 
         # Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface
         tokenizers = TokenizerHandler(model_config.model_name, self.model_config.max_sequence_length, local_path)
         self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
         self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
-
         # Initialize the models
         self.vae = VAE()
         self.transformer = Transformer(model_config)
@@ -44,8 +48,8 @@ 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)
-
+        weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path,lora_path=lora_path)
+        
         # Set the loaded weights if they are not quantized
         if weights.quantization_level is None:
             self._set_model_weights(weights)
@@ -61,7 +65,7 @@ class Flux1:
         # If loading previously saved quantized weights, the weights must be set after modules have been quantized
         if weights.quantization_level is not None:
             self._set_model_weights(weights)
-
+            
     def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
         # Create a new runtime config based on the model type and input parameters
         config = RuntimeConfig(config, self.model_config)

src/flux_1/weights/weight_handler.py

@@ -6,14 +6,23 @@ from mlx.utils import tree_unflatten
 from safetensors import safe_open
 
 from flux_1.config.config import Config
+import icecream as ic
+import json
+from mlx.utils import tree_flatten
+from functools import reduce
 
-
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super(NumpyEncoder, self).default(obj)
 class WeightHandler:
 
     def __init__(
             self,
             repo_id: str | None = None,
             local_path: str | None = None,
+            lora_path: str | None = None
     ):
         root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
 
@@ -21,6 +30,77 @@ class WeightHandler:
         self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path)
         self.vae, _ = WeightHandler._vae(root_path=root_path)
         self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path)
+        if(lora_path is not None):
+            self.lora_transformer,self.lora_quantization_level= WeightHandler._lora_transformer(lora_path=lora_path)
+            if 'transformer' not in self.lora_transformer:
+                pass
+            self._apply_transformer(self.transformer,self.lora_transformer['transformer'])
+        
+    def _apply_transformer(self,transformer,lora_transformer):
+        lora_weights = tree_flatten(lora_transformer)
+        visited={}
+        
+        for key,weight in lora_weights:
+            splits=key.split(".")
+            target=transformer
+            visiting=[]
+            for splitKey in splits:
+                if isinstance(target,dict) and splitKey in target:
+                    target=target[splitKey]
+                    visiting.append(splitKey)
+                elif isinstance(target,list) and len(target)>0:
+                    if(len(target)< int(splitKey)):
+                        for _ in range(int(splitKey)-len(target)+1):
+                            target.append({})
+                        
+                    target=target[int(splitKey)]
+                    visiting.append(splitKey)
+                else:
+                    parentKey=".".join(visiting)
+                    if(parentKey in visited and 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]):
+                        continue
+                    if not splitKey.startswith("lora_"):
+                        visiting.append(splitKey)
+                        parentKey=".".join(visiting)
+                        if(splitKey=="net"):
+                            target['net']=list({})
+                            target=target['net']
+                        elif (splitKey=="0"):
+                            target.append({})
+                            target=target[0]
+                            continue
+                        elif (splitKey=="proj"):
+                            target[splitKey]=weight
+                            if parentKey not in visited:
+                                visited[parentKey]={}
+                        continue
+                    if parentKey not in visited:
+                        visited[parentKey]={}
+                    visited[parentKey][splitKey]=weight
+                    if not 'weight' in target:
+                        continue
+                    if  'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
+                        lora_a=visited[parentKey]['lora_A']
+                        lora_b=visited[parentKey]['lora_B']
+                        transWeight=target['weight']
+                        weight=transWeight + lora_b @lora_a
+                        target['weight']=weight
+                        
+                        
+
+    
+
+    @staticmethod
+    def _lora_transformer(lora_path: Path) -> (dict, int):
+        weights = []
+        quantization_level = safe_open(lora_path, framework="pt").metadata().get("quantization_level")
+        weight = list(mx.load(str(lora_path)).items())
+        weights.extend(weight)
+        weights = [WeightHandler._reshape_weights(k, v) for k, v in weights]
+        weights = WeightHandler._flatten(weights)
+        unflatten = tree_unflatten(weights)
+        return unflatten, quantization_level
+
 
     @staticmethod
     def _clip_encoder(root_path: Path) -> (dict, int):
@@ -74,6 +154,8 @@ class WeightHandler:
                     "linear2": block["ff_context"]["net"][2]
                 }
         return weights, quantization_level
+    
+    
 
     @staticmethod
     def _vae(root_path: Path) -> (dict, int):