vocoder weights

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -369,14 +369,18 @@ def rename_for_ltx2_vocoder(key):
 
 
 def load_vocoder_weights(
-    pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True, subfolder: str = "vocoder"
+    pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True, subfolder: str = "vocoder", filename: str = None
 ):
-  tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device)
+  tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device, filename=filename)
 
   flax_state_dict = {}
   cpu = jax.local_devices(backend="cpu")[0]
 
   for pt_key, tensor in tensors.items():
+    if filename and not pt_key.startswith("vocoder."):
+      continue
+    if filename and pt_key.startswith("vocoder."):
+      pt_key = pt_key[len("vocoder."):]
     key = rename_for_ltx2_vocoder(pt_key)
     parts = key.split(".")
 

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -537,33 +537,35 @@ def load_vocoder(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, confi
     max_logging.log("Loading Vocoder...")
 
     def create_model(rngs: nnx.Rngs, config: HyperParameters):
+      vocoder_repo = "Lightricks/LTX-2" if getattr(config, "model_name", "") == "ltx2.3" else config.pretrained_model_name_or_path
       if getattr(config, "model_name", "") == "ltx2.3":
         vocoder_class = LTX2VocoderWithBWE
       else:
         vocoder_class = LTX2Vocoder
-
+        
       vocoder = vocoder_class.from_config(
-          config.pretrained_model_name_or_path,
+          vocoder_repo,
           subfolder="vocoder",
           rngs=rngs,
           mesh=mesh,
           dtype=jnp.float32,
           weights_dtype=config.weights_dtype if hasattr(config, "weights_dtype") else jnp.float32,
       )
       return vocoder
-
+ 
     p_model_factory = partial(create_model, config=config)
     vocoder = nnx.eval_shape(p_model_factory, rngs=rngs)
     graphdef, state, rest_of_state = nnx.split(vocoder, nnx.Param, ...)
     rest_of_state = jax.tree_util.tree_map(cls._init_dummy_shape, rest_of_state)
-
+ 
     logical_state_spec = nnx.get_partition_spec(state)
     logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
     logical_state_sharding = dict(nnx.to_flat_state(logical_state_sharding))
     params = state.to_pure_dict()
     state = dict(nnx.to_flat_state(state))
-
-    params = load_vocoder_weights(config.pretrained_model_name_or_path, params, "cpu", subfolder="vocoder")
+ 
+    filename = "ltx-2.3-22b-dev.safetensors" if getattr(config, "model_name", "") == "ltx2.3" else None
+    params = load_vocoder_weights(config.pretrained_model_name_or_path, params, "cpu", subfolder="vocoder", filename=filename)
     if hasattr(config, "weights_dtype"):
       params = jax.tree_util.tree_map(lambda x: x.astype(config.weights_dtype), params)