Refactor Wan model structure by renaming and relocating model imports from model.py to wan2.py, enhancing code organization and clarity across the Wan2 module.

mlx_video/models/ltx_2/__init__.py

@@ -4,4 +4,4 @@ from mlx_video.models.ltx_2.config import (
     LTXModelType,
     TransformerConfig,
 )
-from mlx_video.models.ltx_2.ltx import LTXModel, X0Model
+from mlx_video.models.ltx_2.ltx_2 import LTXModel, X0Model

mlx_video/models/ltx_2/generate.py

@@ -33,7 +33,7 @@ from mlx_video.models.ltx_2.conditioning import (
     apply_conditioning,
 )
 from mlx_video.models.ltx_2.conditioning.latent import LatentState, apply_denoise_mask
-from mlx_video.models.ltx_2.ltx import LTXModel
+from mlx_video.models.ltx_2.ltx_2 import LTXModel
 from mlx_video.models.ltx_2.transformer import Modality
 from mlx_video.models.ltx_2.upsampler import load_upsampler, upsample_latents
 from mlx_video.models.ltx_2.video_vae import VideoEncoder

mlx_video/models/wan2/__init__.py

@@ -1,2 +1,2 @@
 from mlx_video.models.wan2.config import WanModelConfig
-from mlx_video.models.wan2.model import WanModel
+from mlx_video.models.wan2.wan2 import WanModel

mlx_video/models/wan2/convert.py

@@ -594,7 +594,7 @@ def _quantize_saved_model(
 
     import mlx.nn as nn
 
-    from mlx_video.models.wan2.model import WanModel
+    from mlx_video.models.wan2.wan2 import WanModel
 
     if source_dir is None:
         source_dir = output_dir

mlx_video/models/wan2/utils.py

@@ -21,7 +21,7 @@ def load_wan_model(
                       If provided, creates QuantizedLinear stubs before loading.
         loras: Optional list of (lora_path, strength) tuples to apply.
     """
-    from mlx_video.models.wan2.model import WanModel
+    from mlx_video.models.wan2.wan2 import WanModel
 
     model = WanModel(config)