Support other format loras. update readme. Run code_style.

Author: jfacevedo-google

README.md

@@ -17,6 +17,7 @@
 [![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/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.
@@ -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
 
 
@@ -489,6 +489,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/flux/transformers/transformer_flux_flax.py

@@ -144,7 +144,7 @@ def __call__(self, hidden_states, temb, image_rotary_emb=None):
     hidden_states = self.linear2(attn_mlp)
     hidden_states = gate * hidden_states
     hidden_states = residual + hidden_states
-    if hidden_states.dtype == jnp.float16 or hidden_states.dtype == jnp.bfloat16:
+    if hidden_states.dtype == jnp.float16:
       hidden_states = jnp.clip(hidden_states, -65504, 65504)
 
     return hidden_states, temb, image_rotary_emb
@@ -294,7 +294,7 @@ def __call__(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=
 
     context_ff_output = self.txt_mlp(norm_encoder_hidden_states)
     encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output
-    if encoder_hidden_states.dtype == jnp.float16 or encoder_hidden_states.dtype == jnp.bfloat16:
+    if encoder_hidden_states.dtype == jnp.float16:
       encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
     return hidden_states, encoder_hidden_states, temb, image_rotary_emb
 

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

@@ -0,0 +1,88 @@
+import os
+import unittest
+import pytest
+
+import numpy as np
+
+from .. import pyconfig
+from absl.testing import absltest
+from maxdiffusion.generate_flux import run as generate_flux
+from PIL import Image
+from skimage.metrics import structural_similarity as ssim
+from google.cloud import storage
+
+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+
+JAX_CACHE_DIR = "gs://maxdiffusion-github-runner-test-assets/cache_dir"
+
+
+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)
+
+
+class GenerateFlux(unittest.TestCase):
+  """Smoke test."""
+
+  def setUp(self):
+    GenerateFlux.dummy_data = {}
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions")
+  def test_flux_dev(self):
+    img_url = os.path.join(THIS_DIR, "images", "test_flux_dev.png")
+    base_image = np.array(Image.open(img_url)).astype(np.uint8)
+    pyconfig.initialize(
+        [
+            None,
+            os.path.join(THIS_DIR, "..", "configs", "base_flux_dev.yml"),
+            "run_name=flux_test",
+            "output_dir=/tmp/",
+            "jax_cache_dir=/tmp/cache_dir",
+            'prompt="A cute corgi lives in a house made out of sushi, anime"',
+        ],
+        unittest=True,
+    )
+
+    images = generate_flux(pyconfig.config)
+    test_image = np.array(images[0]).astype(np.uint8)
+    ssim_compare = ssim(base_image, test_image, multichannel=True, channel_axis=-1, data_range=255)
+    assert base_image.shape == test_image.shape
+    assert ssim_compare >= 0.80
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions")
+  def test_flux_dev_lora(self):
+    img_url = os.path.join(THIS_DIR, "images", "test_flux_dev_lora.png")
+    base_image = np.array(Image.open(img_url)).astype(np.uint8)
+
+    gcs_lora_path = "gs://maxdiffusion-github-runner-test-assets/flux/lora/anime_lora.safetensors"
+    local_path = "/tmp/anime_lora.safetensors"
+    download_blob(gcs_lora_path, local_path)
+
+    pyconfig.initialize(
+        [
+            None,
+            os.path.join(THIS_DIR, "..", "configs", "base_flux_dev.yml"),
+            "run_name=flux_test",
+            "output_dir=/tmp/",
+            "jax_cache_dir=/tmp/cache_dir",
+            'prompt="A cute corgi lives in a house made out of sushi, anime"',
+            'lora_config={"lora_model_name_or_path" : ["/tmp/anime_lora.safetensors"], "weight_name" : ["anime_lora.safetensors"], "adapter_name" : ["anime"], "scale": [0.8], "from_pt": ["true"]}',
+        ],
+        unittest=True,
+    )
+
+    images = generate_flux(pyconfig.config)
+    test_image = np.array(images[1]).astype(np.uint8)
+    ssim_compare = ssim(base_image, test_image, multichannel=True, channel_axis=-1, data_range=255)
+    assert base_image.shape == test_image.shape
+    assert ssim_compare >= 0.80
+
+
+if __name__ == "__main__":
+  absltest.main()