added model_name in config file

Author: prishajain1

src/maxdiffusion/configs/base_wan_14b.yml

@@ -28,6 +28,7 @@ save_config_to_gcs: False
 log_period: 100
 
 pretrained_model_name_or_path: 'Wan-AI/Wan2.1-T2V-14B-Diffusers'
+model_name: wan2.1
 
 # Overrides the transformer from pretrained_model_name_or_path
 wan_transformer_pretrained_model_name_or_path: ''

src/maxdiffusion/configs/base_wan_27b.yml

@@ -28,7 +28,7 @@ save_config_to_gcs: False
 log_period: 100
 
 pretrained_model_name_or_path: 'Wan-AI/Wan2.2-T2V-A14B-Diffusers'
-run_wan2_2: True
+model_name: wan2.2
 
 # Overrides the transformer from pretrained_model_name_or_path
 wan_transformer_pretrained_model_name_or_path: ''

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -213,7 +213,7 @@ def __init__(
     self.devices_array = devices_array
     self.mesh = mesh
     self.config = config
-    self.run_wan2_2 = config.run_wan2_2 if 'run_wan2_2' in self.config.__dict__ else False
+    self.run_wan2_2 = config.model_name == "wan2.2"
 
     self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
     self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
@@ -379,7 +379,7 @@ def from_checkpoint(cls, config: HyperParameters, restored_checkpoint=None, vae_
     mesh = Mesh(devices_array, config.mesh_axes)
     rng = jax.random.key(config.seed)
     rngs = nnx.Rngs(rng)
-    run_wan2_2 = config.run_wan2_2 if 'run_wan2_2' in config.__dict__ else False
+    run_wan2_2 = config.model_name == "wan2.2"
     low_noise_transformer = None
     high_noise_transformer = None
     tokenizer = None
@@ -421,7 +421,7 @@ def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transform
     mesh = Mesh(devices_array, config.mesh_axes)
     rng = jax.random.key(config.seed)
     rngs = nnx.Rngs(rng)
-    run_wan2_2 = config.run_wan2_2 if 'run_wan2_2' in config.__dict__ else False
+    run_wan2_2 = config.model_name == "wan2.2"
     low_noise_transformer = None
     high_noise_transformer = None
     tokenizer = None