convert pt weights to flax and load transformer state.

Author: jfacevedo-google

src/maxdiffusion/generate_flux.py

@@ -38,7 +38,8 @@
   get_memory_allocations,
   create_device_mesh,
   get_flash_block_sizes,
-  get_precision
+  get_precision,
+  setup_initial_state
 )
 
 def prepare_latent_image_ids(height, width):
@@ -195,7 +196,7 @@ def run(config):
   devices_array = create_device_mesh(config)
   mesh = Mesh(devices_array, config.mesh_axes)
 
-  per_host_number_of_images = 1#config.per_device_batch_size * jax.local_device_count()
+  per_host_number_of_images = config.per_device_batch_size * jax.local_device_count()
 
   # LOAD VAE
 
@@ -233,16 +234,6 @@ def run(config):
     rng=rng
   )
 
-  #load_flow_model("flux-dev", "cpu")
-
-  # transformer, params = FluxTransformer2DModel.from_pretrained(
-  #   config.pretrained_model_name_or_path,
-  #   subfolder="text_encoder_2",
-  #   from_pt=True,
-  #   dtype=config.weights_dtype
-  # )
-
-
   # LOAD TEXT ENCODERS - t5 on cpu
   clip_text_encoder = FlaxCLIPTextModel.from_pretrained(
     config.pretrained_model_name_or_path,
@@ -303,17 +294,35 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
     pooled_prompt_embeds
   )
   get_memory_allocations()
-  transformer_params = transformer.init_weights(rng, True)
-  # transformer_params = transformer.init(
-  #   {"params" : rng},
-  #   img=latents,
-  #   img_ids=latent_image_ids,
-  #   txt=prompt_embeds,
-  #   txt_ids=text_ids,
-  #   timesteps=timesteps,
-  #   guidance=guidance,
-  #   y=pooled_prompt_embeds
-  # )["params"]
+  # evaluate shapes
+  transformer_eval_params = transformer.init_weights(rngs=rng, max_sequence_length=512, eval_only=True)
+  
+  # loads pretrained weights
+  transformer_params = load_flow_model("flux-dev", transformer_eval_params, "cpu")
+  get_memory_allocations()
+  # create transformer state
+  weights_init_fn = functools.partial(transformer.init_weights, rngs=rng, max_sequence_length=512, eval_only=False)
+  transformer_state, transformer_state_shardings = setup_initial_state(
+    model=transformer,
+    tx=None,
+    config=config,
+    mesh=mesh,
+    weights_init_fn=weights_init_fn,
+    model_params=None,
+    training=False
+  )
+  breakpoint()
+  transformer_state = transformer_state.replace(params=transformer_params)
+  img = transformer.apply(
+    {"params" : transformer_state.params},
+    img=latents,
+    img_ids=latent_image_ids,
+    txt=prompt_embeds,
+    txt_ids=text_ids,
+    timesteps=timesteps,
+    guidance=guidance,
+    y=pooled_prompt_embeds
+  )
   get_memory_allocations()
   breakpoint()
 

src/maxdiffusion/models/flux/modules/layers.py

@@ -53,11 +53,13 @@ class QKNorm(nn.Module):
   def __call__(self, q: Array, k: Array, v: Array) -> tuple[Array, Array]:
     q = nn.RMSNorm(
       dtype=self.dtype,
-      param_dtype=self.weights_dtype
+      param_dtype=self.weights_dtype,
+      name="query_norm"
     )(q)
     k = nn.RMSNorm(
       dtype=self.dtype,
-      param_dtype=self.weights_dtype
+      param_dtype=self.weights_dtype,
+      name="key_norm"
     )(k)
     return q, k
 
@@ -173,7 +175,8 @@ def __call__(self, vec: Array) -> tuple[ModulationOut, ModulationOut | None]:
       kernel_init=nn.with_logical_partitioning(
         nn.initializers.lecun_normal(),
         ("embed", "heads")
-      )
+      ),
+      name="lin"
     )(nn.silu(vec))
 
     out = jnp.split(lin[:, None, :], multiplier, axis=-1)
@@ -205,14 +208,16 @@ def __call__(self, x: Array, vec: Array, pe: Array) -> Array:
       double=False,
       dtype=self.dtype,
       weights_dtype=self.weights_dtype,
-      precision=self.precision
+      precision=self.precision,
+      name="modulation"
     )(vec)
     x_mod = (1 + mod.scale) * nn.LayerNorm(
       use_scale=False,
       use_bias=False,
       epsilon=1e-6,
       dtype=self.dtype,
-      param_dtype=self.weights_dtype
+      param_dtype=self.weights_dtype,
+      name="pre_norm"
     )(x) + mod.shift
 
     x_mod = nn.Dense(
@@ -231,7 +236,8 @@ def __call__(self, x: Array, vec: Array, pe: Array) -> Array:
     q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
     q, k = QKNorm(
       dtype=self.dtype,
-      weights_dtype=self.weights_dtype
+      weights_dtype=self.weights_dtype,
+      name="norm"
     )(q, k, v)
 
     q, k = apply_rope(q, k, pe)
@@ -286,15 +292,17 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
       double=True,
       dtype=self.dtype,
       weights_dtype=self.weights_dtype,
-      precision=self.precision
+      precision=self.precision,
+      name="img_mod"
     )(vec)
 
     txt_mod1, txt_mod2 = Modulation(
       self.hidden_size,
       double=True,
       dtype=self.dtype,
       weights_dtype=self.weights_dtype,
-      precision=self.precision
+      precision=self.precision,
+      name="txt_mod"
     )(vec)
 
     # prepare image for attention
@@ -303,7 +311,8 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
       use_bias=False,
       epsilon=1e-6,
       dtype=self.dtype,
-      param_dtype=self.weights_dtype
+      param_dtype=self.weights_dtype,
+      name="img_norm1"
     )(img)
     img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
     img_qkv = nn.Dense(
@@ -315,14 +324,16 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
       kernel_init=nn.with_logical_partitioning(
         nn.initializers.lecun_normal(),
         ("embed", "heads")
-      )
+      ),
+      name="img_attn_qkv"
     )(img_modulated)
     img_q, img_k, img_v = rearrange(
       img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
     )
     img_q, img_k = QKNorm(
       dtype=self.dtype,
-      weights_dtype=self.weights_dtype
+      weights_dtype=self.weights_dtype,
+      name="img_attn_norm"
     )(img_q, img_k, img_v)
 
     # prepare text for attention
@@ -331,7 +342,8 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
       use_bias=False,
       epsilon=1e-6,
       dtype=self.dtype,
-      param_dtype=self.weights_dtype
+      param_dtype=self.weights_dtype,
+      name="txt_norm1"
     )(txt)
     txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
     txt_qkv = nn.Dense(
@@ -343,14 +355,16 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
       kernel_init=nn.with_logical_partitioning(
         nn.initializers.lecun_normal(),
         ("embed", "heads")
-      )
+      ),
+      name="txt_attn_qkv"
     )(txt_modulated)
     txt_q, txt_k, txt_v = rearrange(
       txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
     )
     txt_q, txt_k = QKNorm(
       dtype=self.dtype,
-      weights_dtype=self.weights_dtype
+      weights_dtype=self.weights_dtype,
+      name="txt_attn_norm"
     )(txt_q, txt_k, txt_v)
 
     # run actual attention
@@ -385,6 +399,7 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
         nn.initializers.lecun_normal(),
         ("heads", "embed")
       ),
+      name="img_attn_proj"
     )(img_attn)
     img = img + img_mod2.gate * nn.Sequential(
       [
@@ -397,7 +412,8 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
           kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(),
             ("embed", "heads")
-          )
+          ),
+          name="img_mlp_0"
         ),
         nn.gelu,
         nn.Dense(
@@ -408,14 +424,16 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
           kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(),
             ("heads", "embed")
-          )
-        )
-      ]
+          ),
+          name="img_mlp_2"
+        ),
+      ],
     )(
       (1 + img_mod2.scale) * nn.LayerNorm(
         use_scale=False,
         use_bias=False,
-        param_dtype=self.weights_dtype
+        param_dtype=self.weights_dtype,
+        name="img_norm2"
       )(img) + img_mod2.shift
     )
 
@@ -430,6 +448,7 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
         nn.initializers.lecun_normal(),
         ("heads", "embed")
       ),
+      name="txt_attn_proj"
     )(txt_attn)
     txt = txt + txt_mod1.gate * txt_proj
 
@@ -444,7 +463,8 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
           kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(),
             ("embed", "heads")
-          )
+          ),
+          name="txt_mlp_0"
         ),
         nn.gelu,
         nn.Dense(
@@ -455,14 +475,16 @@ def __call__(self, img: Array, txt: Array, vec: Array, pe: Array) -> tuple[Array
           kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(),
             ("heads", "embed")
-          )
-        )
-      ]
+          ),
+          name="txt_mlp_2"
+        ),
+      ],
     )(
       (1 + txt_mod2.scale) * nn.LayerNorm(
         use_scale=False,
         use_bias=False,
-        param_dtype=self.weights_dtype
+        param_dtype=self.weights_dtype,
+        name="txt_norm2"
       )(txt) + txt_mod2.shift
     )
 
@@ -491,8 +513,9 @@ def __call__(self, x: Array, vec: Array) -> Array:
             kernel_init=nn.with_logical_partitioning(
               nn.initializers.lecun_normal(),
               ("embed", "heads")
-            )
-          )
+            ),
+            name="adaLN_modulation_1"
+          ),
         ]
       )(vec), 2, axis=1
     )
@@ -515,5 +538,5 @@ def __call__(self, x: Array, vec: Array) -> Array:
         ("heads", "embed")
       ),
       name="linear"
-    )
+    )(x)
     return x