initial lora implementation for flux

Author: jfacevedo-google

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -54,27 +54,17 @@ precision: "DEFAULT"
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash
-flash_block_sizes: {
-  "block_q" : 256,
-  "block_kv_compute" : 256,
-  "block_kv" : 256,
-  "block_q_dkv" : 256,
-  "block_kv_dkv" : 256,
-  "block_kv_dkv_compute" : 256,
-  "block_q_dq" : 256,
-  "block_kv_dq" : 256
-}
-
-# Use the following flash_block_sizes on v6e (Trillium).
+flash_block_sizes: {}
+# Use the following flash_block_sizes on v6e (Trillium) due to larger vmem.
 # flash_block_sizes: {
-#   "block_q" : 2176,
-#   "block_kv_compute" : 2176,
-#   "block_kv" : 2176,
-#   "block_q_dkv" : 2176,
-#   "block_kv_dkv" : 2176,
-#   "block_kv_dkv_compute" : 2176,
-#   "block_q_dq" : 2176,
-#   "block_kv_dq" : 2176
+#   "block_q" : 1536,
+#   "block_kv_compute" : 1536,
+#   "block_kv" : 1536,
+#   "block_q_dkv" : 1536,
+#   "block_kv_dkv" : 1536,
+#   "block_kv_dkv_compute" : 1536,
+#   "block_q_dq" : 1536,
+#   "block_kv_dq" : 1536
 # }
 # GroupNorm groups
 norm_num_groups: 32

src/maxdiffusion/generate_flux.py

@@ -16,6 +16,7 @@
 
 from typing import Callable, List, Union, Sequence
 from absl import app
+from contextlib import ExitStack
 import functools
 import math
 import time
@@ -24,6 +25,7 @@
 import jax
 from jax.sharding import Mesh, PositionalSharding, PartitionSpec as P
 import jax.numpy as jnp
+import flax.linen as nn
 from chex import Array
 from einops import rearrange
 from flax.linen import partitioning as nn_partitioning
@@ -39,6 +41,28 @@
     get_precision,
     setup_initial_state,
 )
+from maxdiffusion.loaders.flux_lora_pipeline import FluxLoraLoaderMixin
+
+def maybe_load_flux_lora(config, lora_loader, params):
+  def _noop_interceptor(next_fn, args, kwargs, context):
+    return next_fn(*args, **kwargs)
+
+  lora_config = config.lora_config
+  interceptors= [_noop_interceptor]
+  if len(lora_config["lora_model_name_or_path"]) > 0:
+    interceptors = []
+    for i in range(len(lora_config["lora_model_name_or_path"])):
+      params, rank, network_alphas = lora_loader.load_lora_weights(
+        config,
+        lora_config["lora_model_name_or_path"][i],
+        weight_name=lora_config["weight_name"][i],
+        params=params,
+        adapter_name=lora_config["adapter_name"][i],
+      )
+      interceptor = lora_loader.make_lora_interceptor(params, rank, network_alphas, lora_config["adapter_name"][i])
+      interceptors.append(interceptor)
+
+  return params, interceptors
 
 
 def unpack(x: Array, height: int, width: int) -> Array:
@@ -400,21 +424,29 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
 
   # loads pretrained weights
   transformer_params = load_flow_model(config.flux_name, transformer_eval_params, "cpu")
+  params = {}
+  params["transformer"] = transformer_params
+  # maybe load lora and create interceptor
+  lora_loader = FluxLoraLoaderMixin()
+  params, lora_interceptors = maybe_load_flux_lora(config, lora_loader, params)
+  transformer_params = params["transformer"]
   # create transformer state
   weights_init_fn = functools.partial(
       transformer.init_weights, rngs=rng, max_sequence_length=config.max_sequence_length, eval_only=False
   )
-  transformer_state, transformer_state_shardings = setup_initial_state(
-      model=transformer,
-      tx=None,
-      config=config,
-      mesh=mesh,
-      weights_init_fn=weights_init_fn,
-      model_params=None,
-      training=False,
-  )
-  transformer_state = transformer_state.replace(params=transformer_params)
-  transformer_state = jax.device_put(transformer_state, transformer_state_shardings)
+  with ExitStack() as stack:
+    _ = [stack.enter_context(nn.intercept_methods(interceptor)) for interceptor in lora_interceptors]
+    transformer_state, transformer_state_shardings = setup_initial_state(
+        model=transformer,
+        tx=None,
+        config=config,
+        mesh=mesh,
+        weights_init_fn=weights_init_fn,
+        model_params=None,
+        training=False,
+    )
+    transformer_state = transformer_state.replace(params=transformer_params)
+    transformer_state = jax.device_put(transformer_state, transformer_state_shardings)
   get_memory_allocations()
 
   states = {}
@@ -444,17 +476,23 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
       out_shardings=None,
   )
   t0 = time.perf_counter()
-  p_run_inference(states).block_until_ready()
+  with ExitStack() as stack:
+    _ = [stack.enter_context(nn.intercept_methods(interceptor)) for interceptor in lora_interceptors]
+    p_run_inference(states).block_until_ready()
   t1 = time.perf_counter()
   max_logging.log(f"Compile time: {t1 - t0:.1f}s.")
 
   t0 = time.perf_counter()
-  imgs = p_run_inference(states).block_until_ready()
+  with ExitStack() as stack, jax.profiler.trace("/home/jfacevedo/trace/"):
+    _ = [stack.enter_context(nn.intercept_methods(interceptor)) for interceptor in lora_interceptors]
+    imgs = p_run_inference(states).block_until_ready()
   t1 = time.perf_counter()
   max_logging.log(f"Inference time: {t1 - t0:.1f}s.")
 
   t0 = time.perf_counter()
-  imgs = p_run_inference(states).block_until_ready()
+  with ExitStack() as stack:
+    _ = [stack.enter_context(nn.intercept_methods(interceptor)) for interceptor in lora_interceptors]
+    imgs = p_run_inference(states).block_until_ready()
   imgs = jax.experimental.multihost_utils.process_allgather(imgs, tiled=True)
   t1 = time.perf_counter()
   max_logging.log(f"Inference time: {t1 - t0:.1f}s.")

src/maxdiffusion/loaders/__init__.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 from .lora_pipeline import StableDiffusionLoraLoaderMixin
-from .flux_lora_pipeline import FluxLoraLoaderMixin
+from .flux_lora_pipeline import FluxLoraLoaderMixin

src/maxdiffusion/loaders/flux_lora_pipeline.py

@@ -16,30 +16,30 @@
 from .lora_base import LoRABaseMixin
 from ..models.lora import LoRALinearLayer, BaseLoRALayer
 import jax.numpy as jnp
-from flax.traverse_util import flatten_dict
+from flax.traverse_util import flatten_dict, unflatten_dict
+from flax.core.frozen_dict import unfreeze
 from ..models.modeling_flax_pytorch_utils import convert_flux_lora_pytorch_state_dict_to_flax
 from huggingface_hub.utils import validate_hf_hub_args
-
-
+from maxdiffusion.models.modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor)
 class FluxLoraLoaderMixin(LoRABaseMixin):
 
   _lora_lodable_modules = ["transformer", "text_encoder"]
-
+  
   def load_lora_weights(
       self,
       config,
       pretrained_model_name_or_path_or_dict: Union[str, Dict[str, jnp.ndarray]],
       params,
       adapter_name=None,
-      **kwargs,
+      **kwargs
   ):
     state_dict = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
 
     params, rank, network_alphas = self.load_lora(
-        config,
-        state_dict,
-        params=params,
-        adapter_name=adapter_name,
+      config, 
+      state_dict,
+      params=params,
+      adapter_name=adapter_name,
     )
 
     return params, rank, network_alphas
@@ -53,7 +53,7 @@ def rename_for_interceptor(params_keys, network_alphas, adapter_name):
         new_layer_lora = layer_lora[: layer_lora.index(lora_name)]
         if new_layer_lora not in new_params_keys:
           new_params_keys.append(new_layer_lora)
-          network_alpha = network_alphas.get(layer_lora, None)
+          network_alpha = network_alphas[layer_lora]
           new_network_alphas[new_layer_lora] = network_alpha
     return new_params_keys, new_network_alphas
 
@@ -64,7 +64,7 @@ def make_lora_interceptor(cls, params, rank, network_alphas, adapter_name):
     transformer_keys = flatten_dict(params["transformer"]).keys()
     lora_keys, transformer_alphas = cls.rename_for_interceptor(transformer_keys, network_alphas, adapter_name)
     network_alphas_for_interceptor.update(transformer_alphas)
-
+  
     def _intercept(next_fn, args, kwargs, context):
       mod = context.module
       while mod is not None:
@@ -107,6 +107,7 @@ def lora_state_dict(cls, pretrained_model_name_or_path: str, **kwargs):
     revision = kwargs.pop("revision", None)
     subfolder = kwargs.pop("subfolder", None)
     weight_name = kwargs.pop("weight_name", None)
+    unet_config = kwargs.pop("unet_config", None)
     use_safetensors = kwargs.pop("use_safetensors", None)
     resume_download = kwargs.pop("resume_download", False)
 
@@ -137,8 +138,8 @@ def lora_state_dict(cls, pretrained_model_name_or_path: str, **kwargs):
     )
 
     return state_dict
-
+  
   @classmethod
   def load_lora(cls, config, state_dict, params, adapter_name=None):
     params, rank, network_alphas = convert_flux_lora_pytorch_state_dict_to_flax(config, state_dict, params, adapter_name)
-    return params, rank, network_alphas
+    return params, rank, network_alphas

src/maxdiffusion/models/modeling_flax_pytorch_utils.py

@@ -222,6 +222,51 @@ def create_flax_params_from_pytorch_state(
       renamed_network_alphas[tuple(flax_key_list)] = network_alpha_value
   return unet_state_dict, text_encoder_state_dict, text_encoder_2_state_dict, rank, renamed_network_alphas
 
+def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params, adapter_name):
+  pt_state_dict = {k: v.float().numpy() for k, v in pt_state_dict.items()}
+  transformer_params = flatten_dict(unfreeze(params["transformer"]))
+  network_alphas = {}
+  rank = None
+  for pt_key, tensor in pt_state_dict.items():
+    renamed_pt_key = rename_key(pt_key)
+    print("renamed_pt_key:", renamed_pt_key)
+    renamed_pt_key = renamed_pt_key.replace("lora_unet_", "")
+    renamed_pt_key = renamed_pt_key.replace("lora_down", f"lora-{adapter_name}.down")
+    renamed_pt_key = renamed_pt_key.replace("lora_up", f"lora-{adapter_name}.up")
+
+    if "double_blocks" in renamed_pt_key:
+      renamed_pt_key = renamed_pt_key.replace("_img_attn_proj", ".attn.i_proj")
+      renamed_pt_key = renamed_pt_key.replace("_img_attn_qkv", ".attn.i_qkv")
+      renamed_pt_key = renamed_pt_key.replace("_img_mlp_0", ".img_mlp.layers_0")
+      renamed_pt_key = renamed_pt_key.replace("_img_mlp_2", ".img_mlp.layers_2")
+      renamed_pt_key = renamed_pt_key.replace("_img_mod_lin", ".img_norm1.lin")
+      renamed_pt_key = renamed_pt_key.replace("_txt_attn_proj", ".attn.e_proj")
+      renamed_pt_key = renamed_pt_key.replace("_txt_attn_qkv", ".attn.e_qkv")
+      renamed_pt_key = renamed_pt_key.replace("_txt_mlp_0", ".txt_mlp.layers_0")
+      renamed_pt_key = renamed_pt_key.replace("_txt_mlp_2", ".txt_mlp.layers_2")
+      renamed_pt_key = renamed_pt_key.replace("_txt_mod_lin", ".txt_norm1.lin")
+    elif "single_blocks" in renamed_pt_key:
+      renamed_pt_key = renamed_pt_key.replace("_linear1", ".linear1")
+      renamed_pt_key = renamed_pt_key.replace("_linear2", ".linear2")
+      renamed_pt_key = renamed_pt_key.replace("_modulation_lin", ".norm.lin")
+    
+    renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+    
+    pt_tuple_key = tuple(renamed_pt_key.split("."))
+    if "alpha" in pt_tuple_key:
+      pt_tuple_key = pt_tuple_key[:-1] + (f"lora-{adapter_name}", 'down', 'kernel')
+      network_alphas[tuple([*pt_tuple_key])] = tensor.item()
+      pt_tuple_key = pt_tuple_key[:-1] + (f"lora-{adapter_name}", 'up', 'kernel')
+      network_alphas[tuple([*pt_tuple_key])] = tensor.item()
+    else:
+      if pt_tuple_key[-2] == "up":
+        rank = tensor.shape[1]
+      transformer_params[tuple([*pt_tuple_key])] = jnp.asarray(tensor.T, dtype=config.weights_dtype)
+  
+  params["transformer"] = unflatten_dict(transformer_params)
+
+  return params, rank, network_alphas
+
 
 def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params, adapter_name):
   pt_state_dict = {k: v.float().numpy() for k, v in pt_state_dict.items()}