monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Code style improvements and some restructuring

Browse Source
This commit is contained in:
filipstrand 2024-09-04 19:48:15 +02:00
parent b4c27a6fab
commit 8983070728
3 changed files with 60 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
main.py
View File

@ -26,13 +26,11 @@ def main():
parser.add_argument('--apply-lora', type=str, nargs='*', default=[], help='Local safetensors for applying LORA from disk') parser.add_argument('--apply-lora', type=str, nargs='*', default=[], help='Local safetensors for applying LORA from disk')
parser.add_argument('--lora-scales', type=float,nargs='*', default=[1.0], help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.') parser.add_argument('--lora-scales', type=float,nargs='*', default=[1.0], help='Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.')
args = parser.parse_args() args = parser.parse_args()
if args.path and args.model is None: if args.path and args.model is None:
parser.error("--model must be specified when using --path") parser.error("--model must be specified when using --path")
seed = int(time.time()) if args.seed is None else args.seed
flux = Flux1( flux = Flux1(
model_config=ModelConfig.from_alias(args.model), model_config=ModelConfig.from_alias(args.model),
quantize_full_weights=args.quantize, quantize_full_weights=args.quantize,
@ -42,7 +40,7 @@ def main():
) )
image = flux.generate_image( image = flux.generate_image(
seed=seed, seed=int(time.time()) if args.seed is None else args.seed,
prompt=args.prompt, prompt=args.prompt,
config=Config( config=Config(
num_inference_steps=args.steps, num_inference_steps=args.steps,

10
src/flux_1/flux.py
View File

@ -21,7 +21,6 @@ from flux_1.tokenizer.tokenizer_handler import TokenizerHandler
from flux_1.weights.weight_handler import WeightHandler from flux_1.weights.weight_handler import WeightHandler
class Flux1: class Flux1:
def __init__( def __init__(
@ -29,17 +28,17 @@ class Flux1:
model_config: ModelConfig, model_config: ModelConfig,
quantize_full_weights: int | None = None, quantize_full_weights: int | None = None,
local_path: str | None = None, local_path: str | None = None,
lora_files: [str] =[], lora_files: [str] = [],
lora_scales: [float] = [1.0] lora_scales: [float] = [1.0]
): ):
self.model_config = model_config self.model_config = model_config
self.quantize_full_weights = quantize_full_weights self.quantize_full_weights = quantize_full_weights
self.lora_files = lora_files
# Load and initialize the tokenizers from disk, huggingface cache, or download from huggingface # 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) 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.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
self.clip_tokenizer = TokenizerCLIP(tokenizers.clip) self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
# Initialize the models # Initialize the models
self.vae = VAE() self.vae = VAE()
self.transformer = Transformer(model_config) self.transformer = Transformer(model_config)
@ -47,7 +46,8 @@ class Flux1:
self.clip_text_encoder = CLIPEncoder() self.clip_text_encoder = CLIPEncoder()
# Load the weights from disk, huggingface cache, or download from huggingface # Load the weights from disk, huggingface cache, or download from huggingface
weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path,lora_files=lora_files,lora_scales=lora_scales) weights = WeightHandler(repo_id=model_config.model_name, local_path=local_path, lora_files=lora_files, lora_scales=lora_scales)
# Set the loaded weights if they are not quantized # Set the loaded weights if they are not quantized
if weights.quantization_level is None: if weights.quantization_level is None:
self._set_model_weights(weights) self._set_model_weights(weights)
@ -63,7 +63,7 @@ class Flux1:
# If loading previously saved quantized weights, the weights must be set after modules have been quantized # If loading previously saved quantized weights, the weights must be set after modules have been quantized
if weights.quantization_level is not None: if weights.quantization_level is not None:
self._set_model_weights(weights) self._set_model_weights(weights)
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image: 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 # Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config) config = RuntimeConfig(config, self.model_config)

109
src/flux_1/weights/weight_handler.py
View File

@ -1,102 +1,104 @@
import logging
from pathlib import Path from pathlib import Path
import mlx.core as mx import mlx.core as mx
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
from mlx.utils import tree_flatten
from mlx.utils import tree_unflatten from mlx.utils import tree_unflatten
from safetensors import safe_open from safetensors import safe_open
from flux_1.config.config import Config from flux_1.config.config import Config
from mlx.utils import tree_flatten
import logging
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
class WeightHandler: class WeightHandler:
def __init__( def __init__(
self, self,
repo_id: str | None = None, repo_id: str | None = None,
local_path: str | None = None, local_path: str | None = None,
lora_files: [str] =[], lora_files=None,
lora_scales: [float] = [1.0] lora_scales=None
): ):
if lora_files is None:
lora_files = []
if lora_scales is None:
lora_scales = [1.0]
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id) root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
self.clip_encoder, _ = WeightHandler._clip_encoder(root_path=root_path) self.clip_encoder, _ = WeightHandler._clip_encoder(root_path=root_path)
self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path) self.t5_encoder, _ = WeightHandler._t5_encoder(root_path=root_path)
self.vae, _ = WeightHandler._vae(root_path=root_path) self.vae, _ = WeightHandler._vae(root_path=root_path)
self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path) self.transformer, self.quantization_level = WeightHandler._transformer(root_path=root_path)
if(lora_files): if lora_files:
if(len(lora_files)< len(lora_scales)): if len(lora_files) < len(lora_scales):
lora_scales=lora_scales[0:len(lora_files)] lora_scales = lora_scales[0:len(lora_files)]
if(len(lora_scales)<len(lora_files)): if len(lora_scales) < len(lora_files):
lora_scales= lora_scales + (len(lora_files)-len(lora_scales)) * [1.0] lora_scales = lora_scales + (len(lora_files) - len(lora_scales)) * [1.0]
for lora_file, lora_scale in zip(lora_files, lora_scales): for lora_file, lora_scale in zip(lora_files, lora_scales):
if( lora_scale<0.0 or lora_scale>1.0): if lora_scale < 0.0 or lora_scale > 1.0:
raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0-1.0] ") raise Exception(f"Invalid scale {lora_scale} provided for {lora_file}. Valid Range [0.0-1.0] ")
try: try:
lora_transformer,_ = WeightHandler._lora_transformer(lora_file=lora_file) lora_transformer, _ = WeightHandler._lora_transformer(lora_file=lora_file)
if 'transformer' not in lora_transformer: if 'transformer' not in lora_transformer:
raise Exception("The key `transformer` is missing in the LoRA safetensors file. Please ensure that the file is correctly formatted and contains the expected keys.") raise Exception(
self._apply_transformer(self.transformer,lora_transformer['transformer'],lora_scale) "The key `transformer` is missing in the LoRA safetensors file. Please ensure that the file is correctly formatted and contains the expected keys.")
WeightHandler._apply_transformer(self.transformer, lora_transformer['transformer'], lora_scale)
except Exception as e: except Exception as e:
log.error(f"Error loading the LoRA safetensors file: {e}") log.error(f"Error loading the LoRA safetensors file: {e}")
@staticmethod
def _apply_transformer(self,transformer,lora_transformer,lora_scale): def _apply_transformer(transformer, lora_transformer, lora_scale):
lora_weights = tree_flatten(lora_transformer) lora_weights = tree_flatten(lora_transformer)
visited={} visited = {}
for key,weight in lora_weights: for key, weight in lora_weights:
splits=key.split(".") splits = key.split(".")
target=transformer target = transformer
visiting=[] visiting = []
for splitKey in splits: for splitKey in splits:
if isinstance(target,dict) and splitKey in target: if isinstance(target, dict) and splitKey in target:
target=target[splitKey] target = target[splitKey]
visiting.append(splitKey) visiting.append(splitKey)
elif isinstance(target,list) and len(target)>0: elif isinstance(target, list) and len(target) > 0:
if(len(target)< int(splitKey)): if len(target) < int(splitKey):
for _ in range(int(splitKey)-len(target)+1): for _ in range(int(splitKey) - len(target) + 1):
target.append({}) target.append({})
target=target[int(splitKey)] target = target[int(splitKey)]
visiting.append(splitKey) visiting.append(splitKey)
else: else:
parentKey=".".join(visiting) parentKey = ".".join(visiting)
if(parentKey in visited and 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]): if parentKey in visited and 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
continue continue
if not splitKey.startswith("lora_"): if not splitKey.startswith("lora_"):
visiting.append(splitKey) visiting.append(splitKey)
parentKey=".".join(visiting) parentKey = ".".join(visiting)
if(splitKey=="net"): if splitKey == "net":
target['net']=list({}) target['net'] = list({})
target=target['net'] target = target['net']
elif (splitKey=="0"): elif splitKey == "0":
target.append({}) target.append({})
target=target[0] target = target[0]
continue continue
elif (splitKey=="proj"): elif splitKey == "proj":
target[splitKey]=weight target[splitKey] = weight
if parentKey not in visited: if parentKey not in visited:
visited[parentKey]={} visited[parentKey] = {}
continue continue
if parentKey not in visited: if parentKey not in visited:
visited[parentKey]={} visited[parentKey] = {}
visited[parentKey][splitKey]=weight visited[parentKey][splitKey] = weight
if not 'weight' in target: if not 'weight' in target:
raise ValueError(f"LoRA weights for layer {parentKey} cannot be loaded into the model.") raise ValueError(f"LoRA weights for layer {parentKey} cannot be loaded into the model.")
if 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]: if 'lora_A' in visited[parentKey] and 'lora_B' in visited[parentKey]:
lora_a=visited[parentKey]['lora_A'] lora_a = visited[parentKey]['lora_A']
lora_b=visited[parentKey]['lora_B'] lora_b = visited[parentKey]['lora_B']
transWeight=target['weight'] transWeight = target['weight']
weight=transWeight + lora_scale* (lora_b @lora_a) weight = transWeight + lora_scale * (lora_b @ lora_a)
target['weight']=weight target['weight'] = weight
@staticmethod @staticmethod
def _lora_transformer(lora_file: Path) -> (dict, int): def _lora_transformer(lora_file: Path) -> (dict, int):
@ -117,7 +119,6 @@ class WeightHandler:
} }
return unflatten, quantization_level return unflatten, quantization_level
@staticmethod @staticmethod
def _clip_encoder(root_path: Path) -> (dict, int): def _clip_encoder(root_path: Path) -> (dict, int):
weights, quantization_level = WeightHandler._get_weights("text_encoder", root_path) weights, quantization_level = WeightHandler._get_weights("text_encoder", root_path)
@ -170,8 +171,6 @@ class WeightHandler:
"linear2": block["ff_context"]["net"][2] "linear2": block["ff_context"]["net"][2]
} }
return weights, quantization_level return weights, quantization_level
@staticmethod @staticmethod
def _vae(root_path: Path) -> (dict, int): def _vae(root_path: Path) -> (dict, int):