diff --git a/.gitignore b/.gitignore index ee9ae73..afe2251 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ *.png *.jpg *.pyc +*.safetensors diff --git a/main.py b/main.py index 126e783..d6a1790 100644 --- a/main.py +++ b/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( diff --git a/src/flux_1/flux.py b/src/flux_1/flux.py index 74b58a8..08e1036 100644 --- a/src/flux_1/flux.py +++ b/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) diff --git a/src/flux_1/weights/weight_handler.py b/src/flux_1/weights/weight_handler.py index 7ef50b6..9a3b821 100644 --- a/src/flux_1/weights/weight_handler.py +++ b/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):