Merge pull request #4 from borisdayma/fix-flax

fix: handle flax 0.3.6

Author: patil-suraj

vqgan_jax/convert_pt_model_to_jax.py

@@ -50,17 +50,13 @@ def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model):
 
     # Correctly rename weight parameters
     if ("norm" in pt_key and (pt_tuple_key[-1] == "bias")
-        and (pt_tuple_key[:-1] + ("bias", ) in random_flax_state_dict)):
-      pt_tensor = pt_tensor[None, None, None, :]
-    elif ("norm" in pt_key and (pt_tuple_key[-1] == "bias")
+          and (pt_tuple_key[:-1] + ("bias", ) not in random_flax_state_dict)
           and (pt_tuple_key[:-1] + ("scale", ) in random_flax_state_dict)):
       pt_tuple_key = pt_tuple_key[:-1] + ("scale", )
-      pt_tensor = pt_tensor[None, None, None, :]
     elif pt_tuple_key[-1] in [
         "weight", "gamma"
     ] and pt_tuple_key[:-1] + ("scale", ) in random_flax_state_dict:
       pt_tuple_key = pt_tuple_key[:-1] + ("scale", )
-      pt_tensor = pt_tensor[None, None, None, :]
     if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + (
         "embedding", ) in random_flax_state_dict:
       pt_tuple_key = pt_tuple_key[:-1] + ("embedding", )

vqgan_jax/modeling_flax_vqgan.py

@@ -399,8 +399,6 @@ def setup(self):
                                                     - 1]
     curr_res = self.config.resolution // 2**(self.config.num_resolutions - 1)
     self.z_shape = (1, self.config.z_channels, curr_res, curr_res)
-    print("Working with z of shape {} = {} dimensions.".format(
-        self.z_shape, np.prod(self.z_shape)))
 
     # z to block_in
     self.conv_in = nn.Conv(