e2e wan vae with weights loading. Still not fully working.

Author: jfacevedo-google

src/maxdiffusion/configuration_utils.py

@@ -464,7 +464,8 @@ def extract_init_dict(cls, config_dict, **kwargs):
     # remove flax internal keys
     if hasattr(cls, "_flax_internal_args"):
       for arg in cls._flax_internal_args:
-        expected_keys.remove(arg)
+        if arg in expected_keys:
+          expected_keys.remove(arg)
 
     # 2. Remove attributes that cannot be expected from expected config attributes
     # remove keys to be ignored

src/maxdiffusion/generate_wan.py

@@ -26,6 +26,7 @@
 from absl import app
 from transformers import AutoTokenizer, UMT5EncoderModel
 from maxdiffusion import pyconfig, max_logging
+from maxdiffusion.models.wan.autoencoder_kl_wan import AutoencoderKLWan
 from maxdiffusion.models.wan.transformers.transformer_flux_wan_nnx import WanModel
 from maxdiffusion.pipelines.wan.pipeline_wan import WanPipeline
 

src/maxdiffusion/models/flux/util.py

@@ -11,7 +11,11 @@
 from jax import numpy as jnp
 from safetensors import safe_open
 
-from maxdiffusion.models.modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor)
+from ..modeling_flax_pytorch_utils import (
+  rename_key,
+  rename_key_and_reshape_tensor,
+  torch2jax
+)
 from maxdiffusion import max_logging
 
 
@@ -32,21 +36,6 @@ class FluxParams:
   rngs: Array
   param_dtype: DTypeLike
 
-
-def torch2jax(torch_tensor: torch.Tensor) -> Array:
-  is_bfloat16 = torch_tensor.dtype == torch.bfloat16
-  if is_bfloat16:
-    # upcast the tensor to fp32
-    torch_tensor = torch_tensor.float()
-
-  if torch.device.type != "cpu":
-    torch_tensor = torch_tensor.to("cpu")
-
-  numpy_value = torch_tensor.numpy()
-  jax_array = jnp.array(numpy_value, dtype=jnp.bfloat16 if is_bfloat16 else None)
-  return jax_array
-
-
 @dataclass
 class ModelSpec:
   params: FluxParams

src/maxdiffusion/models/modeling_flax_pytorch_utils.py

@@ -15,18 +15,57 @@
 """ PyTorch - Flax general utilities."""
 import re
 
+import torch
 import jax
 import jax.numpy as jnp
 from flax.linen import Partitioned
 from flax.traverse_util import flatten_dict, unflatten_dict
 from flax.core.frozen_dict import unfreeze
 from jax.random import PRNGKey
-
+from chex import Array
 from ..utils import logging
+from .. import max_logging
 
 
 logger = logging.get_logger(__name__)
 
+def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
+  """
+  expected_pytree: dict - a pytree that comes from initializing the model.
+  new_pytree: dict - a pytree that has been created from pytorch weights.
+  """
+  expected_pytree = flatten_dict(expected_pytree)
+  if len(expected_pytree.keys()) != len(new_pytree.keys()):
+    set1 = set(expected_pytree.keys())
+    set2 = set(new_pytree.keys())
+    missing_keys = set1 ^ set2
+    max_logging.log(f"missing keys : {missing_keys}")
+  for key in expected_pytree.keys():
+    if key in new_pytree.keys():
+      try:
+        expected_pytree_shape = expected_pytree[key].shape
+      except Exception:
+        expected_pytree_shape = expected_pytree[key].value.shape
+      if expected_pytree_shape != new_pytree[key].shape:
+        max_logging.log(f"shape mismatch for {key}")
+        max_logging.log(
+            f"shape mismatch, expected shape of {expected_pytree[key].shape}, but got shape of {new_pytree[key].shape}"
+        )
+    else:
+      max_logging.log(f"key: {key} not found...")
+
+def torch2jax(torch_tensor: torch.Tensor) -> Array:
+  is_bfloat16 = torch_tensor.dtype == torch.bfloat16
+  if is_bfloat16:
+    # upcast the tensor to fp32
+    torch_tensor = torch_tensor.float()
+
+  if torch.device.type != "cpu":
+    torch_tensor = torch_tensor.to("cpu")
+
+  numpy_value = torch_tensor.numpy()
+  jax_array = jnp.array(numpy_value, dtype=jnp.bfloat16 if is_bfloat16 else None)
+  return jax_array
 
 def rename_key(key):
   regex = r"\w+[.]\d+"
@@ -93,6 +132,12 @@ def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dic
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
     pt_tensor = pt_tensor.transpose(2, 3, 1, 0)
     return renamed_pt_tuple_key, pt_tensor
+  
+  # 3d conv layer
+  renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
+  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 5:
+    pt_tensor = pt_tensor.transpose(2, 3, 4, 1, 0)
+    return renamed_pt_tuple_key, pt_tensor
 
   # linear layer
   renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
@@ -103,6 +148,8 @@ def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dic
   # old PyTorch layer norm weight
   renamed_pt_tuple_key = pt_tuple_key[:-1] + ("weight",)
   if pt_tuple_key[-1] == "gamma":
+    renamed_pt_tuple_key = pt_tuple_key
+    pt_tensor = pt_tensor.flatten()
     return renamed_pt_tuple_key, pt_tensor
 
   # old PyTorch layer norm bias

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -200,8 +200,6 @@ def __init__(
       precision: jax.lax.Precision = None,
       attention: str = "dot_product",
   ):
-    kernel_size = _canonicalize_tuple(kernel_size, 3, "kernel_size")
-    stride = _canonicalize_tuple(stride, 3, "stride")
     self.conv = nnx.Conv(dim, dim, kernel_size=kernel_size, strides=stride, use_bias=True, rngs=rngs)
 
   def __call__(self, x):
@@ -233,19 +231,19 @@ def __init__(
           nnx.Conv(
               dim,
               dim // 2,
-              kernel_size=(1, 3, 3),
+              kernel_size=(3, 3),
               padding="SAME",
               use_bias=True,
               rngs=rngs,
           ),
       )
     elif mode == "upsample3d":
       self.resample = nnx.Sequential(
-          WanUpsample(scale_factor=(2.0, 2.0, 2.0), method="nearest"),
+          WanUpsample(scale_factor=(2.0, 2.0), method="nearest"),
           nnx.Conv(
               dim,
               dim // 2,
-              kernel_size=(1, 3, 3),
+              kernel_size=(3, 3),
               padding="SAME",
               use_bias=True,
               rngs=rngs,
@@ -259,11 +257,9 @@ def __init__(
           padding=(1, 0, 0),
       )
     elif mode == "downsample2d":
-      # TODO - do I need to transpose?
-      self.resample = ZeroPaddedConv2D(dim=dim, rngs=rngs, kernel_size=(1, 3, 3), stride=(1, 2, 2))
+      self.resample = ZeroPaddedConv2D(dim=dim, rngs=rngs, kernel_size=(3, 3), stride=(2, 2))
     elif mode == "downsample3d":
-      # TODO - do I need to transpose?
-      self.resample = ZeroPaddedConv2D(dim=dim, rngs=rngs, kernel_size=(1, 3, 3), stride=(1, 2, 2))
+      self.resample = ZeroPaddedConv2D(dim=dim, rngs=rngs, kernel_size=(3, 3), stride=(2, 2))
       self.time_conv = WanCausalConv3d(
           rngs=rngs, in_channels=dim, out_channels=dim, kernel_size=(3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)
       )
@@ -334,7 +330,6 @@ def __init__(
     self.norm1 = WanRMS_norm(dim=in_dim, rngs=rngs, images=False, channel_first=False)
     self.conv1 = WanCausalConv3d(rngs=rngs, in_channels=in_dim, out_channels=out_dim, kernel_size=3, padding=1)
     self.norm2 = WanRMS_norm(dim=out_dim, rngs=rngs, images=False, channel_first=False)
-    self.dropout = nnx.Dropout(dropout, rngs=rngs)
     self.conv2 = WanCausalConv3d(rngs=rngs, in_channels=out_dim, out_channels=out_dim, kernel_size=3, padding=1)
     self.conv_shortcut = (
         WanCausalConv3d(rngs=rngs, in_channels=in_dim, out_channels=out_dim, kernel_size=1)
@@ -363,7 +358,6 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=[0]):
 
     x = self.norm2(x)
     x = self.nonlinearity(x)
-    x = self.dropout(x)
 
     if feat_cache is not None:
       idx = feat_idx[0]
@@ -384,8 +378,8 @@ class WanAttentionBlock(nnx.Module):
   def __init__(self, dim: int, rngs: nnx.Rngs):
     self.dim = dim
     self.norm = WanRMS_norm(rngs=rngs, dim=dim, channel_first=False)
-    self.to_qkv = nnx.Conv(in_features=dim, out_features=dim * 3, kernel_size=1, rngs=rngs)
-    self.proj = nnx.Conv(in_features=dim, out_features=dim, kernel_size=1, rngs=rngs)
+    self.to_qkv = nnx.Conv(in_features=dim, out_features=dim * 3, kernel_size=(1, 1), rngs=rngs)
+    self.proj = nnx.Conv(in_features=dim, out_features=dim, kernel_size=(1, 1), rngs=rngs)
 
   def __call__(self, x: jax.Array):
     batch_size, time, height, width, channels = x.shape
@@ -801,8 +795,6 @@ def _encode(self, x: jax.Array):
       x = jnp.transpose(x, (0, 2, 3, 4, 1))
       assert x.shape[-1] == 3, f"Expected input shape (N, D, H, W, 3), got {x.shape}"
 
-    # self.clear_cache()
-
     t = x.shape[1]
     iter_ = 1 + (t - 1) // 4
     for i in range(iter_):
@@ -854,8 +846,8 @@ def _decode(self, z: jax.Array, return_dict: bool = True) -> Union[FlaxDecoderOu
   def decode(self, z: jax.Array, return_dict: bool = True) -> Union[FlaxDecoderOutput, jax.Array]:
     if z.shape[-1] != self.z_dim:
       # reshape channel last for JAX
-      x = jnp.transpose(x, (0, 2, 3, 4, 1))
-      assert x.shape[-1] == self.z_dim, f"Expected input shape (N, D, H, W, {self.z_dim}, got {x.shape}"
+      z = jnp.transpose(z, (0, 2, 3, 4, 1))
+      assert z.shape[-1] == self.z_dim, f"Expected input shape (N, D, H, W, {self.z_dim}, got {z.shape}"
     decoded = self._decode(z).sample
     if not return_dict:
       return (decoded,)

src/maxdiffusion/models/wan/wan_utils.py

@@ -0,0 +1,78 @@
+import jax
+import jax.numpy as jnp
+from maxdiffusion import max_logging
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from flax.traverse_util import flatten_dict, unflatten_dict
+from ..modeling_flax_pytorch_utils import (
+  rename_key,
+  rename_key_and_reshape_tensor,
+  torch2jax,
+  validate_flax_state_dict
+)
+
+def _tuple_str_to_int(in_tuple):
+  out_list = []
+  for item in in_tuple:
+    try:
+      out_list.append(int(item))
+    except:
+      out_list.append(item)
+  return tuple(out_list)
+
+
+def load_wan_vae(pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True):
+  device = jax.devices(device)[0]
+  with jax.default_device(device):
+    if hf_download:
+      ckpt_path = hf_hub_download(pretrained_model_name_or_path, subfolder="vae", filename="diffusion_pytorch_model.safetensors")
+    #breakpoint()
+    max_logging.log(f"Load and port Wan 2.1 VAE on {device}")
+
+    if ckpt_path is not None:
+      tensors = {}
+      with safe_open(ckpt_path, framework="pt") as f:
+        for k in f.keys():
+          tensors[k] = torch2jax(f.get_tensor(k))
+      flax_state_dict = {}
+      cpu = jax.local_devices(backend="cpu")[0]
+      for pt_key, tensor in tensors.items():
+        renamed_pt_key = rename_key(pt_key)
+        # Order matters
+        renamed_pt_key = renamed_pt_key.replace("up_blocks_", "up_blocks.")
+        renamed_pt_key = renamed_pt_key.replace("mid_block_", "mid_block.")
+        renamed_pt_key = renamed_pt_key.replace("down_blocks_", "down_blocks.")
+
+        renamed_pt_key = renamed_pt_key.replace("conv_in.bias", "conv_in.conv.bias")
+        renamed_pt_key = renamed_pt_key.replace("conv_in.weight", "conv_in.conv.weight")
+        renamed_pt_key = renamed_pt_key.replace("conv_out.bias", "conv_out.conv.bias")
+        renamed_pt_key = renamed_pt_key.replace("conv_out.weight", "conv_out.conv.weight")
+        renamed_pt_key = renamed_pt_key.replace("attentions_", "attentions.")
+        renamed_pt_key = renamed_pt_key.replace("resnets_", "resnets.")
+        renamed_pt_key = renamed_pt_key.replace("upsamplers_", "upsamplers.")
+        renamed_pt_key = renamed_pt_key.replace("resample_", "resample.")
+        renamed_pt_key = renamed_pt_key.replace("conv1.bias", "conv1.conv.bias")
+        renamed_pt_key = renamed_pt_key.replace("conv1.weight", "conv1.conv.weight")
+        renamed_pt_key = renamed_pt_key.replace("conv2.bias", "conv2.conv.bias")
+        renamed_pt_key = renamed_pt_key.replace("conv2.weight", "conv2.conv.weight")
+        renamed_pt_key = renamed_pt_key.replace("time_conv.bias", "time_conv.conv.bias")
+        renamed_pt_key = renamed_pt_key.replace("time_conv.weight", "time_conv.conv.weight")
+        renamed_pt_key = renamed_pt_key.replace("quant_conv", "quant_conv.conv")
+        renamed_pt_key = renamed_pt_key.replace("conv_shortcut", "conv_shortcut.conv")
+        if "decoder" in renamed_pt_key:
+          renamed_pt_key = renamed_pt_key.replace("resample.1.bias", "resample.layers.1.bias")
+          renamed_pt_key = renamed_pt_key.replace("resample.1.weight", "resample.layers.1.weight")
+        if "encoder" in renamed_pt_key:
+          renamed_pt_key = renamed_pt_key.replace("resample.1", "resample.conv")
+        pt_tuple_key = tuple(renamed_pt_key.split("."))
+        flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes)
+        flax_key = _tuple_str_to_int(flax_key)
+        flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)
+      validate_flax_state_dict(eval_shapes, flax_state_dict)
+      flax_state_dict = unflatten_dict(flax_state_dict)
+      del tensors
+      jax.clear_caches()
+    else:
+      raise FileNotFoundError(f"Path {ckpt_path} was not found")
+    
+    return flax_state_dict

src/maxdiffusion/tests/wan_vae_test.py

@@ -29,14 +29,14 @@
     WanCausalConv3d,
     WanUpsample,
     AutoencoderKLWan,
-    WanEncoder3d,
     WanMidBlock,
     WanResidualBlock,
     WanRMS_norm,
     WanResample,
     ZeroPaddedConv2D,
     WanAttentionBlock,
 )
+from ..models.wan.wan_utils import load_wan_vae
 
 CACHE_T = 2
 
@@ -421,6 +421,20 @@ def test_wan_encode(self):
     output = wan_vae.encode(input)
     assert output.latent_dist.sample(key).shape == (1, 13, 60, 90, 16)
 
+  # def test_load_checkpoint(self):
+  #   pretrained_model_name_or_path = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
+  #   key = jax.random.key(0)
+  #   rngs = nnx.Rngs(key)
+  #   wan_vae = AutoencoderKLWan.from_config(
+  #     pretrained_model_name_or_path,
+  #     subfolder="vae",
+  #     rngs=rngs
+  #   )
+  #   graphdef, state = nnx.split(wan_vae)
+  #   params = state.to_pure_dict()
+  #   # This replaces random params with the model.
+  #   params = load_wan_vae(pretrained_model_name_or_path, params, "cpu")
+
 
 if __name__ == "__main__":
   absltest.main()

src/maxdiffusion/utils/__init__.py

@@ -83,7 +83,7 @@
     is_xformers_available,
     requires_backends,
 )
-from .loading_utils import load_image
+from .loading_utils import load_image, load_video
 from .logging import get_logger
 from .outputs import BaseOutput
 from .peft_utils import (
@@ -103,7 +103,6 @@
     convert_unet_state_dict_to_peft,
 )
 
-
 logger = get_logger(__name__)