remove redundant code, fix documentation

Author: Gleb Sterkin

video/wan2.1/README.md

@@ -5,7 +5,7 @@ Wan2.1 text-to-video and image-to-video implementation in MLX. The model
 weights are downloaded directly from the [Hugging Face
 Hub](https://huggingface.co/Wan-AI).
 
-| Model | Task | HF Repo | RAM (unquantized) | Single DiT step on M4 Pro chip |
+| Model | Task | HF Repo | RAM (unquantized), 81 frames | Single DiT step on M4 Max chip, 81 frames |
 |-------|------|---------|-----------------|---|
 | 1.3B | T2V | [Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) | ~10GB | ~100 s/it |
 | 14B | T2V | [Wan-AI/Wan2.1-T2V-14B](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) | ~36GB | ~230 s/it |
@@ -82,7 +82,7 @@ python txt2video.py 'A cat playing piano' --quantize --output out_quantized.mp4
 ```
 
 ### Disabling the cache
-To get additional memory savings at the expense of a bit of speed use `--no-cache` argument that will prevent MLX from utilizing the cache (sets `mx.set_cache_limit(0)` under the hood). See [documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.set_cache_limit.html) for more info
+To get additional memory savings at the expense of a bit of speed use `--no-cache` argument. It will prevent MLX from utilizing the cache (sets `mx.set_cache_limit(0)` under the hood). See [documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.set_cache_limit.html) for more info
 ```shell
 python txt2video.py 'A cat playing piano' --output out.mp4 --no-cache
 ```
@@ -107,15 +107,15 @@ python txt2video.py 'A cat playing piano' \
     --quantize --output out_t2v_distilled.mp4
 ```
 
-For image to video pipeline we use [4 steps distilled i2v model](https://huggingface.co/lightx2v/Wan2.1-Distill-Models/blob/main/wan2.1_i2v_480p_scaled_fp8_e4m3_lightx2v_4step.safetensors)
+For image to video pipeline we use [4 steps distilled i2v model](https://huggingface.co/lightx2v/Wan2.1-Distill-Models/resolve/main/wan2.1_i2v_480p_lightx2v_4step.safetensors)
 
 ```shell
-wget https://huggingface.co/lightx2v/Wan2.1-Distill-Models/blob/main/wan2.1_i2v_480p_scaled_fp8_e4m3_lightx2v_4step.safetensors
+wget https://huggingface.co/lightx2v/Wan2.1-Distill-Models/resolve/main/wan2.1_i2v_480p_lightx2v_4step.safetensors
 ```
 
 ```shell
 python img2video.py 'Astronaut riding a horse' \
-    --image ./inputs/astronaut-on-a-horse.png --checkpoint ./wan2.1_i2v_480p_scaled_fp8_e4m3_lightx2v_4step.safetensors \
+    --image ./inputs/astronaut-on-a-horse.png --checkpoint ./wan2.1_i2v_480p_lightx2v_4step.safetensors \
     --sampler euler --steps 4 --guidance 1.0 --shift 5.0 \
     --quantize --output out_i2v_distilled.mp4
 ```

video/wan2.1/img2video.py

@@ -8,7 +8,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 from tqdm import tqdm
-from wan import WanI2VPipeline
+from wan import WanPipeline
 from wan.utils import save_video
 
 
@@ -51,6 +51,12 @@ def quantization_predicate(name, m):
         "--n-prompt",
         default="镜头晃动，色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
     )
+    parser.add_argument(
+        "--teacache",
+        type=float,
+        default=0.0,
+        help="TeaCache threshold for step skipping (0=off, 0.26=recommended for i2v)",
+    )
     parser.add_argument(
         "--checkpoint",
         type=str,
@@ -94,7 +100,7 @@ def quantization_predicate(name, m):
         logging.getLogger("wan").addHandler(handler)
 
     # Load pipeline
-    pipeline = WanI2VPipeline(args.model, checkpoint=args.checkpoint)
+    pipeline = WanPipeline(args.model, checkpoint=args.checkpoint)
 
     # Quantize DiT
     if args.quantize:
@@ -117,6 +123,7 @@ def quantization_predicate(name, m):
         guidance=args.guidance,
         shift=args.shift,
         seed=args.seed,
+        teacache=args.teacache,
         verbose=args.verbose,
         denoising_step_list=denoising_step_list,
     )
@@ -129,7 +136,8 @@ def quantization_predicate(name, m):
 
     # Free T5 and CLIP memory
     del pipeline.t5
-    del pipeline.clip
+    if pipeline.clip is not None:
+        del pipeline.clip
     mx.clear_cache()
 
     # 2. Denoising loop

video/wan2.1/txt2video.py

@@ -8,7 +8,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 from tqdm import tqdm
-from wan import WanT2VPipeline
+from wan import WanPipeline
 from wan.utils import save_video
 
 
@@ -99,7 +99,7 @@ def quantization_predicate(name, m):
         logging.getLogger("wan").addHandler(handler)
 
     # Load pipeline
-    pipeline = WanT2VPipeline(args.model, checkpoint=args.checkpoint)
+    pipeline = WanPipeline(args.model, checkpoint=args.checkpoint)
 
     # Quantize DiT
     if args.quantize:

video/wan2.1/wan/__init__.py

@@ -1,3 +1,3 @@
 # Copyright © 2026 Apple Inc.
 
-from .pipeline import WanI2VPipeline, WanT2VPipeline
+from .pipeline import WanPipeline