ruff

Author: jfacevedo-google

src/maxdiffusion/loaders/flux_lora_pipeline.py

@@ -16,11 +16,9 @@
 from .lora_base import LoRABaseMixin
 from ..models.lora import LoRALinearLayer, BaseLoRALayer
 import jax.numpy as jnp
-from flax.traverse_util import flatten_dict, unflatten_dict
-from flax.core.frozen_dict import unfreeze
+from flax.traverse_util import flatten_dict
 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):
@@ -109,7 +107,6 @@ 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)
 

src/maxdiffusion/models/modeling_flax_pytorch_utils.py

@@ -265,13 +265,13 @@ def convert_flux_lora_pytorch_state_dict_to_flax(config, pt_state_dict, params,
     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()
+      network_alphas[tuple([*pt_tuple_key])] = tensor.item()  # noqa: C409
       pt_tuple_key = pt_tuple_key[:-1] + (f"lora-{adapter_name}", "up", "kernel")
-      network_alphas[tuple([*pt_tuple_key])] = tensor.item()
+      network_alphas[tuple([*pt_tuple_key])] = tensor.item()  # noqa: C409
     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)
+      transformer_params[tuple([*pt_tuple_key])] = jnp.asarray(tensor.T, dtype=config.weights_dtype)  # noqa: C409
 
   params["transformer"] = unflatten_dict(transformer_params)
 

src/maxdiffusion/tests/generate_flux_smoke_test.py

@@ -22,7 +22,6 @@ 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)