Support other format loras. update readme. Run code_style.

Author: jfacevedo-google

README.md

@@ -17,7 +17,8 @@
 [![Unit Tests](https://github.com/google/maxtext/actions/workflows/UnitTests.yml/badge.svg)](https://github.com/google/maxdiffusion/actions/workflows/UnitTests.yml)
 
 # What's new?
-- **`2025/02/08**: Flux schnell & dev inference.
+- **`2025/02/12`**: Flux LoRA for inference.
+- **`2025/02/08`**: Flux schnell & dev inference.
 - **`2024/12/12`**: Load multiple LoRAs for inference.
 - **`2024/10/22`**: LoRA support for Hyper SDXL.
 - **`2024/8/1`**: Orbax is the new default checkpointer. You can still use `pipeline.save_pretrained` after training to save in diffusers format.
@@ -47,7 +48,8 @@ MaxDiffusion supports
     * [Training](#training)
       * [Dreambooth](#dreambooth)
     * [Inference](#inference)
-      * [Flux](#flux) 
+      * [Flux](#flux)
+        * [Flux LoRA](#flux-lora)
       * [Hyper-SD XL LoRA](#hyper-sdxl-lora)
       * [Load Multiple LoRA](#load-multiple-lora)
       * [SDXL Lightning](#sdxl-lightning)
@@ -169,6 +171,24 @@ To generate images, run the following command:
   python src/maxdiffusion/generate_flux.py src/maxdiffusion/configs/base_flux_schnell.yml jax_cache_dir=/tmp/cache_dir run_name=flux_test output_dir=/tmp/ prompt="photograph of an electronics chip in the shape of a race car with trillium written on its side" per_device_batch_size=1 ici_data_parallelism=1 ici_fsdp_parallelism=-1 offload_encoders=False
   ```
 
+    ## Flux LoRA
+
+    Disclaimer: not all LoRA formats have been tested. If there is a specific LoRA that doesn't load, please let us know.
+
+    Tested with [Amateur Photography](https://civitai.com/models/652699/amateur-photography-flux-dev) and [XLabs-AI](https://huggingface.co/XLabs-AI/flux-lora-collection/tree/main) LoRA collection.
+
+    First download the LoRA file to a local directory, for example, `/home/jfacevedo/anime_lora.safetensors`. Then run as follows:
+
+    ```bash
+    python src/maxdiffusion/generate_flux.py src/maxdiffusion/configs/base_flux_dev.yml jax_cache_dir=/tmp/cache_dir run_name=flux_test output_dir=/tmp/ prompt='A cute corgi lives in a house made out of sushi, anime' num_inference_steps=28 ici_data_parallelism=1 ici_fsdp_parallelism=-1 split_head_dim=True lora_config='{"lora_model_name_or_path" : ["/home/jfacevedo/anime_lora.safetensors"], "weight_name" : ["anime_lora.safetensors"], "adapter_name" : ["anime"], "scale": [0.8], "from_pt": ["true"]}'
+    ```
+
+    Loading multiple LoRAs is supported as follows:
+
+    ```bash
+    python src/maxdiffusion/generate_flux.py src/maxdiffusion/configs/base_flux_dev.yml jax_cache_dir=/tmp/cache_dir run_name=flux_test output_dir=/tmp/ prompt='A cute corgi lives in a house made out of sushi, anime' num_inference_steps=28 ici_data_parallelism=1 ici_fsdp_parallelism=-1 split_head_dim=True lora_config='{"lora_model_name_or_path" : ["/home/jfacevedo/anime_lora.safetensors", "/home/jfacevedo/amateurphoto-v6-forcu.safetensors"], "weight_name" : ["anime_lora.safetensors","amateurphoto-v6-forcu.safetensors"], "adapter_name" : ["anime","realistic"], "scale": [0.6, 0.6], "from_pt": ["true","true"]}'
+    ```
+
   ## Hyper SDXL LoRA
 
   Supports Hyper-SDXL models from [ByteDance](https://huggingface.co/ByteDance/Hyper-SD)

src/maxdiffusion/generate_flux.py

@@ -43,25 +43,25 @@
 )
 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]
+  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],
+          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
 
 
@@ -501,6 +501,8 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
   for i, image in enumerate(imgs):
     Image.fromarray(image).save(f"flux_{i}.png")
 
+  return imgs
+
 
 def main(argv: Sequence[str]) -> None:
   pyconfig.initialize(argv)

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

@@ -21,25 +21,27 @@
 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
@@ -64,7 +66,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:
@@ -138,8 +140,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,7 @@ 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"]))
@@ -243,7 +244,7 @@ def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params,
       renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora1.up", f"attn.i_qkv.lora-{adapter_name}.up")
       renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora2.down", f"attn.e_qkv.lora-{adapter_name}.down")
       renamed_pt_key = renamed_pt_key.replace("processor.qkv_lora2.up", f"attn.e_qkv.lora-{adapter_name}.up")
-      
+
       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")
@@ -258,20 +259,20 @@ def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params,
       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')
+      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')
+      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

src/maxdiffusion/tests/generate_flux_smoke_test.py

@@ -22,6 +22,7 @@ def download_blob(gcs_file, local_file):
   gcs_dir_arr = gcs_file.replace("gs://", "").split("/")
   storage_client = storage.Client()
   bucket = storage_client.get_bucket(gcs_dir_arr[0])
+  blob_loc = "/".join(gcs_dir_arr[1:])
   blob = bucket.blob("/".join(gcs_dir_arr[1:]))
   blob.download_to_filename(local_file)