Fixed comments and rebased on main

Author: ksikiric

src/maxdiffusion/checkpointing/flux_checkpointer.py

@@ -61,10 +61,8 @@ def __init__(self, config, checkpoint_type):
     self.mesh = Mesh(self.devices_array, self.config.mesh_axes)
     self.total_train_batch_size = self.config.total_train_batch_size
 
-    checkpoint_dir = os.path.abspath(self.config.checkpoint_dir)
-
     self.checkpoint_manager = create_orbax_checkpoint_manager(
-        checkpoint_dir,
+        self.config.checkpoint_dir,
         enable_checkpointing=True,
         save_interval_steps=1,
         checkpoint_type=checkpoint_type,

src/maxdiffusion/configs/base_flux_dev.yml

@@ -75,7 +75,7 @@ flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
 
-# If train_new_unet, unet weights will be randomly initialized to train the unet from scratch
+# If train_new_flux, flux weights will be randomly initialized to train flux from scratch
 # else they will be loaded from pretrained_model_name_or_path
 train_new_flux: False
 
@@ -223,8 +223,8 @@ skip_first_n_steps_for_profiler: 5
 profiler_steps: 10
 
 # Generation parameters
-prompt: "A Cubone, the lonely Pokémon, sits clutching its signature bone, its face hidden by a skull helmet."
-prompt_2: "A Cubone, the lonely Pokémon, sits clutching its signature bone, its face hidden by a skull helmet."
+prompt: "A magical castle in the middle of a forest, artistic drawing"
+prompt_2: "A magical castle in the middle of a forest, artistic drawing"
 negative_prompt: "purple, red"
 do_classifier_free_guidance: True
 guidance_scale: 3.5

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -83,9 +83,9 @@ flash_block_sizes: {
 # GroupNorm groups
 norm_num_groups: 32
 
-# If train_new_unet, unet weights will be randomly initialized to train the unet from scratch
+# If train_new_flux, flux weights will be randomly initialized to train flux from scratch
 # else they will be loaded from pretrained_model_name_or_path
-train_new_unet: False
+train_new_flux: False
 
 # train text_encoder - Currently not supported for SDXL
 train_text_encoder: False
@@ -123,7 +123,7 @@ diffusion_scheduler_config: {
 base_output_directory: ""
 
 # Hardware
-hardware: 'gpu' # Supported hardware types are 'tpu', 'gpu'
+hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
 
 # Parallelism
 mesh_axes: ['data', 'fsdp', 'tensor']
@@ -238,7 +238,7 @@ do_classifier_free_guidance: True
 guidance_scale: 0.0
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0
-num_inference_steps: 50
+num_inference_steps: 4
 
 # SDXL Lightning parameters
 lightning_from_pt: True

src/maxdiffusion/schedulers/scheduling_euler_discrete_flax.py

@@ -17,7 +17,6 @@
 
 import flax
 import jax.numpy as jnp
-import max_logging
 
 from ..configuration_utils import ConfigMixin, register_to_config
 from .scheduling_utils_flax import (
@@ -170,7 +169,6 @@ def set_timesteps(
       timesteps = (jnp.arange(self.config.num_train_timesteps, 0, -step_ratio)).round()
       timesteps -= 1
     elif timestep_spacing == "flux":
-      max_logging.log("Using flux timestep spacing")
       timesteps = jnp.linspace(1, 0, num_inference_steps + 1)
     else:
       raise ValueError(