ruff + code_style

Author: ksikiric

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -58,9 +58,7 @@ def create_orbax_checkpoint_manager(
   p = epath.Path(checkpoint_dir)
 
   if checkpoint_type == FLUX_CHECKPOINT:
-    item_names = ("flux_state", "flux_config",
-                  "vae_state", "vae_config",
-                  "scheduler", "scheduler_config")
+    item_names = ("flux_state", "flux_config", "vae_state", "vae_config", "scheduler", "scheduler_config")
   else:
     item_names = (
         "unet_config",

src/maxdiffusion/checkpointing/flux_checkpointer.py

@@ -18,9 +18,7 @@
 from contextlib import nullcontext
 import functools
 import json
-import os
 import jax
-import jax.numpy as jnp
 from jax.sharding import Mesh
 import orbax.checkpoint as ocp
 import grain.python as grain
@@ -32,11 +30,9 @@
 from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
 from ..pipelines.flux.flux_pipeline import FluxPipeline
 
-from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer)
+from transformers import (CLIPTokenizer, FlaxCLIPTextModel, FlaxT5EncoderModel, AutoTokenizer)
 
-from maxdiffusion.checkpointing.checkpointing_utils import (
-    create_orbax_checkpoint_manager
-)
+from maxdiffusion.checkpointing.checkpointing_utils import (create_orbax_checkpoint_manager)
 from maxdiffusion.models.flux.util import load_flow_model
 
 FLUX_CHECKPOINT = "FLUX_CHECKPOINT"
@@ -49,6 +45,7 @@
 VAE_STATE_KEY = "vae_state"
 VAE_STATE_SHARDINGS_KEY = "vae_state_shardings"
 
+
 class FluxCheckpointer(ABC):
 
   def __init__(self, config, checkpoint_type):
@@ -87,12 +84,14 @@ def create_flux_state(self, pipeline, params, checkpoint_item_name, is_training)
       tx, learning_rate_scheduler = self._create_optimizer(self.config, learning_rate)
 
     transformer_eval_params = transformer.init_weights(
-      rngs=self.rng, max_sequence_length=self.config.max_sequence_length, eval_only=True
+        rngs=self.rng, max_sequence_length=self.config.max_sequence_length, eval_only=True
     )
 
     transformer_params = load_flow_model(self.config.flux_name, transformer_eval_params, "cpu")
 
-    weights_init_fn = functools.partial(pipeline.flux.init_weights, rngs=self.rng, max_sequence_length=self.config.max_sequence_length)
+    weights_init_fn = functools.partial(
+        pipeline.flux.init_weights, rngs=self.rng, max_sequence_length=self.config.max_sequence_length
+    )
     flux_state, state_mesh_shardings = max_utils.setup_initial_state(
         model=pipeline.flux,
         tx=tx,
@@ -150,10 +149,11 @@ def _set_checkpoint_format(self, checkpoint_format):
   def save_checkpoint(self, train_step, pipeline, train_states):
     def config_to_json(model_or_config):
       return json.loads(model_or_config.to_json_string())
+
     items = {
         "flux_config": ocp.args.JsonSave(config_to_json(pipeline.flux)),
         "vae_config": ocp.args.JsonSave(config_to_json(pipeline.vae)),
-        "scheduler_config": ocp.args.JsonSave(config_to_json(pipeline.scheduler))
+        "scheduler_config": ocp.args.JsonSave(config_to_json(pipeline.scheduler)),
     }
 
     items[FLUX_STATE_KEY] = ocp.args.PyTreeSave(train_states[FLUX_STATE_KEY])
@@ -165,7 +165,7 @@ def config_to_json(model_or_config):
   def load_params(self, step=None):
 
     self.checkpoint_format = _CHECKPOINT_FORMAT_ORBAX
-  
+
   def load_flux_configs_from_orbax(self, step):
     max_logging.log("Restoring stable diffusion configs")
     if step is None:
@@ -188,68 +188,57 @@ def load_diffusers_checkpoint(self):
       context = jax.default_device(jax.devices("cpu")[0])
     else:
       context = nullcontext()
-    
+
     with context:
-      clip_encoder = FlaxCLIPTextModel.from_pretrained(
-        self.config.clip_model_name_or_path, dtype=self.config.weights_dtype
-      )
-      clip_tokenizer = CLIPTokenizer.from_pretrained(
-        self.config.clip_model_name_or_path,
-        max_length=77,
-        use_fast=True
-      )
+      clip_encoder = FlaxCLIPTextModel.from_pretrained(self.config.clip_model_name_or_path, dtype=self.config.weights_dtype)
+      clip_tokenizer = CLIPTokenizer.from_pretrained(self.config.clip_model_name_or_path, max_length=77, use_fast=True)
       t5_encoder = FlaxT5EncoderModel.from_pretrained(self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype)
       t5_tokenizer = AutoTokenizer.from_pretrained(
-        self.config.t5xxl_model_name_or_path,
-        max_length=self.config.max_sequence_length,
-        use_fast=True
+          self.config.t5xxl_model_name_or_path, max_length=self.config.max_sequence_length, use_fast=True
       )
 
       vae, vae_params = FlaxAutoencoderKL.from_pretrained(
-        self.config.pretrained_model_name_or_path,
-        subfolder="vae",
-        from_pt=True,
-        use_safetensors=True,
-        dtype=self.config.weights_dtype
+          self.config.pretrained_model_name_or_path,
+          subfolder="vae",
+          from_pt=True,
+          use_safetensors=True,
+          dtype=self.config.weights_dtype,
       )
 
       # loading from pretrained here causes a crash when trying to compile the model
       # Failed to load HSACO: HIP_ERROR_NoBinaryForGpu
       transformer = FluxTransformer2DModel.from_config(
-        self.config.pretrained_model_name_or_path,
-        subfolder="transformer",
-        mesh=self.mesh,
-        split_head_dim=self.config.split_head_dim,
-        attention_kernel=self.config.attention,
-        flash_block_sizes=flash_block_sizes,
-        dtype=self.config.activations_dtype,
-        weights_dtype=self.config.weights_dtype,
-        precision=max_utils.get_precision(self.config),
+          self.config.pretrained_model_name_or_path,
+          subfolder="transformer",
+          mesh=self.mesh,
+          split_head_dim=self.config.split_head_dim,
+          attention_kernel=self.config.attention,
+          flash_block_sizes=flash_block_sizes,
+          dtype=self.config.activations_dtype,
+          weights_dtype=self.config.weights_dtype,
+          precision=max_utils.get_precision(self.config),
       )
       transformer_eval_params = transformer.init_weights(
-        rngs=self.rng, max_sequence_length=self.config.max_sequence_length, eval_only=True
+          rngs=self.rng, max_sequence_length=self.config.max_sequence_length, eval_only=True
       )
-      
+
       transformer_params = load_flow_model(self.config.flux_name, transformer_eval_params, "cpu")
 
     pipeline = FluxPipeline(
-      t5_encoder,
-      clip_encoder,
-      vae,
-      t5_tokenizer,
-      clip_tokenizer,
-      transformer,
-      None,
-      dtype=self.config.activations_dtype,
-      mesh=self.mesh,
-      config=self.config,
-      rng=self.rng
+        t5_encoder,
+        clip_encoder,
+        vae,
+        t5_tokenizer,
+        clip_tokenizer,
+        transformer,
+        None,
+        dtype=self.config.activations_dtype,
+        mesh=self.mesh,
+        config=self.config,
+        rng=self.rng,
     )
 
-    params = {
-      FLUX_VAE_PARAMS_KEY : vae_params,
-      FLUX_TRANSFORMER_PARAMS_KEY : transformer_params
-    }
+    params = {FLUX_VAE_PARAMS_KEY: vae_params, FLUX_TRANSFORMER_PARAMS_KEY: transformer_params}
 
     return pipeline, params
 
@@ -267,55 +256,50 @@ def load_checkpoint(self, step=None, scheduler_class=None):
 
       with context:
         clip_encoder = FlaxCLIPTextModel.from_pretrained(
-          self.config.clip_model_name_or_path, dtype=self.config.weights_dtype
+            self.config.clip_model_name_or_path, dtype=self.config.weights_dtype
         )
-        clip_tokenizer = CLIPTokenizer.from_pretrained(
-          self.config.clip_model_name_or_path,
-          max_length=77,
-          use_fast=True
+        clip_tokenizer = CLIPTokenizer.from_pretrained(self.config.clip_model_name_or_path, max_length=77, use_fast=True)
+        t5_encoder = FlaxT5EncoderModel.from_pretrained(
+            self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype
         )
-        t5_encoder = FlaxT5EncoderModel.from_pretrained(self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype)
         t5_tokenizer = AutoTokenizer.from_pretrained(
-          self.config.t5xxl_model_name_or_path,
-          max_length=self.config.max_sequence_length,
-          use_fast=True
+            self.config.t5xxl_model_name_or_path, max_length=self.config.max_sequence_length, use_fast=True
         )
 
         vae = FlaxAutoencoderKL.from_config(
-          model_configs[0]["vae_config"],
-          dtype=self.config.activations_dtype,
-          weights_dtype=self.config.weights_dtype,
-          from_pt=self.config.from_pt,
+            model_configs[0]["vae_config"],
+            dtype=self.config.activations_dtype,
+            weights_dtype=self.config.weights_dtype,
+            from_pt=self.config.from_pt,
         )
 
         transformer = FluxTransformer2DModel.from_config(
-          model_configs[0]["flux_config"],
-          mesh=self.mesh,
-          split_head_dim=self.config.split_head_dim,
-          attention_kernel=self.config.attention,
-          flash_block_sizes=max_utils.get_flash_block_sizes(self.config),
-          dtype=self.config.activations_dtype,
-          weights_dtype=self.config.weights_dtype,
-          precision=max_utils.get_precision(self.config),
-          from_pt=self.config.from_pt,
+            model_configs[0]["flux_config"],
+            mesh=self.mesh,
+            split_head_dim=self.config.split_head_dim,
+            attention_kernel=self.config.attention,
+            flash_block_sizes=max_utils.get_flash_block_sizes(self.config),
+            dtype=self.config.activations_dtype,
+            weights_dtype=self.config.weights_dtype,
+            precision=max_utils.get_precision(self.config),
+            from_pt=self.config.from_pt,
         )
 
         pipeline = FluxPipeline(
-          t5_encoder,
-          clip_encoder,
-          vae,
-          t5_tokenizer,
-          clip_tokenizer,
-          transformer,
-          None,
-          dtype=self.config.activations_dtype,
-          mesh=self.mesh,
-          config=self.config,
-          rng=self.rng
+            t5_encoder,
+            clip_encoder,
+            vae,
+            t5_tokenizer,
+            clip_tokenizer,
+            transformer,
+            None,
+            dtype=self.config.activations_dtype,
+            mesh=self.mesh,
+            config=self.config,
+            rng=self.rng,
         )
 
     else:
       pipeline, params = self.load_diffusers_checkpoint()
-    
-    return pipeline, params
 
+    return pipeline, params

src/maxdiffusion/generate_flux_pipeline.py

@@ -28,8 +28,10 @@
 from maxdiffusion.checkpointing.checkpointing_utils import load_params_from_path
 from maxdiffusion.max_utils import setup_initial_state
 
+
 def run(config):
   from maxdiffusion.checkpointing.flux_checkpointer import FluxCheckpointer
+
   checkpoint_loader = FluxCheckpointer(config, "FLUX_CHECKPOINT")
   pipeline, params = checkpoint_loader.load_checkpoint()
 
@@ -47,9 +49,9 @@ def run(config):
     vae_state = {"params": vae_params}
 
     ## Flux
-    weights_init_fn = functools.partial(pipeline.flux.init_weights,
-                                        rngs=checkpoint_loader.rng,
-                                        max_sequence_length=config.max_sequence_length)
+    weights_init_fn = functools.partial(
+        pipeline.flux.init_weights, rngs=checkpoint_loader.rng, max_sequence_length=config.max_sequence_length
+    )
 
     unboxed_abstract_state, _, _ = max_utils.get_abstract_state(
         pipeline.flux, None, config, checkpoint_loader.mesh, weights_init_fn, False
@@ -61,10 +63,10 @@ def run(config):
     flux_state = {"params": flux_params}
   else:
     weights_init_fn = functools.partial(
-      pipeline.flux.init_weights,
-      rngs=checkpoint_loader.rng,
-      max_sequence_length=config.max_sequence_length,
-      eval_only=False
+        pipeline.flux.init_weights,
+        rngs=checkpoint_loader.rng,
+        max_sequence_length=config.max_sequence_length,
+        eval_only=False,
     )
     transformer_state, flux_state_shardings = setup_initial_state(
         model=pipeline.flux,
@@ -85,26 +87,22 @@ def run(config):
         config=config,
         mesh=checkpoint_loader.mesh,
         weights_init_fn=weights_init_fn,
-        model_params=params['flux_vae'],
+        model_params=params["flux_vae"],
         training=False,
     )
 
     vae_state = {"params": vae_state.params}
     flux_state = {"params": transformer_state.params}
 
   t0 = time.perf_counter()
-  with ExitStack() as stack:
-    imgs = pipeline(flux_params=flux_state,
-                    timesteps=50,
-                    vae_params=vae_state).block_until_ready()
+  with ExitStack():
+    imgs = pipeline(flux_params=flux_state, timesteps=50, vae_params=vae_state).block_until_ready()
   t1 = time.perf_counter()
   max_logging.log(f"Compile time: {t1 - t0:.1f}s.")
 
   t0 = time.perf_counter()
-  with ExitStack() as stack:
-    imgs = pipeline(flux_params=flux_state,
-                    timesteps=50,
-                    vae_params=vae_state).block_until_ready()
+  with ExitStack():
+    imgs = pipeline(flux_params=flux_state, timesteps=50, vae_params=vae_state).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/maxdiffusion_utils.py

@@ -255,19 +255,20 @@ def calculate_unet_tflops(config, pipeline, batch_size, rngs, train):
       / jax.local_device_count()
   )
 
+
 def get_dummy_flux_inputs(config, pipeline, batch_size):
   """Returns randomly initialized flux inputs."""
   latents, latents_ids = pipeline.prepare_latents(
-    batch_size=batch_size,
-    num_channels_latents=pipeline.flux.in_channels // 4,
-    height=config.resolution,
-    width=config.resolution,
-    vae_scale_factor=pipeline.vae_scale_factor,
-    dtype=config.activations_dtype,
-    rng=pipeline.rng
+      batch_size=batch_size,
+      num_channels_latents=pipeline.flux.in_channels // 4,
+      height=config.resolution,
+      width=config.resolution,
+      vae_scale_factor=pipeline.vae_scale_factor,
+      dtype=config.activations_dtype,
+      rng=pipeline.rng,
   )
   guidance_vec = jnp.asarray([config.guidance_scale] * batch_size, dtype=config.activations_dtype)
-  
+
   timesteps = jnp.ones((batch_size,), dtype=config.weights_dtype)
   t5_hidden_states_shape = (
       batch_size,
@@ -282,7 +283,7 @@ def get_dummy_flux_inputs(config, pipeline, batch_size):
       768,
   )
   clip_hidden_states = jnp.zeros(clip_hidden_states_shape, dtype=config.weights_dtype)
-  
+
   return (latents, timesteps, latents_ids, guidance_vec, t5_hidden_states, t5_ids, clip_hidden_states)
 
 
@@ -293,7 +294,9 @@ def calculate_flux_tflops(config, pipeline, batch_size, rngs, train):
   cache the compilation when flash is enabled.
   """
 
-  (latents, timesteps, latents_ids, guidance_vec, t5_hidden_states, t5_ids, clip_hidden_states) = get_dummy_flux_inputs(config, pipeline, batch_size)
+  (latents, timesteps, latents_ids, guidance_vec, t5_hidden_states, t5_ids, clip_hidden_states) = get_dummy_flux_inputs(
+      config, pipeline, batch_size
+  )
   return (
       max_utils.calculate_model_tflops(
           pipeline.flux,

src/maxdiffusion/pipelines/flux/__init__.py

@@ -1,5 +1,5 @@
-_import_structure = { "pipeline_jflux" : "JfluxPipeline" }
+_import_structure = {"pipeline_jflux": "JfluxPipeline"}
 
 from .flux_pipeline import (
     FluxPipeline,
-)
+)